| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242 | 
							- import numpy as np
 
- import torch
 
- import torch.nn as nn
 
- import torch.nn.functional as F
 
- import os
 
- from torch.autograd import Variable
 
- from util.image_pool import ImagePool
 
- from .base_model import BaseModel
 
- from . import networks
 
- from .fs_networks import Generator_Adain_Upsample, Discriminator
 
- class SpecificNorm(nn.Module):
 
-     def __init__(self, epsilon=1e-8):
 
-         """
 
-             @notice: avoid in-place ops.
 
-             https://discuss.pytorch.org/t/encounter-the-runtimeerror-one-of-the-variables-needed-for-gradient-computation-has-been-modified-by-an-inplace-operation/836/3
 
-         """
 
-         super(SpecificNorm, self).__init__()
 
-         self.mean = np.array([0.485, 0.456, 0.406])
 
-         self.mean = torch.from_numpy(self.mean).float().cuda()
 
-         self.mean = self.mean.view([1, 3, 1, 1])
 
-         self.std = np.array([0.229, 0.224, 0.225])
 
-         self.std = torch.from_numpy(self.std).float().cuda()
 
-         self.std = self.std.view([1, 3, 1, 1])
 
-     def forward(self, x):
 
-         mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
 
-         std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
 
-         x = (x - mean) / std
 
-         return x
 
- class fsModel(BaseModel):
 
-     def name(self):
 
-         return 'fsModel'
 
-     def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
 
-         flags = (True, use_gan_feat_loss, use_vgg_loss, True, True, True, True, True)
 
-         def loss_filter(g_gan, g_gan_feat, g_vgg, g_id, g_rec, g_mask, d_real, d_fake):
 
-             return [l for (l, f) in zip((g_gan, g_gan_feat, g_vgg, g_id, g_rec, g_mask, d_real, d_fake), flags) if f]
 
-         return loss_filter
 
-     def initialize(self, opt):
 
-         BaseModel.initialize(self, opt)
 
-         if opt.resize_or_crop != 'none' or not opt.isTrain:  # when training at full res this causes OOM
 
-             torch.backends.cudnn.benchmark = True
 
-         self.isTrain = opt.isTrain
 
-         device = torch.device("cuda:0")
 
-         # Generator network
 
-         self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False)
 
-         self.netG.to(device)
 
-         # Id network
 
-         netArc_checkpoint = opt.Arc_path
 
-         netArc_checkpoint = torch.load(netArc_checkpoint)
 
-         self.netArc = netArc_checkpoint['model'].module
 
-         self.netArc = self.netArc.to(device)
 
-         self.netArc.eval()
 
-         if not self.isTrain:
 
-             pretrained_path = '' if not self.isTrain else opt.load_pretrain
 
-             self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
 
-             return
 
-         # Discriminator network
 
-         if opt.gan_mode == 'original':
 
-             use_sigmoid = True
 
-         else:
 
-             use_sigmoid = False
 
-         self.netD1 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
 
-         self.netD2 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
 
-         self.netD1.to(device)
 
-         self.netD2.to(device)
 
-         #
 
-         self.spNorm =SpecificNorm()
 
-         self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
 
-         # load networks
 
-         if opt.continue_train or opt.load_pretrain:
 
-             pretrained_path = '' if not self.isTrain else opt.load_pretrain
 
-             # print (pretrained_path)
 
-             self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
 
-             self.load_network(self.netD1, 'D1', opt.which_epoch, pretrained_path)
 
-             self.load_network(self.netD2, 'D2', opt.which_epoch, pretrained_path)
 
-         if self.isTrain:
 
-             # define loss functions
 
-             self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
 
-             self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.Tensor, opt=self.opt)
 
-             self.criterionFeat = nn.L1Loss()
 
-             self.criterionRec = nn.L1Loss()
 
-             # Names so we can breakout loss
 
-             self.loss_names = self.loss_filter('G_GAN', 'G_GAN_Feat', 'G_VGG', 'G_ID', 'G_Rec', 'D_GP',
 
-                                                'D_real', 'D_fake')
 
-            # initialize optimizers
 
-             # optimizer G
 
-             params = list(self.netG.parameters())
 
-             self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
 
-             # optimizer D
 
-             params = list(self.netD1.parameters()) + list(self.netD2.parameters())
 
-             self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
 
-     def _gradinet_penalty_D(self, netD, img_att, img_fake):
 
-         # interpolate sample
 
-         bs = img_fake.shape[0]
 
-         alpha = torch.rand(bs, 1, 1, 1).expand_as(img_fake).cuda()
 
-         interpolated = Variable(alpha * img_att + (1 - alpha) * img_fake, requires_grad=True)
 
-         pred_interpolated = netD.forward(interpolated)
 
-         pred_interpolated = pred_interpolated[-1]
 
-         # compute gradients
 
