| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845 | import torchimport torch.nn as nnimport functoolsfrom torch.autograd import Variableimport numpy as npfrom torchvision import transformsimport 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_layerdef 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 netGdef 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 netGdef 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 netGdef 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 netGdef 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 netGdef 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 netDdef 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_prevclass 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 blockclass 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 outclass 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 * tmpclass 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 xclass 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 xclass 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 outclass 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 yclass 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_meanclass 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 xclass 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_outclass 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 xclass 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 xclass 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 modelsclass 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
 |