123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845 |
- import torch
- import torch.nn as nn
- import functools
- from torch.autograd import Variable
- import numpy as np
- from torchvision import transforms
- import torch.nn.functional as F
- ###############################################################################
- # Functions
- ###############################################################################
- def weights_init(m):
- classname = m.__class__.__name__
- if classname.find('Conv') != -1:
- m.weight.data.normal_(0.0, 0.02)
- elif classname.find('BatchNorm2d') != -1:
- m.weight.data.normal_(1.0, 0.02)
- m.bias.data.fill_(0)
- def get_norm_layer(norm_type='instance'):
- if norm_type == 'batch':
- norm_layer = functools.partial(nn.BatchNorm2d, affine=True)
- elif norm_type == 'instance':
- norm_layer = functools.partial(nn.InstanceNorm2d, affine=False)
- else:
- raise NotImplementedError('normalization layer [%s] is not found' % norm_type)
- return norm_layer
- def define_G(input_nc, output_nc, ngf, netG, n_downsample_global=3, n_blocks_global=9, n_local_enhancers=1,
- n_blocks_local=3, norm='instance', gpu_ids=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- if netG == 'global':
- netG = GlobalGenerator(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global, norm_layer)
- elif netG == 'local':
- netG = LocalEnhancer(input_nc, output_nc, ngf, n_downsample_global, n_blocks_global,
- n_local_enhancers, n_blocks_local, norm_layer)
- elif netG == 'encoder':
- netG = Encoder(input_nc, output_nc, ngf, n_downsample_global, norm_layer)
- else:
- raise('generator not implemented!')
- print(netG)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netG.cuda(gpu_ids[0])
- netG.apply(weights_init)
- return netG
- def define_G_Adain(input_nc, output_nc, latent_size, ngf, netG, n_downsample_global=2, n_blocks_global=4, norm='instance', gpu_ids=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- netG = Generator_Adain(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
- print(netG)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netG.cuda(gpu_ids[0])
- netG.apply(weights_init)
- return netG
- 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=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- netG = Generator_Adain_Mask(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
- print(netG)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netG.cuda(gpu_ids[0])
- netG.apply(weights_init)
- return netG
- 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=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- netG = Generator_Adain_Upsample(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
- print(netG)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netG.cuda(gpu_ids[0])
- netG.apply(weights_init)
- return netG
- 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=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- netG = Generator_Adain_2(input_nc, output_nc, latent_size, ngf, n_downsample_global, n_blocks_global, norm_layer)
- print(netG)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netG.cuda(gpu_ids[0])
- netG.apply(weights_init)
- return netG
- def define_D(input_nc, ndf, n_layers_D, norm='instance', use_sigmoid=False, num_D=1, getIntermFeat=False, gpu_ids=[]):
- norm_layer = get_norm_layer(norm_type=norm)
- netD = MultiscaleDiscriminator(input_nc, ndf, n_layers_D, norm_layer, use_sigmoid, num_D, getIntermFeat)
- print(netD)
- if len(gpu_ids) > 0:
- assert(torch.cuda.is_available())
- netD.cuda(gpu_ids[0])
- netD.apply(weights_init)
- return netD
- def print_network(net):
- if isinstance(net, list):
- net = net[0]
- num_params = 0
- for param in net.parameters():
- num_params += param.numel()
- print(net)
- print('Total number of parameters: %d' % num_params)
- ##############################################################################
- # Losses
- ##############################################################################
- class GANLoss(nn.Module):
- def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0,
- tensor=torch.FloatTensor, opt=None):
- super(GANLoss, self).__init__()
- self.real_label = target_real_label
- self.fake_label = target_fake_label
- self.real_label_tensor = None
- self.fake_label_tensor = None
- self.zero_tensor = None
- self.Tensor = tensor
- self.gan_mode = gan_mode
- self.opt = opt
- if gan_mode == 'ls':
- pass
- elif gan_mode == 'original':
- pass
- elif gan_mode == 'w':
- pass
- elif gan_mode == 'hinge':
- pass
- else:
- raise ValueError('Unexpected gan_mode {}'.format(gan_mode))
- def get_target_tensor(self, input, target_is_real):
- if target_is_real:
- if self.real_label_tensor is None:
- self.real_label_tensor = self.Tensor(1).fill_(self.real_label)
- self.real_label_tensor.requires_grad_(False)
- return self.real_label_tensor.expand_as(input)
- else:
- if self.fake_label_tensor is None:
- self.fake_label_tensor = self.Tensor(1).fill_(self.fake_label)
- self.fake_label_tensor.requires_grad_(False)
- return self.fake_label_tensor.expand_as(input)
- def get_zero_tensor(self, input):
- if self.zero_tensor is None:
- self.zero_tensor = self.Tensor(1).fill_(0)
- self.zero_tensor.requires_grad_(False)
- return self.zero_tensor.expand_as(input)
- def loss(self, input, target_is_real, for_discriminator=True):
- if self.gan_mode == 'original': # cross entropy loss
- target_tensor = self.get_target_tensor(input, target_is_real)
- loss = F.binary_cross_entropy_with_logits(input, target_tensor)
- return loss
- elif self.gan_mode == 'ls':
- target_tensor = self.get_target_tensor(input, target_is_real)
- return F.mse_loss(input, target_tensor)
- elif self.gan_mode == 'hinge':
- if for_discriminator:
- if target_is_real:
- minval = torch.min(input - 1, self.get_zero_tensor(input))
- loss = -torch.mean(minval)
- else:
- minval = torch.min(-input - 1, self.get_zero_tensor(input))
- loss = -torch.mean(minval)
- else:
- assert target_is_real, "The generator's hinge loss must be aiming for real"
- loss = -torch.mean(input)
- return loss
- else:
- # wgan
- if target_is_real:
- return -input.mean()
- else:
- return input.mean()
- def __call__(self, input, target_is_real, for_discriminator=True):
- # computing loss is a bit complicated because |input| may not be
- # a tensor, but list of tensors in case of multiscale discriminator
- if isinstance(input, list):
- loss = 0
- for pred_i in input:
- if isinstance(pred_i, list):
- pred_i = pred_i[-1]
- loss_tensor = self.loss(pred_i, target_is_real, for_discriminator)
- bs = 1 if len(loss_tensor.size()) == 0 else loss_tensor.size(0)
- new_loss = torch.mean(loss_tensor.view(bs, -1), dim=1)
- loss += new_loss
- return loss / len(input)
- else:
- return self.loss(input, target_is_real, for_discriminator)
- class VGGLoss(nn.Module):
- def __init__(self, gpu_ids):
- super(VGGLoss, self).__init__()
- self.vgg = Vgg19().cuda()
- self.criterion = nn.L1Loss()
- self.weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]
- def forward(self, x, y):
- x_vgg, y_vgg = self.vgg(x), self.vgg(y)
- loss = 0
- for i in range(len(x_vgg)):
- loss += self.weights[i] * self.criterion(x_vgg[i], y_vgg[i].detach())
- return loss
- ##############################################################################
- # Generator
- ##############################################################################
- class LocalEnhancer(nn.Module):
- def __init__(self, input_nc, output_nc, ngf=32, n_downsample_global=3, n_blocks_global=9,
- n_local_enhancers=1, n_blocks_local=3, norm_layer=nn.BatchNorm2d, padding_type='reflect'):
- super(LocalEnhancer, self).__init__()
- self.n_local_enhancers = n_local_enhancers
-
- ###### global generator model #####
- ngf_global = ngf * (2**n_local_enhancers)
- model_global = GlobalGenerator(input_nc, output_nc, ngf_global, n_downsample_global, n_blocks_global, norm_layer).model
- model_global = [model_global[i] for i in range(len(model_global)-3)] # get rid of final convolution layers
- self.model = nn.Sequential(*model_global)
- ###### local enhancer layers #####
- for n in range(1, n_local_enhancers+1):
- ### downsample
- ngf_global = ngf * (2**(n_local_enhancers-n))
- model_downsample = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf_global, kernel_size=7, padding=0),
- norm_layer(ngf_global), nn.ReLU(True),
- nn.Conv2d(ngf_global, ngf_global * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf_global * 2), nn.ReLU(True)]
- ### residual blocks
- model_upsample = []
- for i in range(n_blocks_local):
- model_upsample += [ResnetBlock(ngf_global * 2, padding_type=padding_type, norm_layer=norm_layer)]
- ### upsample
- model_upsample += [nn.ConvTranspose2d(ngf_global * 2, ngf_global, kernel_size=3, stride=2, padding=1, output_padding=1),
- norm_layer(ngf_global), nn.ReLU(True)]
- ### final convolution
- if n == n_local_enhancers:
- model_upsample += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
-
- setattr(self, 'model'+str(n)+'_1', nn.Sequential(*model_downsample))
- setattr(self, 'model'+str(n)+'_2', nn.Sequential(*model_upsample))
-
- self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
- def forward(self, input):
- ### create input pyramid
- input_downsampled = [input]
- for i in range(self.n_local_enhancers):
- input_downsampled.append(self.downsample(input_downsampled[-1]))
- ### output at coarest level
- output_prev = self.model(input_downsampled[-1])
- ### build up one layer at a time
- for n_local_enhancers in range(1, self.n_local_enhancers+1):
- model_downsample = getattr(self, 'model'+str(n_local_enhancers)+'_1')
- model_upsample = getattr(self, 'model'+str(n_local_enhancers)+'_2')
- input_i = input_downsampled[self.n_local_enhancers-n_local_enhancers]
- output_prev = model_upsample(model_downsample(input_i) + output_prev)
- return output_prev
- class GlobalGenerator(nn.Module):
- def __init__(self, input_nc, output_nc, ngf=64, n_downsampling=3, n_blocks=9, norm_layer=nn.BatchNorm2d,
- padding_type='reflect'):
- assert(n_blocks >= 0)
- super(GlobalGenerator, self).__init__()
- activation = nn.ReLU(True)
- model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
- ### downsample
- for i in range(n_downsampling):
- mult = 2**i
- model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), activation]
- ### resnet blocks
- mult = 2**n_downsampling
- for i in range(n_blocks):
- model += [ResnetBlock(ngf * mult, padding_type=padding_type, activation=activation, norm_layer=norm_layer)]
-
- ### upsample
- for i in range(n_downsampling):
- mult = 2**(n_downsampling - i)
- model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
- norm_layer(int(ngf * mult / 2)), activation]
- model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- self.model = nn.Sequential(*model)
-
- def forward(self, input):
- return self.model(input)
-
- # Define a resnet block
- class ResnetBlock(nn.Module):
- def __init__(self, dim, padding_type, norm_layer, activation=nn.ReLU(True), use_dropout=False):
- super(ResnetBlock, self).__init__()
- self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, activation, use_dropout)
- def build_conv_block(self, dim, padding_type, norm_layer, activation, use_dropout):
- conv_block = []
- p = 0
- if padding_type == 'reflect':
- conv_block += [nn.ReflectionPad2d(1)]
- elif padding_type == 'replicate':
- conv_block += [nn.ReplicationPad2d(1)]
- elif padding_type == 'zero':
- p = 1
- else:
- raise NotImplementedError('padding [%s] is not implemented' % padding_type)
- conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
- norm_layer(dim),
- activation]
- if use_dropout:
- conv_block += [nn.Dropout(0.5)]
- p = 0
- if padding_type == 'reflect':
- conv_block += [nn.ReflectionPad2d(1)]
- elif padding_type == 'replicate':
- conv_block += [nn.ReplicationPad2d(1)]
- elif padding_type == 'zero':
- p = 1
- else:
- raise NotImplementedError('padding [%s] is not implemented' % padding_type)
- conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p),
- norm_layer(dim)]
- return nn.Sequential(*conv_block)
- def forward(self, x):
- out = x + self.conv_block(x)
- return out
- class InstanceNorm(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(InstanceNorm, self).__init__()
- self.epsilon = epsilon
- def forward(self, x):
- x = x - torch.mean(x, (2, 3), True)
- tmp = torch.mul(x, x) # or x ** 2
- tmp = torch.rsqrt(torch.mean(tmp, (2, 3), True) + self.epsilon)
- return x * tmp
- 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 ApplyStyle(nn.Module):
- """
- @ref: https://github.com/lernapparat/lernapparat/blob/master/style_gan/pytorch_style_gan.ipynb
- """
- def __init__(self, latent_size, channels):
- super(ApplyStyle, self).__init__()
- self.linear = nn.Linear(latent_size, channels * 2)
- def forward(self, x, latent):
- style = self.linear(latent) # style => [batch_size, n_channels*2]
- shape = [-1, 2, x.size(1), 1, 1]
- style = style.view(shape) # [batch_size, 2, n_channels, ...]
- x = x * (style[:, 0] + 1.) + style[:, 1]
- return x
- class ResnetBlock_Adain(nn.Module):
- def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
- super(ResnetBlock_Adain, self).__init__()
- p = 0
- conv1 = []
- if padding_type == 'reflect':
- conv1 += [nn.ReflectionPad2d(1)]
- elif padding_type == 'replicate':
- conv1 += [nn.ReplicationPad2d(1)]
- elif padding_type == 'zero':
- p = 1
- else:
- raise NotImplementedError('padding [%s] is not implemented' % padding_type)
- conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
- self.conv1 = nn.Sequential(*conv1)
- self.style1 = ApplyStyle(latent_size, dim)
- self.act1 = activation
- p = 0
- conv2 = []
- if padding_type == 'reflect':
- conv2 += [nn.ReflectionPad2d(1)]
- elif padding_type == 'replicate':
- conv2 += [nn.ReplicationPad2d(1)]
- elif padding_type == 'zero':
- p = 1
- else:
- raise NotImplementedError('padding [%s] is not implemented' % padding_type)
- conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
- self.conv2 = nn.Sequential(*conv2)
- self.style2 = ApplyStyle(latent_size, dim)
- def forward(self, x, dlatents_in_slice):
- y = self.conv1(x)
- y = self.style1(y, dlatents_in_slice)
- y = self.act1(y)
- y = self.conv2(y)
- y = self.style2(y, dlatents_in_slice)
- out = x + y
- return out
- class UpBlock_Adain(nn.Module):
- def __init__(self, dim_in, dim_out, latent_size, padding_type, activation=nn.ReLU(True)):
- super(UpBlock_Adain, self).__init__()
- p = 0
- conv1 = [nn.Upsample(scale_factor=2, mode='bilinear')]
- if padding_type == 'reflect':
- conv1 += [nn.ReflectionPad2d(1)]
- elif padding_type == 'replicate':
- conv1 += [nn.ReplicationPad2d(1)]
- elif padding_type == 'zero':
- p = 1
- else:
- raise NotImplementedError('padding [%s] is not implemented' % padding_type)
- conv1 += [nn.Conv2d(dim_in, dim_out, kernel_size=3, padding = p), InstanceNorm()]
- self.conv1 = nn.Sequential(*conv1)
- self.style1 = ApplyStyle(latent_size, dim_out)
- self.act1 = activation
- def forward(self, x, dlatents_in_slice):
- y = self.conv1(x)
- y = self.style1(y, dlatents_in_slice)
- y = self.act1(y)
- return y
- class Encoder(nn.Module):
- def __init__(self, input_nc, output_nc, ngf=32, n_downsampling=4, norm_layer=nn.BatchNorm2d):
- super(Encoder, self).__init__()
- self.output_nc = output_nc
- model = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0),
- norm_layer(ngf), nn.ReLU(True)]
- ### downsample
- for i in range(n_downsampling):
- mult = 2**i
- model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), nn.ReLU(True)]
- ### upsample
- for i in range(n_downsampling):
- mult = 2**(n_downsampling - i)
- model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1, output_padding=1),
- norm_layer(int(ngf * mult / 2)), nn.ReLU(True)]
- model += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- self.model = nn.Sequential(*model)
- def forward(self, input, inst):
- outputs = self.model(input)
- # instance-wise average pooling
- outputs_mean = outputs.clone()
- inst_list = np.unique(inst.cpu().numpy().astype(int))
- for i in inst_list:
- for b in range(input.size()[0]):
- indices = (inst[b:b+1] == int(i)).nonzero() # n x 4
- for j in range(self.output_nc):
- output_ins = outputs[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]]
- mean_feat = torch.mean(output_ins).expand_as(output_ins)
- outputs_mean[indices[:,0] + b, indices[:,1] + j, indices[:,2], indices[:,3]] = mean_feat
- return outputs_mean
- class Generator_Adain(nn.Module):
- def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
- padding_type='reflect'):
- assert (n_blocks >= 0)
- super(Generator_Adain, self).__init__()
- activation = nn.ReLU(True)
- Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
- ### downsample
- for i in range(n_downsampling):
- mult = 2 ** i
- Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), activation]
- self.Encoder = nn.Sequential(*Enc)
- ### resnet blocks
- BN = []
- mult = 2 ** n_downsampling
- for i in range(n_blocks):
- BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
- self.BottleNeck = nn.Sequential(*BN)
- '''self.ResBlockAdain1 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
- activation=activation)
- self.ResBlockAdain2 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
- activation=activation)
- self.ResBlockAdain3 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
- activation=activation)
- self.ResBlockAdain4 = ResnetBlock_Adain(ngf * mult, latent_size=latent_size, padding_type=padding_type,
- activation=activation)'''
- ### upsample
- Dec = []
- for i in range(n_downsampling):
- mult = 2 ** (n_downsampling - i)
- Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
- output_padding=1),
- norm_layer(int(ngf * mult / 2)), activation]
- Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- self.Decoder = nn.Sequential(*Dec)
- #self.model = nn.Sequential(*model)
- self.spNorm = SpecificNorm()
- def forward(self, input, dlatents):
- x = input
- x = self.Encoder(x)
- for i in range(len(self.BottleNeck)):
- x = self.BottleNeck[i](x, dlatents)
- '''x = self.ResBlockAdain1(x, dlatents)
- x = self.ResBlockAdain2(x, dlatents)
- x = self.ResBlockAdain3(x, dlatents)
- x = self.ResBlockAdain4(x, dlatents)'''
- x = self.Decoder(x)
- x = (x + 1) / 2
- x = self.spNorm(x)
- return x
- class Generator_Adain_Mask(nn.Module):
- def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
- padding_type='reflect'):
- assert (n_blocks >= 0)
- super(Generator_Adain_Mask, self).__init__()
- activation = nn.ReLU(True)
- Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
- ### downsample
- for i in range(n_downsampling):
- mult = 2 ** i
- Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), activation]
- self.Encoder = nn.Sequential(*Enc)
- ### resnet blocks
- BN = []
- mult = 2 ** n_downsampling
- for i in range(n_blocks):
- BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
- self.BottleNeck = nn.Sequential(*BN)
- ### upsample
- Dec = []
- for i in range(n_downsampling):
- mult = 2 ** (n_downsampling - i)
- Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
- output_padding=1),
- norm_layer(int(ngf * mult / 2)), activation]
- Fake_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- Mast_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, 1, kernel_size=7, padding=0), nn.Sigmoid()]
- self.Decoder = nn.Sequential(*Dec)
- #self.model = nn.Sequential(*model)
- self.spNorm = SpecificNorm()
- self.Fake_out = nn.Sequential(*Fake_out)
- self.Mask_out = nn.Sequential(*Mast_out)
- def forward(self, input, dlatents):
- x = input
- x = self.Encoder(x)
- for i in range(len(self.BottleNeck)):
- x = self.BottleNeck[i](x, dlatents)
- x = self.Decoder(x)
- fake_out = self.Fake_out(x)
- mask_out = self.Mask_out(x)
- fake_out = (fake_out + 1) / 2
- fake_out = self.spNorm(fake_out)
- generated = fake_out * mask_out + input * (1-mask_out)
- return generated, mask_out
- class Generator_Adain_Upsample(nn.Module):
- def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
- padding_type='reflect'):
- assert (n_blocks >= 0)
- super(Generator_Adain_Upsample, self).__init__()
- activation = nn.ReLU(True)
- Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
- ### downsample
- for i in range(n_downsampling):
- mult = 2 ** i
- Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), activation]
- self.Encoder = nn.Sequential(*Enc)
- ### resnet blocks
- BN = []
- mult = 2 ** n_downsampling
- for i in range(n_blocks):
- BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
- self.BottleNeck = nn.Sequential(*BN)
- ### upsample
- Dec = []
- for i in range(n_downsampling):
- mult = 2 ** (n_downsampling - i)
- '''Dec += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=2, padding=1,
- output_padding=1),
- norm_layer(int(ngf * mult / 2)), activation]'''
- Dec += [nn.Upsample(scale_factor=2, mode='bilinear'),
- nn.Conv2d(ngf * mult, int(ngf * mult / 2), kernel_size=3, stride=1, padding=1),
- norm_layer(int(ngf * mult / 2)), activation]
- Dec += [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- self.Decoder = nn.Sequential(*Dec)
- self.spNorm = SpecificNorm()
- def forward(self, input, dlatents):
- x = input
- x = self.Encoder(x)
- for i in range(len(self.BottleNeck)):
- x = self.BottleNeck[i](x, dlatents)
- x = self.Decoder(x)
- x = (x + 1) / 2
- x = self.spNorm(x)
- return x
- class Generator_Adain_2(nn.Module):
- def __init__(self, input_nc, output_nc, latent_size, ngf=64, n_downsampling=2, n_blocks=4, norm_layer=nn.BatchNorm2d,
- padding_type='reflect'):
- assert (n_blocks >= 0)
- super(Generator_Adain_2, self).__init__()
- activation = nn.ReLU(True)
- Enc = [nn.ReflectionPad2d(3), nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0), norm_layer(ngf), activation]
- ### downsample
- for i in range(n_downsampling):
- mult = 2 ** i
- Enc += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1),
- norm_layer(ngf * mult * 2), activation]
- self.Encoder = nn.Sequential(*Enc)
- ### resnet blocks
- BN = []
- mult = 2 ** n_downsampling
- for i in range(n_blocks):
- BN += [ResnetBlock_Adain(ngf*mult, latent_size=latent_size, padding_type=padding_type, activation=activation)]
- self.BottleNeck = nn.Sequential(*BN)
- ### upsample
- Dec = []
- for i in range(n_downsampling):
- mult = 2 ** (n_downsampling - i)
- Dec += [UpBlock_Adain(dim_in=ngf * mult, dim_out=int(ngf * mult / 2), latent_size=latent_size, padding_type=padding_type)]
- layer_out = [nn.ReflectionPad2d(3), nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0), nn.Tanh()]
- self.Decoder = nn.Sequential(*Dec)
- #self.model = nn.Sequential(*model)
- self.spNorm = SpecificNorm()
- self.layer_out = nn.Sequential(*layer_out)
- def forward(self, input, dlatents):
- x = input
- x = self.Encoder(x)
- for i in range(len(self.BottleNeck)):
- x = self.BottleNeck[i](x, dlatents)
- for i in range(len(self.Decoder)):
- x = self.Decoder[i](x, dlatents)
- x = self.layer_out(x)
- x = (x + 1) / 2
- x = self.spNorm(x)
- return x
- class MultiscaleDiscriminator(nn.Module):
- def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d,
- use_sigmoid=False, num_D=3, getIntermFeat=False):
- super(MultiscaleDiscriminator, self).__init__()
- self.num_D = num_D
- self.n_layers = n_layers
- self.getIntermFeat = getIntermFeat
-
- for i in range(num_D):
- netD = NLayerDiscriminator(input_nc, ndf, n_layers, norm_layer, use_sigmoid, getIntermFeat)
- if getIntermFeat:
- for j in range(n_layers+2):
- setattr(self, 'scale'+str(i)+'_layer'+str(j), getattr(netD, 'model'+str(j)))
- else:
- setattr(self, 'layer'+str(i), netD.model)
- self.downsample = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
- def singleD_forward(self, model, input):
- if self.getIntermFeat:
- result = [input]
- for i in range(len(model)):
- result.append(model[i](result[-1]))
- return result[1:]
- else:
- return [model(input)]
- def forward(self, input):
- num_D = self.num_D
- result = []
- input_downsampled = input
- for i in range(num_D):
- if self.getIntermFeat:
- model = [getattr(self, 'scale'+str(num_D-1-i)+'_layer'+str(j)) for j in range(self.n_layers+2)]
- else:
- model = getattr(self, 'layer'+str(num_D-1-i))
- result.append(self.singleD_forward(model, input_downsampled))
- if i != (num_D-1):
- input_downsampled = self.downsample(input_downsampled)
- return result
-
- # Defines the PatchGAN discriminator with the specified arguments.
- class NLayerDiscriminator(nn.Module):
- def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, use_sigmoid=False, getIntermFeat=False):
- super(NLayerDiscriminator, self).__init__()
- self.getIntermFeat = getIntermFeat
- self.n_layers = n_layers
- kw = 4
- padw = 1
- sequence = [[nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]]
- nf = ndf
- for n in range(1, n_layers):
- nf_prev = nf
- nf = min(nf * 2, 512)
- sequence += [[
- nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=2, padding=padw),
- norm_layer(nf), nn.LeakyReLU(0.2, True)
- ]]
- nf_prev = nf
- nf = min(nf * 2, 512)
- sequence += [[
- nn.Conv2d(nf_prev, nf, kernel_size=kw, stride=1, padding=padw),
- norm_layer(nf),
- nn.LeakyReLU(0.2, True)
- ]]
- if use_sigmoid:
- sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()]]
- else:
- sequence += [[nn.Conv2d(nf, 1, kernel_size=kw, stride=1, padding=padw)]]
- if getIntermFeat:
- for n in range(len(sequence)):
- setattr(self, 'model'+str(n), nn.Sequential(*sequence[n]))
- else:
- sequence_stream = []
- for n in range(len(sequence)):
- sequence_stream += sequence[n]
- self.model = nn.Sequential(*sequence_stream)
- def forward(self, input):
- if self.getIntermFeat:
- res = [input]
- for n in range(self.n_layers+2):
- model = getattr(self, 'model'+str(n))
- res.append(model(res[-1]))
- return res[1:]
- else:
- return self.model(input)
- from torchvision import models
- class Vgg19(torch.nn.Module):
- def __init__(self, requires_grad=False):
- super(Vgg19, self).__init__()
- vgg_pretrained_features = models.vgg19(pretrained=True).features
- self.slice1 = torch.nn.Sequential()
- self.slice2 = torch.nn.Sequential()
- self.slice3 = torch.nn.Sequential()
- self.slice4 = torch.nn.Sequential()
- self.slice5 = torch.nn.Sequential()
- for x in range(2):
- self.slice1.add_module(str(x), vgg_pretrained_features[x])
- for x in range(2, 7):
- self.slice2.add_module(str(x), vgg_pretrained_features[x])
- for x in range(7, 12):
- self.slice3.add_module(str(x), vgg_pretrained_features[x])
- for x in range(12, 21):
- self.slice4.add_module(str(x), vgg_pretrained_features[x])
- for x in range(21, 30):
- self.slice5.add_module(str(x), vgg_pretrained_features[x])
- if not requires_grad:
- for param in self.parameters():
- param.requires_grad = False
- def forward(self, X):
- h_relu1 = self.slice1(X)
- h_relu2 = self.slice2(h_relu1)
- h_relu3 = self.slice3(h_relu2)
- h_relu4 = self.slice4(h_relu3)
- h_relu5 = self.slice5(h_relu4)
- out = [h_relu1, h_relu2, h_relu3, h_relu4, h_relu5]
- return out
|