-         grad = torch.autograd.grad(outputs=pred_interpolated,
 
-                                    inputs=interpolated,
 
-                                    grad_outputs=torch.ones(pred_interpolated.size()).cuda(),
 
-                                    retain_graph=True,
 
-                                    create_graph=True,
 
-                                    only_inputs=True)[0]
 
-         # penalize gradients
 
-         grad = grad.view(grad.size(0), -1)
 
-         grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
 
-         loss_d_gp = torch.mean((grad_l2norm - 1) ** 2)
 
-         return loss_d_gp
 
-     def cosin_metric(self, x1, x2):
 
-         #return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))
 
-         return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1))
 
-     def forward(self, img_id, img_att, latent_id, latent_att, for_G=False):
 
-         loss_D_fake, loss_D_real, loss_D_GP = 0, 0, 0
 
-         loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_ID, loss_G_Rec = 0,0,0,0,0
 
-         img_fake = self.netG.forward(img_att, latent_id)
 
-         if not self.isTrain:
 
-             return img_fake
 
-         img_fake_downsample = self.downsample(img_fake)
 
-         img_att_downsample = self.downsample(img_att)
 
-         # D_Fake
 
-         fea1_fake = self.netD1.forward(img_fake.detach())
 
-         fea2_fake = self.netD2.forward(img_fake_downsample.detach())
 
-         pred_fake = [fea1_fake, fea2_fake]
 
-         loss_D_fake = self.criterionGAN(pred_fake, False, for_discriminator=True)
 
-         # D_Feal
 
-         fea1_real = self.netD1.forward(img_att)
 
-         fea2_real = self.netD2.forward(img_att_downsample)
 
-         pred_real = [fea1_real, fea2_real]
 
-         fea_real = [fea1_real, fea2_real]
 
-         loss_D_real = self.criterionGAN(pred_real, True, for_discriminator=True)
 
-         #print('=====================D_Real========================')
 
-         # D_GP
 
-         loss_D_GP = 0
 
-         # G_GAN
 
-         fea1_fake = self.netD1.forward(img_fake)
 
-         fea2_fake = self.netD2.forward(img_fake_downsample)
 
-         #pred_fake = [fea1_fake[-1], fea2_fake[-1]]
 
-         pred_fake = [fea1_fake, fea2_fake]
 
-         fea_fake = [fea1_fake, fea2_fake]
 
-         loss_G_GAN = self.criterionGAN(pred_fake, True, for_discriminator=False)
 
-         # GAN feature matching loss
 
-         n_layers_D = 4
 
-         num_D = 2
 
-         if not self.opt.no_ganFeat_loss:
 
-             feat_weights = 4.0 / (n_layers_D + 1)
 
-             D_weights = 1.0 / num_D
 
-             for i in range(num_D):
 
-                 for j in range(0, len(fea_fake[i]) - 1):
 
-                     loss_G_GAN_Feat += D_weights * feat_weights * \
 
-                                        self.criterionFeat(fea_fake[i][j],
 
-                                                           fea_real[i][j].detach()) * self.opt.lambda_feat
 
-         #G_ID
 
-         img_fake_down = F.interpolate(img_fake, scale_factor=0.5)
 
-         img_fake_down = self.spNorm(img_fake_down)
 
-         latent_fake = self.netArc(img_fake_down)
 
-         loss_G_ID = (1 - self.cosin_metric(latent_fake, latent_id))
 
-         #print('=====================G_ID========================')
 
-         #print(loss_G_ID)
 
-         #G_Rec
 
-         loss_G_Rec = self.criterionRec(img_fake, img_att) * self.opt.lambda_rec
 
-         # Only return the fake_B image if necessary to save BW
 
-         return [self.loss_filter(loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_ID, loss_G_Rec, loss_D_GP, loss_D_real, loss_D_fake),
 
-                 img_fake]
 
-     def save(self, which_epoch):
 
-         self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
 
-         self.save_network(self.netD1, 'D1', which_epoch, self.gpu_ids)
 
-         self.save_network(self.netD2, 'D2', which_epoch, self.gpu_ids)
 
-         '''if self.gen_features:
 
-             self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)'''
 
-     def update_fixed_params(self):
 
-         # after fixing the global generator for a number of iterations, also start finetuning it
 
-         params = list(self.netG.parameters())
 
-         if self.gen_features:
 
-             params += list(self.netE.parameters())
 
-         self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
 
-         if self.opt.verbose:
 
-             print('------------ Now also finetuning global generator -----------')
 
-     def update_learning_rate(self):
 
-         lrd = self.opt.lr / self.opt.niter_decay
 
-         lr = self.old_lr - lrd
 
-         for param_group in self.optimizer_D.param_groups:
 
-             param_group['lr'] = lr
 
-         for param_group in self.optimizer_G.param_groups:
 
-             param_group['lr'] = lr
 
-         if self.opt.verbose:
 
-             print('update learning rate: %f -> %f' % (self.old_lr, lr))
 
-         self.old_lr = lr
 
 
  |