File size: 2,946 Bytes
a0bcaae |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
import numpy as np
import wandb
from pti.pti_configs import hyperparameters, global_config
l2_criterion = torch.nn.MSELoss(reduction='mean')
def l2_loss(real_images, generated_images):
loss = l2_criterion(real_images, generated_images)
return loss
class Space_Regulizer:
def __init__(self, original_G, lpips_net):
self.original_G = original_G
self.morphing_regulizer_alpha = hyperparameters.regulizer_alpha
self.lpips_loss = lpips_net
def get_morphed_w_code(self, new_w_code, fixed_w):
interpolation_direction = new_w_code - fixed_w
interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
result_w = fixed_w + direction_to_move
self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code
return result_w
def get_image_from_ws(self, w_codes, G):
return torch.cat([G.synthesis(w_code, noise_mode='none', force_fp32=True) for w_code in w_codes])
def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch, use_wandb=False):
loss = 0.0
z_samples = np.random.randn(num_of_sampled_latents, self.original_G.z_dim)
w_samples = self.original_G.mapping(torch.from_numpy(z_samples).to(global_config.device), None,
truncation_psi=0.5)
territory_indicator_ws = [self.get_morphed_w_code(w_code.unsqueeze(0), w_batch) for w_code in w_samples]
for w_code in territory_indicator_ws:
new_img = new_G.synthesis(w_code, noise_mode='none', force_fp32=True)
with torch.no_grad():
old_img = self.original_G.synthesis(w_code, noise_mode='none', force_fp32=True)
if hyperparameters.regulizer_l2_lambda > 0:
l2_loss_val = l2_loss.l2_loss(old_img, new_img)
if use_wandb:
wandb.log({f'space_regulizer_l2_loss_val': l2_loss_val.detach().cpu()},
step=global_config.training_step)
loss += l2_loss_val * hyperparameters.regulizer_l2_lambda
if hyperparameters.regulizer_lpips_lambda > 0:
loss_lpips = self.lpips_loss(old_img, new_img)
loss_lpips = torch.mean(torch.squeeze(loss_lpips))
if use_wandb:
wandb.log({f'space_regulizer_lpips_loss_val': loss_lpips.detach().cpu()},
step=global_config.training_step)
loss += loss_lpips * hyperparameters.regulizer_lpips_lambda
return loss / len(territory_indicator_ws)
def space_regulizer_loss(self, new_G, w_batch, use_wandb):
ret_val = self.ball_holder_loss_lazy(new_G, hyperparameters.latent_ball_num_of_samples, w_batch, use_wandb)
return ret_val
|