fs_model.py 9.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242
  1. import numpy as np
  2. import torch
  3. import torch.nn as nn
  4. import torch.nn.functional as F
  5. import os
  6. from torch.autograd import Variable
  7. from util.image_pool import ImagePool
  8. from .base_model import BaseModel
  9. from . import networks
  10. from .fs_networks import Generator_Adain_Upsample, Discriminator
  11. class SpecificNorm(nn.Module):
  12. def __init__(self, epsilon=1e-8):
  13. """
  14. @notice: avoid in-place ops.
  15. 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
  16. """
  17. super(SpecificNorm, self).__init__()
  18. self.mean = np.array([0.485, 0.456, 0.406])
  19. self.mean = torch.from_numpy(self.mean).float().cuda()
  20. self.mean = self.mean.view([1, 3, 1, 1])
  21. self.std = np.array([0.229, 0.224, 0.225])
  22. self.std = torch.from_numpy(self.std).float().cuda()
  23. self.std = self.std.view([1, 3, 1, 1])
  24. def forward(self, x):
  25. mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
  26. std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
  27. x = (x - mean) / std
  28. return x
  29. class fsModel(BaseModel):
  30. def name(self):
  31. return 'fsModel'
  32. def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
  33. flags = (True, use_gan_feat_loss, use_vgg_loss, True, True, True, True, True)
  34. def loss_filter(g_gan, g_gan_feat, g_vgg, g_id, g_rec, g_mask, d_real, d_fake):
  35. 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]
  36. return loss_filter
  37. def initialize(self, opt):
  38. BaseModel.initialize(self, opt)
  39. if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
  40. torch.backends.cudnn.benchmark = True
  41. self.isTrain = opt.isTrain
  42. device = torch.device("cuda:0")
  43. # Generator network
  44. self.netG = Generator_Adain_Upsample(input_nc=3, output_nc=3, latent_size=512, n_blocks=9, deep=False)
  45. self.netG.to(device)
  46. # Id network
  47. netArc_checkpoint = opt.Arc_path
  48. netArc_checkpoint = torch.load(netArc_checkpoint)
  49. self.netArc = netArc_checkpoint['model'].module
  50. self.netArc = self.netArc.to(device)
  51. self.netArc.eval()
  52. if not self.isTrain:
  53. pretrained_path = '' if not self.isTrain else opt.load_pretrain
  54. self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
  55. return
  56. # Discriminator network
  57. if opt.gan_mode == 'original':
  58. use_sigmoid = True
  59. else:
  60. use_sigmoid = False
  61. self.netD1 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
  62. self.netD2 = Discriminator(input_nc=3, use_sigmoid=use_sigmoid)
  63. self.netD1.to(device)
  64. self.netD2.to(device)
  65. #
  66. self.spNorm =SpecificNorm()
  67. self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
  68. # load networks
  69. if opt.continue_train or opt.load_pretrain:
  70. pretrained_path = '' if not self.isTrain else opt.load_pretrain
  71. # print (pretrained_path)
  72. self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)
  73. self.load_network(self.netD1, 'D1', opt.which_epoch, pretrained_path)
  74. self.load_network(self.netD2, 'D2', opt.which_epoch, pretrained_path)
  75. if self.isTrain:
  76. # define loss functions
  77. self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
  78. self.criterionGAN = networks.GANLoss(opt.gan_mode, tensor=self.Tensor, opt=self.opt)
  79. self.criterionFeat = nn.L1Loss()
  80. self.criterionRec = nn.L1Loss()
  81. # Names so we can breakout loss
  82. self.loss_names = self.loss_filter('G_GAN', 'G_GAN_Feat', 'G_VGG', 'G_ID', 'G_Rec', 'D_GP',
  83. 'D_real', 'D_fake')
  84. # initialize optimizers
  85. # optimizer G
  86. params = list(self.netG.parameters())
  87. self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
  88. # optimizer D
  89. params = list(self.netD1.parameters()) + list(self.netD2.parameters())
  90. self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
  91. def _gradinet_penalty_D(self, netD, img_att, img_fake):
  92. # interpolate sample
  93. bs = img_fake.shape[0]
  94. alpha = torch.rand(bs, 1, 1, 1).expand_as(img_fake).cuda()
  95. interpolated = Variable(alpha * img_att + (1 - alpha) * img_fake, requires_grad=True)
  96. pred_interpolated = netD.forward(interpolated)
  97. pred_interpolated = pred_interpolated[-1]
  98. # compute gradients
  99. grad = torch.autograd.grad(outputs=pred_interpolated,
  100. inputs=interpolated,
  101. grad_outputs=torch.ones(pred_interpolated.size()).cuda(),
  102. retain_graph=True,
  103. create_graph=True,
  104. only_inputs=True)[0]
  105. # penalize gradients
  106. grad = grad.view(grad.size(0), -1)
  107. grad_l2norm = torch.sqrt(torch.sum(grad ** 2, dim=1))
  108. loss_d_gp = torch.mean((grad_l2norm - 1) ** 2)
  109. return loss_d_gp
  110. def cosin_metric(self, x1, x2):
  111. #return np.dot(x1, x2) / (np.linalg.norm(x1) * np.linalg.norm(x2))
  112. return torch.sum(x1 * x2, dim=1) / (torch.norm(x1, dim=1) * torch.norm(x2, dim=1))
  113. def forward(self, img_id, img_att, latent_id, latent_att, for_G=False):
  114. loss_D_fake, loss_D_real, loss_D_GP = 0, 0, 0
  115. loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_G_ID, loss_G_Rec = 0,0,0,0,0
  116. img_fake = self.netG.forward(img_att, latent_id)
  117. if not self.isTrain:
  118. return img_fake
  119. img_fake_downsample = self.downsample(img_fake)
  120. img_att_downsample = self.downsample(img_att)
  121. # D_Fake
  122. fea1_fake = self.netD1.forward(img_fake.detach())
  123. fea2_fake = self.netD2.forward(img_fake_downsample.detach())
  124. pred_fake = [fea1_fake, fea2_fake]
  125. loss_D_fake = self.criterionGAN(pred_fake, False, for_discriminator=True)
  126. # D_Feal
  127. fea1_real = self.netD1.forward(img_att)
  128. fea2_real = self.netD2.forward(img_att_downsample)
  129. pred_real = [fea1_real, fea2_real]
  130. fea_real = [fea1_real, fea2_real]
  131. loss_D_real = self.criterionGAN(pred_real, True, for_discriminator=True)
  132. #print('=====================D_Real========================')
  133. # D_GP
  134. loss_D_GP = 0
  135. # G_GAN
  136. fea1_fake = self.netD1.forward(img_fake)
  137. fea2_fake = self.netD2.forward(img_fake_downsample)
  138. #pred_fake = [fea1_fake[-1], fea2_fake[-1]]
  139. pred_fake = [fea1_fake, fea2_fake]
  140. fea_fake = [fea1_fake, fea2_fake]
  141. loss_G_GAN = self.criterionGAN(pred_fake, True, for_discriminator=False)
  142. # GAN feature matching loss
  143. n_layers_D = 4
  144. num_D = 2
  145. if not self.opt.no_ganFeat_loss:
  146. feat_weights = 4.0 / (n_layers_D + 1)
  147. D_weights = 1.0 / num_D
  148. for i in range(num_D):
  149. for j in range(0, len(fea_fake[i]) - 1):
  150. loss_G_GAN_Feat += D_weights * feat_weights * \
  151. self.criterionFeat(fea_fake[i][j],
  152. fea_real[i][j].detach()) * self.opt.lambda_feat
  153. #G_ID
  154. img_fake_down = F.interpolate(img_fake, scale_factor=0.5)
  155. img_fake_down = self.spNorm(img_fake_down)
  156. latent_fake = self.netArc(img_fake_down)
  157. loss_G_ID = (1 - self.cosin_metric(latent_fake, latent_id))
  158. #print('=====================G_ID========================')
  159. #print(loss_G_ID)
  160. #G_Rec
  161. loss_G_Rec = self.criterionRec(img_fake, img_att) * self.opt.lambda_rec
  162. # Only return the fake_B image if necessary to save BW
  163. 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),
  164. img_fake]
  165. def save(self, which_epoch):
  166. self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
  167. self.save_network(self.netD1, 'D1', which_epoch, self.gpu_ids)
  168. self.save_network(self.netD2, 'D2', which_epoch, self.gpu_ids)
  169. '''if self.gen_features:
  170. self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)'''
  171. def update_fixed_params(self):
  172. # after fixing the global generator for a number of iterations, also start finetuning it
  173. params = list(self.netG.parameters())
  174. if self.gen_features:
  175. params += list(self.netE.parameters())
  176. self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
  177. if self.opt.verbose:
  178. print('------------ Now also finetuning global generator -----------')
  179. def update_learning_rate(self):
  180. lrd = self.opt.lr / self.opt.niter_decay
  181. lr = self.old_lr - lrd
  182. for param_group in self.optimizer_D.param_groups:
  183. param_group['lr'] = lr
  184. for param_group in self.optimizer_G.param_groups:
  185. param_group['lr'] = lr
  186. if self.opt.verbose:
  187. print('update learning rate: %f -> %f' % (self.old_lr, lr))
  188. self.old_lr = lr