init

models/__init__.py

@@ -0,0 +1,4 @@
+from .models import ArcMarginModel
+from .models import ResNet
+from .models import IRBlock
+from .models import SEBlock

models/base_model.py

@@ -0,0 +1,91 @@
+import os
+import torch
+import sys
+
+class BaseModel(torch.nn.Module):
+    def name(self):
+        return 'BaseModel'
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.isTrain = opt.isTrain
+        self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+
+    def set_input(self, input):
+        self.input = input
+
+    def forward(self):
+        pass
+
+    # used in test time, no backprop
+    def test(self):
+        pass
+
+    def get_image_paths(self):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        return self.input
+
+    def get_current_errors(self):
+        return {}
+
+    def save(self, label):
+        pass
+
+    # helper saving function that can be used by subclasses
+    def save_network(self, network, network_label, epoch_label, gpu_ids):
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        save_path = os.path.join(self.save_dir, save_filename)
+        torch.save(network.cpu().state_dict(), save_path)
+        if len(gpu_ids) and torch.cuda.is_available():
+            network.cuda()
+
+    # helper loading function that can be used by subclasses
+    def load_network(self, network, network_label, epoch_label, save_dir=''):        
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        if not save_dir:
+            save_dir = self.save_dir
+        save_path = os.path.join(save_dir, save_filename)        
+        if not os.path.isfile(save_path):
+            print('%s not exists yet!' % save_path)
+            if network_label == 'G':
+                raise('Generator must exist!')
+        else:
+            #network.load_state_dict(torch.load(save_path))
+            try:
+                network.load_state_dict(torch.load(save_path))
+            except:   
+                pretrained_dict = torch.load(save_path)                
+                model_dict = network.state_dict()
+                try:
+                    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}                    
+                    network.load_state_dict(pretrained_dict)
+                    if self.opt.verbose:
+                        print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
+                except:
+                    print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
+                    for k, v in pretrained_dict.items():                      
+                        if v.size() == model_dict[k].size():
+                            model_dict[k] = v
+
+                    if sys.version_info >= (3,0):
+                        not_initialized = set()
+                    else:
+                        from sets import Set
+                        not_initialized = Set()                    
+
+                    for k, v in model_dict.items():
+                        if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
+                            not_initialized.add(k.split('.')[0])
+                    
+                    print(sorted(not_initialized))
+                    network.load_state_dict(model_dict)                  
+
+    def update_learning_rate():
+        pass

models/config.py

@@ -0,0 +1,28 @@
+import os
+
+import torch
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # sets device for model and PyTorch tensors
+
+# Model parameters
+image_w = 112
+image_h = 112
+channel = 3
+emb_size = 512
+
+# Training parameters
+num_workers = 1  # for data-loading; right now, only 1 works with h5py
+grad_clip = 5.  # clip gradients at an absolute value of
+print_freq = 100  # print training/validation stats  every __ batches
+checkpoint = None  # path to checkpoint, None if none
+
+# Data parameters
+num_classes = 93431
+num_samples = 5179510
+DATA_DIR = 'data'
+# faces_ms1m_folder = 'data/faces_ms1m_112x112'
+faces_ms1m_folder = 'data/ms1m-retinaface-t1'
+path_imgidx = os.path.join(faces_ms1m_folder, 'train.idx')
+path_imgrec = os.path.join(faces_ms1m_folder, 'train.rec')
+IMG_DIR = 'data/images'
+pickle_file = 'data/faces_ms1m_112x112.pickle'

models/fs_model.py

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

models/fs_networks.py

