networks.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845
  1. import torch
  2. import torch.nn as nn
  3. import functools
  4. from torch.autograd import Variable
  5. import numpy as np
  6. from torchvision import transforms
  7. import torch.nn.functional as F
  8. ###############################################################################
  9. # Functions
  10. ###############################################################################
  11. def weights_init(m):
  12. classname = m.__class__.__name__
  13. if classname.find('Conv') != -1:
  14. m.weight.data.normal_(0.0, 0.02)
  15. elif classname.find('BatchNorm2d') != -1:
  16. m.weight.data.normal_(1.0, 0.02)
  17. m.bias.data.fill_(0)
  18. def get_norm_layer(norm_type='instance'):
  19. if norm_type == 'batch':
  20. norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
  21. elif norm_type == 'instance':
  22. norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
  23. else:
  24. raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
  25. return norm_layer
  26. def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1,
  27. n_blocks_local=3, norm='instance', gpu_ids=[]):
  28. norm_layer = get_norm_layer(norm_type=norm)
  29. if netG == 'global':
  30. netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
  31. elif netG == 'local':
  32. netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global,
  33. n_local_enhancers, n_blocks_local, norm_layer)
  34. elif netG == 'encoder':
  35. netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer)
  36. else:
  37. raise('generator not implemented!')
  38. print(netG)
  39. if len(gpu_ids) > 0:
  40. assert(torch.cuda.is_available())
  41. netG.cuda(gpu_ids[0])
  42. netG.apply(weights_init)
  43. return netG
  44. def define_G_Adain(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
  45. norm_layer = get_norm_layer(norm_type=norm)
  46. netG = Generator_Adain(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
  47. print(netG)
  48. if len(gpu_ids) > 0:
  49. assert(torch.cuda.is_available())
  50. netG.cuda(gpu_ids[0])
  51. netG.apply(weights_init)
  52. return netG
  53. def define_G_Adain_Mask(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
  54. norm_layer = get_norm_layer(norm_type=norm)
  55. netG = Generator_Adain_Mask(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
  56. print(netG)
  57. if len(gpu_ids) > 0:
  58. assert(torch.cuda.is_available())
  59. netG.cuda(gpu_ids[0])
  60. netG.apply(weights_init)
  61. return netG
  62. def define_G_Adain_Upsample(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
  63. norm_layer = get_norm_layer(norm_type=norm)
  64. netG = Generator_Adain_Upsample(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
  65. print(netG)
  66. if len(gpu_ids) > 0:
  67. assert(torch.cuda.is_available())
  68. netG.cuda(gpu_ids[0])
  69. netG.apply(weights_init)
  70. return netG
  71. def define_G_Adain_2(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
  72. norm_layer = get_norm_layer(norm_type=norm)
  73. netG = Generator_Adain_2(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
  74. print(netG)
  75. if len(gpu_ids) > 0:
  76. assert(torch.cuda.is_available())
  77. netG.cuda(gpu_ids[0])
  78. netG.apply(weights_init)
  79. return netG
  80. def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]):
  81. norm_layer = get_norm_layer(norm_type=norm)
  82. netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
  83. print(netD)
  84. if len(gpu_ids) > 0:
  85. assert(torch.cuda.is_available())
  86. netD.cuda(gpu_ids[0])
  87. netD.apply(weights_init)
  88. return netD
  89. def print_network(net):
  90. if isinstance(net, list):
  91. net = net[0]
  92. num_params = 0
  93. for param in net.parameters():
  94. num_params += param.numel()
  95. print(net)
  96. print('Total number of parameters: %d' % num_params)
  97. ##############################################################################
  98. # Losses
  99. ##############################################################################
  100. class GANLoss(nn.Module):
  101. def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0,
  102. tensor=torch.FloatTensor, opt=None):
  103. super(GANLoss, self).__init__()
  104. self.real_label = target_real_label
  105. self.fake_label = target_fake_label
  106. self.real_label_tensor = None
  107. self.fake_label_tensor = None
  108. self.zero_tensor = None
  109. self.Tensor = tensor
  110. self.gan_mode = gan_mode
  111. self.opt = opt
  112. if gan_mode == 'ls':
  113. pass
  114. elif gan_mode == 'original':
  115. pass
  116. elif gan_mode == 'w':
  117. pass
  118. elif gan_mode == 'hinge':
  119. pass
  120. else:
  121. raise ValueError('Unexpected gan_mode {}'.format(gan_mode))
  122. def get_target_tensor(self, input, target_is_real):
  123. if target_is_real:
  124. if self.real_label_tensor is None:
  125. self.real_label_tensor = self.Tensor(1).fill_(self.real_label)
  126. self.real_label_tensor.requires_grad_(False)
  127. return self.real_label_tensor.expand_as(input)
  128. else:
  129. if self.fake_label_tensor is None:
  130. self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)
  131. self.fake_label_tensor.requires_grad_(False)
  132. return self.fake_label_tensor.expand_as(input)
  133. def get_zero_tensor(self, input):
  134. if self.zero_tensor is None:
  135. self.zero_tensor = self.Tensor(1).fill_(0)
  136. self.zero_tensor.requires_grad_(False)
  137. return self.zero_tensor.expand_as(input)
  138. def loss(self, input, target_is_real, for_discriminator=True):
  139. if self.gan_mode == 'original': # cross entropy loss
  140. target_tensor = self.get_target_tensor(input, target_is_real)
  141. loss = F.binary_cross_entropy_with_logits(input, target_tensor)
  142. return loss
  143. elif self.gan_mode == 'ls':
  144. target_tensor = self.get_target_tensor(input, target_is_real)
  145. return F.mse_loss(input, target_tensor)
  146. elif self.gan_mode == 'hinge':
  147. if for_discriminator:
  148. if target_is_real:
  149. minval = torch.min(input - 1, self.get_zero_tensor(input))
  150. loss = -torch.mean(minval)
  151. else:
  152. minval = torch.min(-input - 1, self.get_zero_tensor(input))
  153. loss = -torch.mean(minval)
  154. else:
  155. assert target_is_real, "The generator's hinge loss must be aiming for real"
  156. loss = -torch.mean(input)
  157. return loss
  158. else:
  159. # wgan
  160. if target_is_real:
  161. return -input.mean()
  162. else:
  163. return input.mean()
  164. def __call__(self, input, target_is_real, for_discriminator=True):
  165. # computing loss is a bit complicated because |input| may not be
  166. # a tensor, but list of tensors in case of multiscale discriminator
  167. if isinstance(input, list):
  168. loss = 0
  169. for pred_i in input:
  170. if isinstance(pred_i, list):
  171. pred_i = pred_i[-1]
  172. loss_tensor = self.loss(pred_i, target_is_real, for_discriminator)
  173. bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)
  174. new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)
  175. loss += new_loss
  176. return loss / len(input)
  177. else:
  178. return self.loss(input, target_is_real, for_discriminator)
  179. class VGGLoss(nn.Module):
  180. def __init__(self, gpu_ids):
  181. super(VGGLoss, self).__init__()
  182. self.vgg = Vgg19().cuda()
  183. self.criterion = nn.L1Loss()
  184. self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
  185. def forward(self, x, y):
  186. x_vgg, y_vgg = self.vgg(x), self.vgg(y)
  187. loss = 0
  188. for i in range(len(x_vgg)):
  189. loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
  190. return loss
  191. ##############################################################################
  192. # Generator
  193. ##############################################################################
  194. class LocalEnhancer(nn.Module):
  195. def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9,
  196. n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'):
  197. super(LocalEnhancer, self).__init__()
  198. self.n_local_enhancers = n_local_enhancers
  199. ###### global generator model #####
  200. ngf_global = ngf * (2**n_local_enhancers)
  201. model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model
  202. model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers
  203. self.model = nn.Sequential(*model_global)
  204. ###### local enhancer layers #####
  205. for n in range(1, n_local_enhancers+1):
  206. ### downsample
  207. ngf_global = ngf * (2**(n_local_enhancers-n))
  208. model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
  209. norm_layer(ngf_global), nn.ReLU(True),
  210. nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1),
  211. norm_layer(ngf_global * 2), nn.ReLU(True)]
  212. ### residual blocks
  213. model_upsample = []
  214. for i in range(n_blocks_local):
  215. model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)]
  216. ### upsample
  217. model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1),
  218. norm_layer(ngf_global), nn.ReLU(True)]
  219. ### final convolution
  220. if n == n_local_enhancers:
  221. model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  222. setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample))
  223. setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample))
  224. self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
  225. def forward(self, input):
  226. ### create input pyramid
  227. input_downsampled = [input]
  228. for i in range(self.n_local_enhancers):
  229. input_downsampled.append(self.downsample(input_downsampled[-1]))
  230. ### output at coarest level
  231. output_prev = self.model(input_downsampled[-1])
  232. ### build up one layer at a time
  233. for n_local_enhancers in range(1, self.n_local_enhancers+1):
  234. model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1')
  235. model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2')
  236. input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers]
  237. output_prev = model_upsample(model_downsample(input_i) + output_prev)
  238. return output_prev
  239. class GlobalGenerator(nn.Module):
  240. def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
  241. padding_type='reflect'):
  242. assert(n_blocks >= 0)
  243. super(GlobalGenerator, self).__init__()
  244. activation = nn.ReLU(True)
  245. model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
  246. ### downsample
  247. for i in range(n_downsampling):
  248. mult = 2**i
  249. model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  250. norm_layer(ngf * mult * 2), activation]
  251. ### resnet blocks
  252. mult = 2**n_downsampling
  253. for i in range(n_blocks):
  254. model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]
  255. ### upsample
  256. for i in range(n_downsampling):
  257. mult = 2**(n_downsampling - i)
  258. model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
  259. norm_layer(int(ngf * mult / 2)), activation]
  260. model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  261. self.model = nn.Sequential(*model)
  262. def forward(self, input):
  263. return self.model(input)
  264. # Define a resnet block
  265. class ResnetBlock(nn.Module):
  266. def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
  267. super(ResnetBlock, self).__init__()
  268. self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)
  269. def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
  270. conv_block = []
  271. p = 0
  272. if padding_type == 'reflect':
  273. conv_block += [nn.ReflectionPad2d(1)]
  274. elif padding_type == 'replicate':
  275. conv_block += [nn.ReplicationPad2d(1)]
  276. elif padding_type == 'zero':
  277. p = 1
  278. else:
  279. raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  280. conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
  281. norm_layer(dim),
  282. activation]
  283. if use_dropout:
  284. conv_block += [nn.Dropout(0.5)]
  285. p = 0
  286. if padding_type == 'reflect':
  287. conv_block += [nn.ReflectionPad2d(1)]
  288. elif padding_type == 'replicate':
  289. conv_block += [nn.ReplicationPad2d(1)]
  290. elif padding_type == 'zero':
  291. p = 1
  292. else:
  293. raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  294. conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
  295. norm_layer(dim)]
  296. return nn.Sequential(*conv_block)
  297. def forward(self, x):
  298. out = x + self.conv_block(x)
  299. return out
  300. class InstanceNorm(nn.Module):
  301. def __init__(self, epsilon=1e-8):
  302. """
  303. @notice: avoid in-place ops.
  304. 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
  305. """
  306. super(InstanceNorm, self).__init__()
  307. self.epsilon = epsilon
  308. def forward(self, x):
  309. x = x - torch.mean(x, (2, 3), True)
  310. tmp = torch.mul(x, x) # or x ** 2
  311. tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
  312. return x * tmp
  313. class SpecificNorm(nn.Module):
  314. def __init__(self, epsilon=1e-8):
  315. """
  316. @notice: avoid in-place ops.
  317. 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
  318. """
  319. super(SpecificNorm, self).__init__()
  320. self.mean = np.array([0.485, 0.456, 0.406])
  321. self.mean = torch.from_numpy(self.mean).float().cuda()
  322. self.mean = self.mean.view([1, 3, 1, 1])
  323. self.std = np.array([0.229, 0.224, 0.225])
  324. self.std = torch.from_numpy(self.std).float().cuda()
  325. self.std = self.std.view([1, 3, 1, 1])
  326. def forward(self, x):
  327. mean = self.mean.expand([1, 3, x.shape[2], x.shape[3]])
  328. std = self.std.expand([1, 3, x.shape[2], x.shape[3]])
  329. x = (x - mean) / std
  330. return x
  331. class ApplyStyle(nn.Module):
  332. """
  333. @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
  334. """
  335. def __init__(self, latent_size, channels):
  336. super(ApplyStyle, self).__init__()
  337. self.linear = nn.Linear(latent_size, channels * 2)
  338. def forward(self, x, latent):
  339. style = self.linear(latent) # style => [batch_size, n_channels*2]
  340. shape = [-1, 2, x.size(1), 1, 1]
  341. style = style.view(shape) # [batch_size, 2, n_channels, ...]
  342. x = x * (style[:, 0] + 1.) + style[:, 1]
  343. return x
  344. class ResnetBlock_Adain(nn.Module):
  345. def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
  346. super(ResnetBlock_Adain, self).__init__()
  347. p = 0
  348. conv1 = []
  349. if padding_type == 'reflect':
  350. conv1 += [nn.ReflectionPad2d(1)]
  351. elif padding_type == 'replicate':
  352. conv1 += [nn.ReplicationPad2d(1)]
  353. elif padding_type == 'zero':
  354. p = 1
  355. else:
  356. raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  357. conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
  358. self.conv1 = nn.Sequential(*conv1)
  359. self.style1 = ApplyStyle(latent_size, dim)
  360. self.act1 = activation
  361. p = 0
  362. conv2 = []
  363. if padding_type == 'reflect':
  364. conv2 += [nn.ReflectionPad2d(1)]
  365. elif padding_type == 'replicate':
  366. conv2 += [nn.ReplicationPad2d(1)]
  367. elif padding_type == 'zero':
  368. p = 1
  369. else:
  370. raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  371. conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
  372. self.conv2 = nn.Sequential(*conv2)
  373. self.style2 = ApplyStyle(latent_size, dim)
  374. def forward(self, x, dlatents_in_slice):
  375. y = self.conv1(x)
  376. y = self.style1(y, dlatents_in_slice)
  377. y = self.act1(y)
  378. y = self.conv2(y)
  379. y = self.style2(y, dlatents_in_slice)
  380. out = x + y
  381. return out
  382. class UpBlock_Adain(nn.Module):
  383. def __init__(self, dim_in, dim_out, latent_size, padding_type, activation=nn.ReLU(True)):
  384. super(UpBlock_Adain, self).__init__()
  385. p = 0
  386. conv1 = [nn.Upsample(scale_factor=2, mode='bilinear')]
  387. if padding_type == 'reflect':
  388. conv1 += [nn.ReflectionPad2d(1)]
  389. elif padding_type == 'replicate':
  390. conv1 += [nn.ReplicationPad2d(1)]
  391. elif padding_type == 'zero':
  392. p = 1
  393. else:
  394. raise NotImplementedError('padding [%s] is not implemented' % padding_type)
  395. conv1 += [nn.Conv2d(dim_in, dim_out, kernel_size=3, padding = p), InstanceNorm()]
  396. self.conv1 = nn.Sequential(*conv1)
  397. self.style1 = ApplyStyle(latent_size, dim_out)
  398. self.act1 = activation
  399. def forward(self, x, dlatents_in_slice):
  400. y = self.conv1(x)
  401. y = self.style1(y, dlatents_in_slice)
  402. y = self.act1(y)
  403. return y
  404. class Encoder(nn.Module):
  405. def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d):
  406. super(Encoder, self).__init__()
  407. self.output_nc = output_nc
  408. model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
  409. norm_layer(ngf), nn.ReLU(True)]
  410. ### downsample
  411. for i in range(n_downsampling):
  412. mult = 2**i
  413. model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  414. norm_layer(ngf * mult * 2), nn.ReLU(True)]
  415. ### upsample
  416. for i in range(n_downsampling):
  417. mult = 2**(n_downsampling - i)
  418. model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
  419. norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]
  420. model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  421. self.model = nn.Sequential(*model)
  422. def forward(self, input, inst):
  423. outputs = self.model(input)
  424. # instance-wise average pooling
  425. outputs_mean = outputs.clone()
  426. inst_list = np.unique(inst.cpu().numpy().astype(int))
  427. for i in inst_list:
  428. for b in range(input.size()[0]):
  429. indices = (inst[b:b+1] == int(i)).nonzero() # n x 4
  430. for j in range(self.output_nc):
  431. output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]]
  432. mean_feat = torch.mean(output_ins).expand_as(output_ins)
  433. outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat
  434. return outputs_mean
  435. class Generator_Adain(nn.Module):
  436. def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
  437. padding_type='reflect'):
  438. assert (n_blocks >= 0)
  439. super(Generator_Adain, self).__init__()
  440. activation = nn.ReLU(True)
  441. Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
  442. ### downsample
  443. for i in range(n_downsampling):
  444. mult = 2 ** i
  445. Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  446. norm_layer(ngf * mult * 2), activation]
  447. self.Encoder = nn.Sequential(*Enc)
  448. ### resnet blocks
  449. BN = []
  450. mult = 2 ** n_downsampling
  451. for i in range(n_blocks):
  452. BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
  453. self.BottleNeck = nn.Sequential(*BN)
  454. '''self.ResBlockAdain1 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
  455. activation=activation)
  456. self.ResBlockAdain2 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
  457. activation=activation)
  458. self.ResBlockAdain3 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
  459. activation=activation)
  460. self.ResBlockAdain4 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
  461. activation=activation)'''
  462. ### upsample
  463. Dec = []
  464. for i in range(n_downsampling):
  465. mult = 2 ** (n_downsampling - i)
  466. Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
  467. output_padding=1),
  468. norm_layer(int(ngf * mult / 2)), activation]
  469. Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  470. self.Decoder = nn.Sequential(*Dec)
  471. #self.model = nn.Sequential(*model)
  472. self.spNorm = SpecificNorm()
  473. def forward(self, input, dlatents):
  474. x = input
  475. x = self.Encoder(x)
  476. for i in range(len(self.BottleNeck)):
  477. x = self.BottleNeck[i](x, dlatents)
  478. '''x = self.ResBlockAdain1(x, dlatents)
  479. x = self.ResBlockAdain2(x, dlatents)
  480. x = self.ResBlockAdain3(x, dlatents)
  481. x = self.ResBlockAdain4(x, dlatents)'''
  482. x = self.Decoder(x)
  483. x = (x + 1) / 2
  484. x = self.spNorm(x)
  485. return x
  486. class Generator_Adain_Mask(nn.Module):
  487. def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
  488. padding_type='reflect'):
  489. assert (n_blocks >= 0)
  490. super(Generator_Adain_Mask, self).__init__()
  491. activation = nn.ReLU(True)
  492. Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
  493. ### downsample
  494. for i in range(n_downsampling):
  495. mult = 2 ** i
  496. Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  497. norm_layer(ngf * mult * 2), activation]
  498. self.Encoder = nn.Sequential(*Enc)
  499. ### resnet blocks
  500. BN = []
  501. mult = 2 ** n_downsampling
  502. for i in range(n_blocks):
  503. BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
  504. self.BottleNeck = nn.Sequential(*BN)
  505. ### upsample
  506. Dec = []
  507. for i in range(n_downsampling):
  508. mult = 2 ** (n_downsampling - i)
  509. Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
  510. output_padding=1),
  511. norm_layer(int(ngf * mult / 2)), activation]
  512. Fake_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  513. Mast_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, 1, kernel_size=7, padding=0), nn.Sigmoid()]
  514. self.Decoder = nn.Sequential(*Dec)
  515. #self.model = nn.Sequential(*model)
  516. self.spNorm = SpecificNorm()
  517. self.Fake_out = nn.Sequential(*Fake_out)
  518. self.Mask_out = nn.Sequential(*Mast_out)
  519. def forward(self, input, dlatents):
  520. x = input
  521. x = self.Encoder(x)
  522. for i in range(len(self.BottleNeck)):
  523. x = self.BottleNeck[i](x, dlatents)
  524. x = self.Decoder(x)
  525. fake_out = self.Fake_out(x)
  526. mask_out = self.Mask_out(x)
  527. fake_out = (fake_out + 1) / 2
  528. fake_out = self.spNorm(fake_out)
  529. generated = fake_out * mask_out + input * (1-mask_out)
  530. return generated, mask_out
  531. class Generator_Adain_Upsample(nn.Module):
  532. def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
  533. padding_type='reflect'):
  534. assert (n_blocks >= 0)
  535. super(Generator_Adain_Upsample, self).__init__()
  536. activation = nn.ReLU(True)
  537. Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
  538. ### downsample
  539. for i in range(n_downsampling):
  540. mult = 2 ** i
  541. Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  542. norm_layer(ngf * mult * 2), activation]
  543. self.Encoder = nn.Sequential(*Enc)
  544. ### resnet blocks
  545. BN = []
  546. mult = 2 ** n_downsampling
  547. for i in range(n_blocks):
  548. BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
  549. self.BottleNeck = nn.Sequential(*BN)
  550. ### upsample
  551. Dec = []
  552. for i in range(n_downsampling):
  553. mult = 2 ** (n_downsampling - i)
  554. '''Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
  555. output_padding=1),
  556. norm_layer(int(ngf * mult / 2)), activation]'''
  557. Dec += [nn.Upsample(scale_factor=2, mode='bilinear'),
  558. nn.Conv2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=1, padding=1),
  559. norm_layer(int(ngf * mult / 2)), activation]
  560. Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  561. self.Decoder = nn.Sequential(*Dec)
  562. self.spNorm = SpecificNorm()
  563. def forward(self, input, dlatents):
  564. x = input
  565. x = self.Encoder(x)
  566. for i in range(len(self.BottleNeck)):
  567. x = self.BottleNeck[i](x, dlatents)
  568. x = self.Decoder(x)
  569. x = (x + 1) / 2
  570. x = self.spNorm(x)
  571. return x
  572. class Generator_Adain_2(nn.Module):
  573. def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
  574. padding_type='reflect'):
  575. assert (n_blocks >= 0)
  576. super(Generator_Adain_2, self).__init__()
  577. activation = nn.ReLU(True)
  578. Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
  579. ### downsample
  580. for i in range(n_downsampling):
  581. mult = 2 ** i
  582. Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
  583. norm_layer(ngf * mult * 2), activation]
  584. self.Encoder = nn.Sequential(*Enc)
  585. ### resnet blocks
  586. BN = []
  587. mult = 2 ** n_downsampling
  588. for i in range(n_blocks):
  589. BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
  590. self.BottleNeck = nn.Sequential(*BN)
  591. ### upsample
  592. Dec = []
  593. for i in range(n_downsampling):
  594. mult = 2 ** (n_downsampling - i)
  595. Dec += [UpBlock_Adain(dim_in=ngf * mult, dim_out=int(ngf * mult / 2), latent_size=latent_size, padding_type=padding_type)]
  596. layer_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
  597. self.Decoder = nn.Sequential(*Dec)
  598. #self.model = nn.Sequential(*model)
  599. self.spNorm = SpecificNorm()
  600. self.layer_out = nn.Sequential(*layer_out)
  601. def forward(self, input, dlatents):
  602. x = input
  603. x = self.Encoder(x)
  604. for i in range(len(self.BottleNeck)):
  605. x = self.BottleNeck[i](x, dlatents)
  606. for i in range(len(self.Decoder)):
  607. x = self.Decoder[i](x, dlatents)
  608. x = self.layer_out(x)
  609. x = (x + 1) / 2
  610. x = self.spNorm(x)
  611. return x
  612. class MultiscaleDiscriminator(nn.Module):
  613. def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d,
  614. use_sigmoid=False, num_D=3, getIntermFeat=False):
  615. super(MultiscaleDiscriminator, self).__init__()
  616. self.num_D = num_D
  617. self.n_layers = n_layers
  618. self.getIntermFeat = getIntermFeat
  619. for i in range(num_D):
  620. netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
  621. if getIntermFeat:
  622. for j in range(n_layers+2):
  623. setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))
  624. else:
  625. setattr(self, 'layer'+str(i), netD.model)
  626. self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
  627. def singleD_forward(self, model, input):
  628. if self.getIntermFeat:
  629. result = [input]
  630. for i in range(len(model)):
  631. result.append(model[i](result[-1]))
  632. return result[1:]
  633. else:
  634. return [model(input)]
  635. def forward(self, input):
  636. num_D = self.num_D
  637. result = []
  638. input_downsampled = input
  639. for i in range(num_D):
  640. if self.getIntermFeat:
  641. model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
  642. else:
  643. model = getattr(self, 'layer'+str(num_D-1-i))
  644. result.append(self.singleD_forward(model, input_downsampled))
  645. if i != (num_D-1):
  646. input_downsampled = self.downsample(input_downsampled)
  647. return result
  648. # Defines the PatchGAN discriminator with the specified arguments.
  649. class NLayerDiscriminator(nn.Module):
  650. def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
  651. super(NLayerDiscriminator, self).__init__()
  652. self.getIntermFeat = getIntermFeat
  653. self.n_layers = n_layers
  654. kw = 4
  655. padw = 1
  656. sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
  657. nf = ndf
  658. for n in range(1, n_layers):
  659. nf_prev = nf
  660. nf = min(nf * 2, 512)
  661. sequence += [[
  662. nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
  663. norm_layer(nf), nn.LeakyReLU(0.2, True)
  664. ]]
  665. nf_prev = nf
  666. nf = min(nf * 2, 512)
  667. sequence += [[
  668. nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
  669. norm_layer(nf),
  670. nn.LeakyReLU(0.2, True)
  671. ]]
  672. if use_sigmoid:
  673. sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()]]
  674. else:
  675. sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
  676. if getIntermFeat:
  677. for n in range(len(sequence)):
  678. setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
  679. else:
  680. sequence_stream = []
  681. for n in range(len(sequence)):
  682. sequence_stream += sequence[n]
  683. self.model = nn.Sequential(*sequence_stream)
  684. def forward(self, input):
  685. if self.getIntermFeat:
  686. res = [input]
  687. for n in range(self.n_layers+2):
  688. model = getattr(self, 'model'+str(n))
  689. res.append(model(res[-1]))
  690. return res[1:]
  691. else:
  692. return self.model(input)
  693. from torchvision import models
  694. class Vgg19(torch.nn.Module):
  695. def __init__(self, requires_grad=False):
  696. super(Vgg19, self).__init__()
  697. vgg_pretrained_features = models.vgg19(pretrained=True).features
  698. self.slice1 = torch.nn.Sequential()
  699. self.slice2 = torch.nn.Sequential()
  700. self.slice3 = torch.nn.Sequential()
  701. self.slice4 = torch.nn.Sequential()
  702. self.slice5 = torch.nn.Sequential()
  703. for x in range(2):
  704. self.slice1.add_module(str(x), vgg_pretrained_features[x])
  705. for x in range(2, 7):
  706. self.slice2.add_module(str(x), vgg_pretrained_features[x])
  707. for x in range(7, 12):
  708. self.slice3.add_module(str(x), vgg_pretrained_features[x])
  709. for x in range(12, 21):
  710. self.slice4.add_module(str(x), vgg_pretrained_features[x])
  711. for x in range(21, 30):
  712. self.slice5.add_module(str(x), vgg_pretrained_features[x])
  713. if not requires_grad:
  714. for param in self.parameters():
  715. param.requires_grad = False
  716. def forward(self, X):
  717. h_relu1 = self.slice1(X)
  718. h_relu2 = self.slice2(h_relu1)
  719. h_relu3 = self.slice3(h_relu2)
  720. h_relu4 = self.slice4(h_relu3)
  721. h_relu5 = self.slice5(h_relu4)
  722. out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
  723. return out