@@ -0,0 +1,215 @@
+"""
+Copyright (C) 2019 NVIDIA Corporation.  All rights reserved.
+Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
+"""
+
+import torch
+import torch.nn as nn
+
+
+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 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]
+        x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 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 Generator_Adain_Upsample(nn.Module):
+    def __init__(self, input_nc, output_nc, latent_size, n_blocks=6, deep=False,
+                 norm_layer=nn.BatchNorm2d,
+                 padding_type='reflect'):
+        assert (n_blocks >= 0)
+        super(Generator_Adain_Upsample, self).__init__()
+        activation = nn.ReLU(True)
+        self.deep = deep
+
+        self.first_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(input_nc, 64, kernel_size=7, padding=0),
+                                         norm_layer(64), activation)
+        ### downsample
+        self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
+                                   norm_layer(128), activation)
+        self.down2 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1),
+                                   norm_layer(256), activation)
+        self.down3 = nn.Sequential(nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1),
+                                   norm_layer(512), activation)
+        if self.deep:
+            self.down4 = nn.Sequential(nn.Conv2d(512, 512, kernel_size=3, stride=2, padding=1),
+                                       norm_layer(512), activation)
+
+        ### resnet blocks
+        BN = []
+        for i in range(n_blocks):
+            BN += [
+                ResnetBlock_Adain(512, latent_size=latent_size, padding_type=padding_type, activation=activation)]
+        self.BottleNeck = nn.Sequential(*BN)
+
+        if self.deep:
+            self.up4 = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode='bilinear'),
+                nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1),
+                nn.BatchNorm2d(512), activation
+            )
+        self.up3 = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            nn.Conv2d(512, 256, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(256), activation
+        )
+        self.up2 = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(128), activation
+        )
+        self.up1 = nn.Sequential(
+            nn.Upsample(scale_factor=2, mode='bilinear'),
+            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
+            nn.BatchNorm2d(64), activation
+        )
+        self.last_layer = nn.Sequential(nn.ReflectionPad2d(3), nn.Conv2d(64, output_nc, kernel_size=7, padding=0),
+                                        nn.Tanh())
+
+    def forward(self, input, dlatents):
+        x = input  # 3*224*224
+
+        skip1 = self.first_layer(x)
+        skip2 = self.down1(skip1)
+        skip3 = self.down2(skip2)
+        if self.deep:
+            skip4 = self.down3(skip3)
+            x = self.down4(skip4)
+        else:
+            x = self.down3(skip3)
+
+        for i in range(len(self.BottleNeck)):
+            x = self.BottleNeck[i](x, dlatents)
+
+        if self.deep:
+            x = self.up4(x)
+        x = self.up3(x)
+        x = self.up2(x)
+        x = self.up1(x)
+        x = self.last_layer(x)
+        x = (x + 1) / 2
+
+        return x
+
+class Discriminator(nn.Module):
+    def __init__(self, input_nc, norm_layer=nn.BatchNorm2d, use_sigmoid=False):
+        super(Discriminator, self).__init__()
+
+        kw = 4
+        padw = 1
+        self.down1 = nn.Sequential(
+            nn.Conv2d(input_nc, 64, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)
+        )
+        self.down2 = nn.Sequential(
+            nn.Conv2d(64, 128, kernel_size=kw, stride=2, padding=padw),
+            norm_layer(128), nn.LeakyReLU(0.2, True)
+        )
+        self.down3 = nn.Sequential(
+            nn.Conv2d(128, 256, kernel_size=kw, stride=2, padding=padw),
+            norm_layer(256), nn.LeakyReLU(0.2, True)
+        )
+        self.down4 = nn.Sequential(
+            nn.Conv2d(256, 512, kernel_size=kw, stride=2, padding=padw),
+            norm_layer(512), nn.LeakyReLU(0.2, True)
+        )
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(512, 512, kernel_size=kw, stride=1, padding=padw),
+            norm_layer(512),
+            nn.LeakyReLU(0.2, True)
+        )
+
+        if use_sigmoid:
+            self.conv2 = nn.Sequential(
+                nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw), nn.Sigmoid()
+            )
+        else:
+            self.conv2 = nn.Sequential(
+                nn.Conv2d(512, 1, kernel_size=kw, stride=1, padding=padw)
+            )
+
+    def forward(self, input):
+        out = []
+        x = self.down1(input)
+        out.append(x)
+        x = self.down2(x)
+        out.append(x)
+        x = self.down3(x)
+        out.append(x)
+        x = self.down4(x)
+        out.append(x)
+        x = self.conv1(x)
+        out.append(x)
+        x = self.conv2(x)
+        out.append(x)
+        
+        return out

models/models.py

@@ -0,0 +1,181 @@
+import math
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import Parameter
+from .config import device, num_classes
+
+
+def create_model(opt):
+    if opt.model == 'pix2pixHD':
+        #from .pix2pixHD_model import Pix2PixHDModel, InferenceModel
+        from .fs_model import fsModel
+        model = fsModel()
+    else:
+        from .ui_model import UIModel
+        model = UIModel()
+
+    model.initialize(opt)
+    if opt.verbose:
+        print("model [%s] was created" % (model.name()))
+
+    if opt.isTrain and len(opt.gpu_ids) and not opt.fp16:
+        model = torch.nn.DataParallel(model, device_ids=opt.gpu_ids)
+
+    return model
+
+
+
+class SEBlock(nn.Module):
+    def __init__(self, channel, reduction=16):
+        super(SEBlock, self).__init__()
+        self.avg_pool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Sequential(
+            nn.Linear(channel, channel // reduction),
+            nn.PReLU(),
+            nn.Linear(channel // reduction, channel),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        b, c, _, _ = x.size()
+        y = self.avg_pool(x).view(b, c)
+        y = self.fc(y).view(b, c, 1, 1)
+        return x * y
+
+
+class IRBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, use_se=True):
+        super(IRBlock, self).__init__()
+        self.bn0 = nn.BatchNorm2d(inplanes)
+        self.conv1 = conv3x3(inplanes, inplanes)
+        self.bn1 = nn.BatchNorm2d(inplanes)
+        self.prelu = nn.PReLU()
+        self.conv2 = conv3x3(inplanes, planes, stride)
+        self.bn2 = nn.BatchNorm2d(planes)
+        self.downsample = downsample
+        self.stride = stride
+        self.use_se = use_se
+        if self.use_se:
+            self.se = SEBlock(planes)
+
+    def forward(self, x):
+        residual = x
+        out = self.bn0(x)
+        out = self.conv1(out)
+        out = self.bn1(out)
+        out = self.prelu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        if self.use_se:
+            out = self.se(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.prelu(out)
+
+        return out
+
+
+class ResNet(nn.Module):
+
+    def __init__(self, block, layers, use_se=True):
+        self.inplanes = 64
+        self.use_se = use_se
+        super(ResNet, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.prelu = nn.PReLU()
+        self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.bn2 = nn.BatchNorm2d(512)
+        self.dropout = nn.Dropout()
+        self.fc = nn.Linear(512 * 7 * 7, 512)
+        self.bn3 = nn.BatchNorm1d(512)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.xavier_normal_(m.weight)
+            elif isinstance(m, nn.BatchNorm2d) or isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.xavier_normal_(m.weight)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, use_se=self.use_se))
+        self.inplanes = planes
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, use_se=self.use_se))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.prelu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.bn2(x)
+        x = self.dropout(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+        x = self.bn3(x)
+
+        return x
+
+
+class ArcMarginModel(nn.Module):
+    def __init__(self, args):
+        super(ArcMarginModel, self).__init__()
+
+        self.weight = Parameter(torch.FloatTensor(num_classes, args.emb_size))
+        nn.init.xavier_uniform_(self.weight)
+
+        self.easy_margin = args.easy_margin
+        self.m = args.margin_m
+        self.s = args.margin_s
+
+        self.cos_m = math.cos(self.m)
+        self.sin_m = math.sin(self.m)
+        self.th = math.cos(math.pi - self.m)
+        self.mm = math.sin(math.pi - self.m) * self.m
+
+    def forward(self, input, label):
+        x = F.normalize(input)
+        W = F.normalize(self.weight)
+        cosine = F.linear(x, W)
+        sine = torch.sqrt(1.0 - torch.pow(cosine, 2))
+        phi = cosine * self.cos_m - sine * self.sin_m  # cos(theta + m)
+        if self.easy_margin:
+            phi = torch.where(cosine > 0, phi, cosine)
+        else:
+            phi = torch.where(cosine > self.th, phi, cosine - self.mm)
+        one_hot = torch.zeros(cosine.size(), device=device)
+        one_hot.scatter_(1, label.view(-1, 1).long(), 1)
+        output = (one_hot * phi) + ((1.0 - one_hot) * cosine)
+        output *= self.s
+        return output

models/networks.py

@@ -0,0 +1,845 @@
+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

models/pix2pixHD_model.py

@@ -0,0 +1,304 @@
+import numpy as np
+import torch
+import os
+from torch.autograd import Variable
+from util.image_pool import ImagePool
+from .base_model import BaseModel
+from . import networks
+
+class Pix2PixHDModel(BaseModel):
+    def name(self):
+        return 'Pix2PixHDModel'
+    
+    def init_loss_filter(self, use_gan_feat_loss, use_vgg_loss):
+        flags = (True, use_gan_feat_loss, use_vgg_loss, True, True)
+        def loss_filter(g_gan, g_gan_feat, g_vgg, d_real, d_fake):
+            return [l for (l,f) in zip((g_gan,g_gan_feat,g_vgg,d_real,d_fake),flags) if f]
+        return loss_filter
+    
+    def initialize(self, opt):
+        BaseModel.initialize(self, opt)
+        if opt.resize_or_crop != 'none' or not opt.isTrain: # when training at full res this causes OOM
+            torch.backends.cudnn.benchmark = True
+        self.isTrain = opt.isTrain
+        self.use_features = opt.instance_feat or opt.label_feat
+        self.gen_features = self.use_features and not self.opt.load_features
+        input_nc = opt.label_nc if opt.label_nc != 0 else opt.input_nc
+
+        ##### define networks        
+        # Generator network
+        netG_input_nc = input_nc        
+        if not opt.no_instance:
+            netG_input_nc += 1
+        if self.use_features:
+            netG_input_nc += opt.feat_num                  
+        self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG, 
+                                      opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 
+                                      opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)        
+
+        # Discriminator network
+        if self.isTrain:
+            use_sigmoid = opt.no_lsgan
+            netD_input_nc = input_nc + opt.output_nc
+            if not opt.no_instance:
+                netD_input_nc += 1
+            self.netD = networks.define_D(netD_input_nc, opt.ndf, opt.n_layers_D, opt.norm, use_sigmoid, 
+                                          opt.num_D, not opt.no_ganFeat_loss, gpu_ids=self.gpu_ids)
+
+        ### Encoder network
+        if self.gen_features:          
+            self.netE = networks.define_G(opt.output_nc, opt.feat_num, opt.nef, 'encoder', 
+                                          opt.n_downsample_E, norm=opt.norm, gpu_ids=self.gpu_ids)  
+        if self.opt.verbose:
+                print('---------- Networks initialized -------------')
+
+        # load networks
+        if not self.isTrain or opt.continue_train or opt.load_pretrain:
+            pretrained_path = '' if not self.isTrain else opt.load_pretrain
+            self.load_network(self.netG, 'G', opt.which_epoch, pretrained_path)            
+            if self.isTrain:
+                self.load_network(self.netD, 'D', opt.which_epoch, pretrained_path)  
+            if self.gen_features:
+                self.load_network(self.netE, 'E', opt.which_epoch, pretrained_path)              
+
+        # set loss functions and optimizers
+        if self.isTrain:
+            if opt.pool_size > 0 and (len(self.gpu_ids)) > 1:
+                raise NotImplementedError("Fake Pool Not Implemented for MultiGPU")
+            self.fake_pool = ImagePool(opt.pool_size)
+            self.old_lr = opt.lr
+
+            # define loss functions
+            self.loss_filter = self.init_loss_filter(not opt.no_ganFeat_loss, not opt.no_vgg_loss)
+            
+            self.criterionGAN = networks.GANLoss(use_lsgan=not opt.no_lsgan, tensor=self.Tensor)   
+            self.criterionFeat = torch.nn.L1Loss()
+            if not opt.no_vgg_loss:             
+                self.criterionVGG = networks.VGGLoss(self.gpu_ids)
+                
+        
+            # Names so we can breakout loss
+            self.loss_names = self.loss_filter('G_GAN','G_GAN_Feat','G_VGG','D_real', 'D_fake')
+
+            # initialize optimizers
+            # optimizer G
+            if opt.niter_fix_global > 0:                
+                import sys
+                if sys.version_info >= (3,0):
+                    finetune_list = set()
+                else:
+                    from sets import Set
+                    finetune_list = Set()
+
+                params_dict = dict(self.netG.named_parameters())
+                params = []
+                for key, value in params_dict.items():       
+                    if key.startswith('model' + str(opt.n_local_enhancers)):                    
+                        params += [value]
+                        finetune_list.add(key.split('.')[0])  
+                print('------------- Only training the local enhancer network (for %d epochs) ------------' % opt.niter_fix_global)
+                print('The layers that are finetuned are ', sorted(finetune_list))                         
+            else:
+                params = list(self.netG.parameters())
+            if self.gen_features:              
+                params += list(self.netE.parameters())         
+            self.optimizer_G = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))                            
+
+            # optimizer D                        
+            params = list(self.netD.parameters())    
+            self.optimizer_D = torch.optim.Adam(params, lr=opt.lr, betas=(opt.beta1, 0.999))
+
+    def encode_input(self, label_map, inst_map=None, real_image=None, feat_map=None, infer=False):             
+        if self.opt.label_nc == 0:
+            input_label = label_map.data.cuda()
+        else:
+            # create one-hot vector for label map 
+            size = label_map.size()
+            oneHot_size = (size[0], self.opt.label_nc, size[2], size[3])
+            input_label = torch.cuda.FloatTensor(torch.Size(oneHot_size)).zero_()
+            input_label = input_label.scatter_(1, label_map.data.long().cuda(), 1.0)
+            if self.opt.data_type == 16:
+                input_label = input_label.half()
+
+        # get edges from instance map
+        if not self.opt.no_instance:
+            inst_map = inst_map.data.cuda()
+            edge_map = self.get_edges(inst_map)
+            input_label = torch.cat((input_label, edge_map), dim=1)         
+        input_label = Variable(input_label, volatile=infer)
+
+        # real images for training
+        if real_image is not None:
+            real_image = Variable(real_image.data.cuda())
+
+        # instance map for feature encoding
+        if self.use_features:
+            # get precomputed feature maps
+            if self.opt.load_features:
+                feat_map = Variable(feat_map.data.cuda())
+            if self.opt.label_feat:
+                inst_map = label_map.cuda()
+
+        return input_label, inst_map, real_image, feat_map
+
+    def discriminate(self, input_label, test_image, use_pool=False):
+        input_concat = torch.cat((input_label, test_image.detach()), dim=1)
+        if use_pool:            
+            fake_query = self.fake_pool.query(input_concat)
+            return self.netD.forward(fake_query)
+        else:
+            return self.netD.forward(input_concat)
+
+    def forward(self, label, inst, image, feat, infer=False):
+        # Encode Inputs
+        input_label, inst_map, real_image, feat_map = self.encode_input(label, inst, image, feat)  
+
+        # Fake Generation
+        if self.use_features:
+            if not self.opt.load_features:
+                feat_map = self.netE.forward(real_image, inst_map)                     
+            input_concat = torch.cat((input_label, feat_map), dim=1)                        
+        else:
+            input_concat = input_label
+        fake_image = self.netG.forward(input_concat)
+
+        # Fake Detection and Loss
+        pred_fake_pool = self.discriminate(input_label, fake_image, use_pool=True)
+        loss_D_fake = self.criterionGAN(pred_fake_pool, False)        
+
+        # Real Detection and Loss        
+        pred_real = self.discriminate(input_label, real_image)
+        loss_D_real = self.criterionGAN(pred_real, True)
+
+        # GAN loss (Fake Passability Loss)        
+        pred_fake = self.netD.forward(torch.cat((input_label, fake_image), dim=1))        
+        loss_G_GAN = self.criterionGAN(pred_fake, True)               
+        
+        # GAN feature matching loss
+        loss_G_GAN_Feat = 0
+        if not self.opt.no_ganFeat_loss:
+            feat_weights = 4.0 / (self.opt.n_layers_D + 1)
+            D_weights = 1.0 / self.opt.num_D
+            for i in range(self.opt.num_D):
+                for j in range(len(pred_fake[i])-1):
+                    loss_G_GAN_Feat += D_weights * feat_weights * \
+                        self.criterionFeat(pred_fake[i][j], pred_real[i][j].detach()) * self.opt.lambda_feat
+                   
+        # VGG feature matching loss
+        loss_G_VGG = 0
+        if not self.opt.no_vgg_loss:
+            loss_G_VGG = self.criterionVGG(fake_image, real_image) * self.opt.lambda_feat
+        
+        # Only return the fake_B image if necessary to save BW
+        return [ self.loss_filter( loss_G_GAN, loss_G_GAN_Feat, loss_G_VGG, loss_D_real, loss_D_fake ), None if not infer else fake_image ]
+
+    def inference(self, label, inst, image=None):
+        # Encode Inputs        
+        image = Variable(image) if image is not None else None
+        input_label, inst_map, real_image, _ = self.encode_input(Variable(label), Variable(inst), image, infer=True)
+
+        # Fake Generation
+        if self.use_features:
+            if self.opt.use_encoded_image:
+                # encode the real image to get feature map
+                feat_map = self.netE.forward(real_image, inst_map)
+            else:
+                # sample clusters from precomputed features             
+                feat_map = self.sample_features(inst_map)
+            input_concat = torch.cat((input_label, feat_map), dim=1)                        
+        else:
+            input_concat = input_label        
+           
+        if torch.__version__.startswith('0.4'):
+            with torch.no_grad():
+                fake_image = self.netG.forward(input_concat)
+        else:
+            fake_image = self.netG.forward(input_concat)
+        return fake_image
+
+    def sample_features(self, inst): 
+        # read precomputed feature clusters 
+        cluster_path = os.path.join(self.opt.checkpoints_dir, self.opt.name, self.opt.cluster_path)        
+        features_clustered = np.load(cluster_path, encoding='latin1').item()
+
+        # randomly sample from the feature clusters
+        inst_np = inst.cpu().numpy().astype(int)                                      
+        feat_map = self.Tensor(inst.size()[0], self.opt.feat_num, inst.size()[2], inst.size()[3])
+        for i in np.unique(inst_np):    
+            label = i if i < 1000 else i//1000
+            if label in features_clustered:
+                feat = features_clustered[label]
+                cluster_idx = np.random.randint(0, feat.shape[0]) 
+                                            
+                idx = (inst == int(i)).nonzero()
+                for k in range(self.opt.feat_num):                                    
+                    feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k]
+        if self.opt.data_type==16:
+            feat_map = feat_map.half()
+        return feat_map
+
+    def encode_features(self, image, inst):
+        image = Variable(image.cuda(), volatile=True)
+        feat_num = self.opt.feat_num
+        h, w = inst.size()[2], inst.size()[3]
+        block_num = 32
+        feat_map = self.netE.forward(image, inst.cuda())
+        inst_np = inst.cpu().numpy().astype(int)
+        feature = {}
+        for i in range(self.opt.label_nc):
+            feature[i] = np.zeros((0, feat_num+1))
+        for i in np.unique(inst_np):
+            label = i if i < 1000 else i//1000
+            idx = (inst == int(i)).nonzero()
+            num = idx.size()[0]
+            idx = idx[num//2,:]
+            val = np.zeros((1, feat_num+1))                        
+            for k in range(feat_num):
+                val[0, k] = feat_map[idx[0], idx[1] + k, idx[2], idx[3]].data[0]            
+            val[0, feat_num] = float(num) / (h * w // block_num)
+            feature[label] = np.append(feature[label], val, axis=0)
+        return feature
+
+    def get_edges(self, t):
+        edge = torch.cuda.ByteTensor(t.size()).zero_()
+        edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
+        edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
+        edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+        edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+        if self.opt.data_type==16:
+            return edge.half()
+        else:
+            return edge.float()
+
+    def save(self, which_epoch):
+        self.save_network(self.netG, 'G', which_epoch, self.gpu_ids)
+        self.save_network(self.netD, 'D', which_epoch, self.gpu_ids)
+        if self.gen_features:
+            self.save_network(self.netE, 'E', which_epoch, self.gpu_ids)
+
+    def update_fixed_params(self):
+        # after fixing the global generator for a number of iterations, also start finetuning it
+        params = list(self.netG.parameters())
+        if self.gen_features:
+            params += list(self.netE.parameters())           
+        self.optimizer_G = torch.optim.Adam(params, lr=self.opt.lr, betas=(self.opt.beta1, 0.999))
+        if self.opt.verbose:
+            print('------------ Now also finetuning global generator -----------')
+
+    def update_learning_rate(self):
+        lrd = self.opt.lr / self.opt.niter_decay
+        lr = self.old_lr - lrd        
+        for param_group in self.optimizer_D.param_groups:
+            param_group['lr'] = lr
+        for param_group in self.optimizer_G.param_groups:
+            param_group['lr'] = lr
+        if self.opt.verbose:
+            print('update learning rate: %f -> %f' % (self.old_lr, lr))
+        self.old_lr = lr
+
+class InferenceModel(Pix2PixHDModel):
+    def forward(self, inp):
+        label, inst = inp
+        return self.inference(label, inst)
+
+        

models/ui_model.py

@@ -0,0 +1,347 @@
+import torch
+from torch.autograd import Variable
+from collections import OrderedDict
+import numpy as np
+import os
+from PIL import Image
+import util.util as util
+from .base_model import BaseModel
+from . import networks
+
+class UIModel(BaseModel):
+    def name(self):
+        return 'UIModel'
+
+    def initialize(self, opt):
+        assert(not opt.isTrain)
+        BaseModel.initialize(self, opt)
+        self.use_features = opt.instance_feat or opt.label_feat
+
+        netG_input_nc = opt.label_nc
+        if not opt.no_instance:
+            netG_input_nc += 1            
+        if self.use_features:   
+            netG_input_nc += opt.feat_num           
+
+        self.netG = networks.define_G(netG_input_nc, opt.output_nc, opt.ngf, opt.netG, 
+                                      opt.n_downsample_global, opt.n_blocks_global, opt.n_local_enhancers, 
+                                      opt.n_blocks_local, opt.norm, gpu_ids=self.gpu_ids)            
+        self.load_network(self.netG, 'G', opt.which_epoch)
+
+        print('---------- Networks initialized -------------')
+
+    def toTensor(self, img, normalize=False):
+        tensor = torch.from_numpy(np.array(img, np.int32, copy=False))
+        tensor = tensor.view(1, img.size[1], img.size[0], len(img.mode))    
+        tensor = tensor.transpose(1, 2).transpose(1, 3).contiguous()
+        if normalize:
+            return (tensor.float()/255.0 - 0.5) / 0.5        
+        return tensor.float()
+
+    def load_image(self, label_path, inst_path, feat_path):
+        opt = self.opt
+        # read label map
+        label_img = Image.open(label_path)    
+        if label_path.find('face') != -1:
+            label_img = label_img.convert('L')
+        ow, oh = label_img.size    
+        w = opt.loadSize
+        h = int(w * oh / ow)    
+        label_img = label_img.resize((w, h), Image.NEAREST)
+        label_map = self.toTensor(label_img)           
+        
+        # onehot vector input for label map
+        self.label_map = label_map.cuda()
+        oneHot_size = (1, opt.label_nc, h, w)
+        input_label = self.Tensor(torch.Size(oneHot_size)).zero_()
+        self.input_label = input_label.scatter_(1, label_map.long().cuda(), 1.0)
+
+        # read instance map
+        if not opt.no_instance:
+            inst_img = Image.open(inst_path)        
+            inst_img = inst_img.resize((w, h), Image.NEAREST)            
+            self.inst_map = self.toTensor(inst_img).cuda()
+            self.edge_map = self.get_edges(self.inst_map)          
+            self.net_input = Variable(torch.cat((self.input_label, self.edge_map), dim=1), volatile=True)
+        else:
+            self.net_input = Variable(self.input_label, volatile=True)  
+        
+        self.features_clustered = np.load(feat_path).item()
+        self.object_map = self.inst_map if opt.instance_feat else self.label_map 
+                       
+        object_np = self.object_map.cpu().numpy().astype(int) 
+        self.feat_map = self.Tensor(1, opt.feat_num, h, w).zero_()                 
+        self.cluster_indices = np.zeros(self.opt.label_nc, np.uint8)
+        for i in np.unique(object_np):    
+            label = i if i < 1000 else i//1000
+            if label in self.features_clustered:
+                feat = self.features_clustered[label]
+                np.random.seed(i+1)
+                cluster_idx = np.random.randint(0, feat.shape[0])
+                self.cluster_indices[label] = cluster_idx
+                idx = (self.object_map == i).nonzero()                    
+                self.set_features(idx, feat, cluster_idx)
+
+        self.net_input_original = self.net_input.clone()        
+        self.label_map_original = self.label_map.clone()
+        self.feat_map_original = self.feat_map.clone()
+        if not opt.no_instance:
+            self.inst_map_original = self.inst_map.clone()        
+
+    def reset(self):
+        self.net_input = self.net_input_prev = self.net_input_original.clone()        
+        self.label_map = self.label_map_prev = self.label_map_original.clone()
+        self.feat_map = self.feat_map_prev = self.feat_map_original.clone()
+        if not self.opt.no_instance:
+            self.inst_map = self.inst_map_prev = self.inst_map_original.clone()
+        self.object_map = self.inst_map if self.opt.instance_feat else self.label_map 
+
+    def undo(self):        
+        self.net_input = self.net_input_prev
+        self.label_map = self.label_map_prev
+        self.feat_map = self.feat_map_prev
+        if not self.opt.no_instance:
+            self.inst_map = self.inst_map_prev
+        self.object_map = self.inst_map if self.opt.instance_feat else self.label_map 
+            
+    # get boundary map from instance map
+    def get_edges(self, t):
+        edge = torch.cuda.ByteTensor(t.size()).zero_()
+        edge[:,:,:,1:] = edge[:,:,:,1:] | (t[:,:,:,1:] != t[:,:,:,:-1])
+        edge[:,:,:,:-1] = edge[:,:,:,:-1] | (t[:,:,:,1:] != t[:,:,:,:-1])
+        edge[:,:,1:,:] = edge[:,:,1:,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+        edge[:,:,:-1,:] = edge[:,:,:-1,:] | (t[:,:,1:,:] != t[:,:,:-1,:])
+        return edge.float()
+
+    # change the label at the source position to the label at the target position
+    def change_labels(self, click_src, click_tgt): 
+        y_src, x_src = click_src[0], click_src[1]
+        y_tgt, x_tgt = click_tgt[0], click_tgt[1]
+        label_src = int(self.label_map[0, 0, y_src, x_src])
+        inst_src = self.inst_map[0, 0, y_src, x_src]
+        label_tgt = int(self.label_map[0, 0, y_tgt, x_tgt])
+        inst_tgt = self.inst_map[0, 0, y_tgt, x_tgt]
+
+        idx_src = (self.inst_map == inst_src).nonzero()         
+        # need to change 3 things: label map, instance map, and feature map
+        if idx_src.shape:
+            # backup current maps
+            self.backup_current_state() 
+
+            # change both the label map and the network input
+            self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+            self.net_input[idx_src[:,0], idx_src[:,1] + label_src, idx_src[:,2], idx_src[:,3]] = 0
+            self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1                                    
+            
+            # update the instance map (and the network input)
+            if inst_tgt > 1000:
+                # if different instances have different ids, give the new object a new id
+                tgt_indices = (self.inst_map > label_tgt * 1000) & (self.inst_map < (label_tgt+1) * 1000)
+                inst_tgt = self.inst_map[tgt_indices].max() + 1
+            self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = inst_tgt
+            self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+
+            # also copy the source features to the target position      
+            idx_tgt = (self.inst_map == inst_tgt).nonzero()    
+            if idx_tgt.shape:
+                self.copy_features(idx_src, idx_tgt[0,:])
+
+        self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+    # add strokes of target label in the image
+    def add_strokes(self, click_src, label_tgt, bw, save):
+        # get the region of the new strokes (bw is the brush width)        
+        size = self.net_input.size()
+        h, w = size[2], size[3]
+        idx_src = torch.LongTensor(bw**2, 4).fill_(0)
+        for i in range(bw):
+            idx_src[i*bw:(i+1)*bw, 2] = min(h-1, max(0, click_src[0]-bw//2 + i))
+            for j in range(bw):
+                idx_src[i*bw+j, 3] = min(w-1, max(0, click_src[1]-bw//2 + j))
+        idx_src = idx_src.cuda()
+        
+        # again, need to update 3 things
+        if idx_src.shape:
+            # backup current maps
+            if save:
+                self.backup_current_state()
+
+            # update the label map (and the network input) in the stroke region            
+            self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+            for k in range(self.opt.label_nc):
+                self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
+            self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1                 
+
+            # update the instance map (and the network input)
+            self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+            self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+            
+            # also update the features if available
+            if self.opt.instance_feat:                                            
+                feat = self.features_clustered[label_tgt]
+                #np.random.seed(label_tgt+1)   
+                #cluster_idx = np.random.randint(0, feat.shape[0])
+                cluster_idx = self.cluster_indices[label_tgt]
+                self.set_features(idx_src, feat, cluster_idx)                                                  
+        
+        self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+    # add an object to the clicked position with selected style
+    def add_objects(self, click_src, label_tgt, mask, style_id=0):
+        y, x = click_src[0], click_src[1]
+        mask = np.transpose(mask, (2, 0, 1))[np.newaxis,...]        
+        idx_src = torch.from_numpy(mask).cuda().nonzero()        
+        idx_src[:,2] += y
+        idx_src[:,3] += x
+
+        # backup current maps
+        self.backup_current_state()
+
+        # update label map
+        self.label_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt        
+        for k in range(self.opt.label_nc):
+            self.net_input[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = 0
+        self.net_input[idx_src[:,0], idx_src[:,1] + label_tgt, idx_src[:,2], idx_src[:,3]] = 1            
+
+        # update instance map
+        self.inst_map[idx_src[:,0], idx_src[:,1], idx_src[:,2], idx_src[:,3]] = label_tgt
+        self.net_input[:,-1,:,:] = self.get_edges(self.inst_map)
+                
+        # update feature map
+        self.set_features(idx_src, self.feat, style_id)                
+        
+        self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+
+    def single_forward(self, net_input, feat_map):
+        net_input = torch.cat((net_input, feat_map), dim=1)
+        fake_image = self.netG.forward(net_input)
+
+        if fake_image.size()[0] == 1:
+            return fake_image.data[0]        
+        return fake_image.data
+
+
+    # generate all outputs for different styles
+    def style_forward(self, click_pt, style_id=-1):           
+        if click_pt is None:            
+            self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))
+            self.crop = None
+            self.mask = None        
+        else:                       
+            instToChange = int(self.object_map[0, 0, click_pt[0], click_pt[1]])
+            self.instToChange = instToChange
+            label = instToChange if instToChange < 1000 else instToChange//1000        
+            self.feat = self.features_clustered[label]
+            self.fake_image = []
+            self.mask = self.object_map == instToChange
+            idx = self.mask.nonzero()
+            self.get_crop_region(idx)            
+            if idx.size():                
+                if style_id == -1:
+                    (min_y, min_x, max_y, max_x) = self.crop
+                    ### original
+                    for cluster_idx in range(self.opt.multiple_output):
+                        self.set_features(idx, self.feat, cluster_idx)
+                        fake_image = self.single_forward(self.net_input, self.feat_map)
+                        fake_image = util.tensor2im(fake_image[:,min_y:max_y,min_x:max_x])
+                        self.fake_image.append(fake_image)    
+                    """### To speed up previewing different style results, either crop or downsample the label maps
+                    if instToChange > 1000:
+                        (min_y, min_x, max_y, max_x) = self.crop                                                
+                        ### crop                                                
+                        _, _, h, w = self.net_input.size()
+                        offset = 512
+                        y_start, x_start = max(0, min_y-offset), max(0, min_x-offset)
+                        y_end, x_end = min(h, (max_y + offset)), min(w, (max_x + offset))
+                        y_region = slice(y_start, y_start+(y_end-y_start)//16*16)
+                        x_region = slice(x_start, x_start+(x_end-x_start)//16*16)
+                        net_input = self.net_input[:,:,y_region,x_region]                    
+                        for cluster_idx in range(self.opt.multiple_output):  
+                            self.set_features(idx, self.feat, cluster_idx)
+                            fake_image = self.single_forward(net_input, self.feat_map[:,:,y_region,x_region])                            
+                            fake_image = util.tensor2im(fake_image[:,min_y-y_start:max_y-y_start,min_x-x_start:max_x-x_start])
+                            self.fake_image.append(fake_image)
+                    else:
+                        ### downsample
+                        (min_y, min_x, max_y, max_x) = [crop//2 for crop in self.crop]                    
+                        net_input = self.net_input[:,:,::2,::2]                    
+                        size = net_input.size()
+                        net_input_batch = net_input.expand(self.opt.multiple_output, size[1], size[2], size[3])             
+                        for cluster_idx in range(self.opt.multiple_output):  
+                            self.set_features(idx, self.feat, cluster_idx)
+                            feat_map = self.feat_map[:,:,::2,::2]
+                            if cluster_idx == 0:
+                                feat_map_batch = feat_map
+                            else:
+                                feat_map_batch = torch.cat((feat_map_batch, feat_map), dim=0)
+                        fake_image_batch = self.single_forward(net_input_batch, feat_map_batch)
+                        for i in range(self.opt.multiple_output):
+                            self.fake_image.append(util.tensor2im(fake_image_batch[i,:,min_y:max_y,min_x:max_x]))"""
+                                        
+                else:
+                    self.set_features(idx, self.feat, style_id)
+                    self.cluster_indices[label] = style_id
+                    self.fake_image = util.tensor2im(self.single_forward(self.net_input, self.feat_map))        
+
+    def backup_current_state(self):
+        self.net_input_prev = self.net_input.clone()
+        self.label_map_prev = self.label_map.clone() 
+        self.inst_map_prev = self.inst_map.clone() 
+        self.feat_map_prev = self.feat_map.clone() 
+
+    # crop the ROI and get the mask of the object
+    def get_crop_region(self, idx):
+        size = self.net_input.size()
+        h, w = size[2], size[3]
+        min_y, min_x = idx[:,2].min(), idx[:,3].min()
+        max_y, max_x = idx[:,2].max(), idx[:,3].max()             
+        crop_min = 128
+        if max_y - min_y < crop_min:
+            min_y = max(0, (max_y + min_y) // 2 - crop_min // 2)
+            max_y = min(h-1, min_y + crop_min)
+        if max_x - min_x < crop_min:
+            min_x = max(0, (max_x + min_x) // 2 - crop_min // 2)
+            max_x = min(w-1, min_x + crop_min)
+        self.crop = (min_y, min_x, max_y, max_x)           
+        self.mask = self.mask[:,:, min_y:max_y, min_x:max_x]
+
+    # update the feature map once a new object is added or the label is changed
+    def update_features(self, cluster_idx, mask=None, click_pt=None):        
+        self.feat_map_prev = self.feat_map.clone()
+        # adding a new object
+        if mask is not None:
+            y, x = click_pt[0], click_pt[1]
+            mask = np.transpose(mask, (2,0,1))[np.newaxis,...]        
+            idx = torch.from_numpy(mask).cuda().nonzero()        
+            idx[:,2] += y
+            idx[:,3] += x    
+        # changing the label of an existing object 
+        else:            
+            idx = (self.object_map == self.instToChange).nonzero()              
+
+        # update feature map
+        self.set_features(idx, self.feat, cluster_idx)        
+
+    # set the class features to the target feature
+    def set_features(self, idx, feat, cluster_idx):        
+        for k in range(self.opt.feat_num):
+            self.feat_map[idx[:,0], idx[:,1] + k, idx[:,2], idx[:,3]] = feat[cluster_idx, k] 
+
+    # copy the features at the target position to the source position
+    def copy_features(self, idx_src, idx_tgt):        
+        for k in range(self.opt.feat_num):
+            val = self.feat_map[idx_tgt[0], idx_tgt[1] + k, idx_tgt[2], idx_tgt[3]]
+            self.feat_map[idx_src[:,0], idx_src[:,1] + k, idx_src[:,2], idx_src[:,3]] = val 
+
+    def get_current_visuals(self, getLabel=False):                              
+        mask = self.mask     
+        if self.mask is not None:
+            mask = np.transpose(self.mask[0].cpu().float().numpy(), (1,2,0)).astype(np.uint8)        
+
+        dict_list = [('fake_image', self.fake_image), ('mask', mask)]
+
+        if getLabel: # only output label map if needed to save bandwidth
+            label = util.tensor2label(self.net_input.data[0], self.opt.label_nc)                    
+            dict_list += [('label', label)]
+
+        return OrderedDict(dict_list)