update

GUI/file_sync/filestate_machine0.json

@@ -1,7 +1,7 @@
 {
     "GUI.py": 1642351532.4558506,
-    "test.py": 1634039043.4872007,
+    "test.py": 1642733759.8015468,
-    "train.py": 1642408973.489742,
+    "train.py": 1643009568.7253726,
     "components\\Generator.py": 1642347735.351465,
     "components\\projected_discriminator.py": 1642348101.4661522,
     "components\\pg_modules\\blocks.py": 1640773190.0,
@@ -23,7 +23,7 @@
     "test_scripts\\tester_common.py": 1625369535.199175,
     "test_scripts\\tester_FastNST.py": 1634041357.607633,
     "train_scripts\\trainer_base.py": 1642396105.3868554,
-    "train_scripts\\trainer_FM.py": 1642396334.407562,
+    "train_scripts\\trainer_FM.py": 1642826710.3532298,
     "train_scripts\\trainer_naiv512.py": 1642315674.9740853,
     "utilities\\checkpoint_manager.py": 1611123530.6624403,
     "utilities\\figure.py": 1611123530.6634378,
@@ -37,7 +37,23 @@
     "utilities\\transfer_checkpoint.py": 1642397157.0163105,
     "utilities\\utilities.py": 1634019485.0783668,
     "utilities\\yaml_config.py": 1611123530.6614666,
-    "train_yamls\\train_512FM.yaml": 1642408586.2747152,
+    "train_yamls\\train_512FM.yaml": 1642412254.0831068,
-    "train_scripts\\trainer_2layer_FM.py": 1642408727.3950334,
+    "train_scripts\\trainer_2layer_FM.py": 1642826548.2530458,
-    "train_yamls\\train_2layer_FM.yaml": 1642408813.9488862
+    "train_yamls\\train_2layer_FM.yaml": 1642411635.5534878,
+    "components\\Generator_reduce.py": 1642690262.572149,
+    "insightface_func\\face_detect_crop_multi.py": 1638370471.789609,
+    "insightface_func\\face_detect_crop_single.py": 1638370471.7967434,
+    "insightface_func\\__init__.py": 1624197300.011183,
+    "insightface_func\\utils\\face_align_ffhqandnewarc.py": 1638370471.850638,
+    "losses\\PatchNCE.py": 1642989384.9713614,
+    "parsing_model\\model.py": 1626745709.554252,
+    "parsing_model\\resnet.py": 1626745709.554252,
+    "test_scripts\\tester_common copy.py": 1625369535.199175,
+    "test_scripts\\tester_video.py": 1642734397.3307388,
+    "train_scripts\\trainer_cycleloss.py": 1642580463.495596,
+    "train_scripts\\trainer_GramFM.py": 1643010471.1077821,
+    "utilities\\ImagenetNorm.py": 1642732910.5280058,
+    "utilities\\reverse2original.py": 1642733688.7976837,
+    "train_yamls\\train_cycleloss.yaml": 1642577741.345273,
+    "train_yamls\\train_GramFM.yaml": 1643011210.82505
 }

components/Generator_reduce.py

@@ -5,7 +5,7 @@
 # Created Date: Sunday January 16th 2022
 # Author: Chen Xuanhong
 # Email: chenxuanhongzju@outlook.com
-# Last Modified:  Thursday, 20th January 2022 10:51:02 pm
+# Last Modified:  Monday, 24th January 2022 6:47:22 pm
 # Modified By: Chen Xuanhong
 # Copyright (c) 2022 Shanghai Jiao Tong University
 #############################################################
@@ -15,17 +15,16 @@ from torch import nn
 from torch.nn import init
 from torch.nn import functional as F
 
-class InstanceNorm(nn.Module):
+class Demodule(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__()
+        super(Demodule, 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
@@ -46,9 +45,22 @@ class ApplyStyle(nn.Module):
         x = x * (style[:, 0] * 1 + 1.) + style[:, 1] * 1
         return x
 
-class ResnetBlock_Adain(nn.Module):
+class Modulation(nn.Module):
+    def __init__(self, latent_size, channels):
+        super(Modulation, self).__init__()
+        self.linear = nn.Linear(latent_size, channels)
+
+    def forward(self, x, latent):
+        style = self.linear(latent)  # style => [batch_size, n_channels*2]
+        shape = [-1, x.size(1), 1, 1]
+        style = style.view(shape)    # [batch_size, 2, n_channels, ...]
+        #x = x * (style[:, 0] + 1.) + style[:, 1]
+        x = x * style
+        return x
+
+class ResnetBlock_Modulation(nn.Module):
     def __init__(self, dim, latent_size, padding_type, activation=nn.ReLU(True)):
-        super(ResnetBlock_Adain, self).__init__()
+        super(ResnetBlock_Modulation, self).__init__()
 
         p = 0
         conv1 = []
@@ -60,9 +72,9 @@ class ResnetBlock_Adain(nn.Module):
             p = 1
         else:
             raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-        conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), InstanceNorm()]
+        conv1 += [nn.Conv2d(dim, dim, kernel_size=3, padding = p), Demodule()]
         self.conv1 = nn.Sequential(*conv1)
-        self.style1 = ApplyStyle(latent_size, dim)
+        self.style1 = Modulation(latent_size, dim)
         self.act1 = activation
 
         p = 0
@@ -75,52 +87,9 @@ class ResnetBlock_Adain(nn.Module):
             p = 1
         else:
             raise NotImplementedError('padding [%s] is not implemented' % padding_type)
-        conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), InstanceNorm()]
+        conv2 += [nn.Conv2d(dim, dim, kernel_size=3, padding=p), Demodule()]
         self.conv2 = nn.Sequential(*conv2)
-        self.style2 = ApplyStyle(latent_size, dim)
+        self.style2 = Modulation(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 newres(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=1), 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):
@@ -148,7 +117,7 @@ class Generator(nn.Module):
         
         activation = nn.ReLU(True)
 
-        self.first_layer = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1),
+        self.stem = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, padding=1),
                                 nn.BatchNorm2d(64), activation)
         ### downsample
         self.down1 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
@@ -167,7 +136,7 @@ class Generator(nn.Module):
         BN = []
         for i in range(res_num):
             BN += [
-                ResnetBlock_Adain(512, latent_size=chn, padding_type=padding_type, activation=activation)]
+                ResnetBlock_Modulation(512, latent_size=chn, padding_type=padding_type, activation=activation)]
         self.BottleNeck = nn.Sequential(*BN)
 
         self.up4 = nn.Sequential(
@@ -210,7 +179,7 @@ class Generator(nn.Module):
 
     def forward(self, input, id):
         x = input  # 3*224*224
-        skip1 = self.first_layer(x)
+        skip1 = self.stem(x)
         skip2 = self.down1(skip1)
         skip3 = self.down2(skip2)
         skip4 = self.down3(skip3)

losses/PatchNCE.py

@@ -5,9 +5,225 @@
 # Created Date: Friday January 21st 2022
 # Author: Chen Xuanhong
 # Email: chenxuanhongzju@outlook.com
-# Last Modified:  Friday, 21st January 2022 5:04:43 pm
+# Last Modified:  Monday, 24th January 2022 9:56:24 am
 # Modified By: Chen Xuanhong
 # Copyright (c) 2022 Shanghai Jiao Tong University
 #############################################################
 
+import torch
+import torch.nn as nn
+from   torch.nn import init
+import torch.nn.functional as F
+from packaging import version
 
+import numpy as np
+
+
+
+def init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], debug=False, initialize_weights=True):
+    """Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights
+    Parameters:
+        net (network)      -- the network to be initialized
+        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        gain (float)       -- scaling factor for normal, xavier and orthogonal.
+        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
+
+    Return an initialized network.
+    """
+    if len(gpu_ids) > 0:
+        assert(torch.cuda.is_available())
+        net.to(gpu_ids[0])
+        # if not amp:
+        # net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs for non-AMP training
+    if initialize_weights:
+        init_weights(net, init_type, init_gain=init_gain, debug=debug)
+    return net
+
+def init_weights(net, init_type='normal', init_gain=0.02, debug=False):
+    """Initialize network weights.
+
+    Parameters:
+        net (network)   -- network to be initialized
+        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal
+        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.
+
+    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might
+    work better for some applications. Feel free to try yourself.
+    """
+    def init_func(m):  # define the initialization function
+        classname = m.__class__.__name__
+        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):
+            if debug:
+                print(classname)
+            if init_type == 'normal':
+                init.normal_(m.weight.data, 0.0, init_gain)
+            elif init_type == 'xavier':
+                init.xavier_normal_(m.weight.data, gain=init_gain)
+            elif init_type == 'kaiming':
+                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')
+            elif init_type == 'orthogonal':
+                init.orthogonal_(m.weight.data, gain=init_gain)
+            else:
+                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)
+            if hasattr(m, 'bias') and m.bias is not None:
+                init.constant_(m.bias.data, 0.0)
+        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.
+            init.normal_(m.weight.data, 1.0, init_gain)
+            init.constant_(m.bias.data, 0.0)
+
+    net.apply(init_func)  # apply the initialization function <init_func>
+
+class Normalize(nn.Module):
+
+    def __init__(self, power=2):
+        super(Normalize, self).__init__()
+        self.power = power
+
+    def forward(self, x):
+        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)
+        out = x.div(norm + 1e-7)
+        return out
+
+class PatchSampleF(nn.Module):
+    def __init__(self, use_mlp=False, init_type='normal', init_gain=0.02, nc=256, gpu_ids=[]):
+        # potential issues: currently, we use the same patch_ids for multiple images in the batch
+        super(PatchSampleF, self).__init__()
+        self.l2norm = Normalize(2)
+        self.use_mlp = use_mlp
+        self.nc = nc  # hard-coded
+        self.mlp_init = False
+        self.init_type = init_type
+        self.init_gain = init_gain
+        self.gpu_ids = gpu_ids
+
+    def create_mlp(self, feats):
+        for mlp_id, feat in enumerate(feats):
+            input_nc = feat.shape[1]
+            mlp = nn.Sequential(*[nn.Linear(input_nc, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)])
+            if len(self.gpu_ids) > 0:
+                mlp.cuda()
+            setattr(self, 'mlp_%d' % mlp_id, mlp)
+        init_net(self, self.init_type, self.init_gain, self.gpu_ids)
+        self.mlp_init = True
+
+    def forward(self, feats, num_patches=64, patch_ids=None):
+        return_ids = []
+        return_feats = []
+        if self.use_mlp and not self.mlp_init:
+            self.create_mlp(feats)
+        for feat_id, feat in enumerate(feats):
+            B, H, W = feat.shape[0], feat.shape[2], feat.shape[3]
+            feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2)
+            if num_patches > 0:
+                if patch_ids is not None:
+                    patch_id = patch_ids[feat_id]
+                else:
+                    # torch.randperm produces cudaErrorIllegalAddress for newer versions of PyTorch. https://github.com/taesungp/contrastive-unpaired-translation/issues/83
+                    #patch_id = torch.randperm(feat_reshape.shape[1], device=feats[0].device)
+                    patch_id = np.random.permutation(feat_reshape.shape[1])
+                    patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))]  # .to(patch_ids.device)
+                x_sample = feat_reshape[:, patch_id, :].flatten(0, 1)  # reshape(-1, x.shape[1])
+            else:
+                x_sample = feat_reshape
+                patch_id = []
+            if self.use_mlp:
+                mlp = getattr(self, 'mlp_%d' % feat_id)
+                x_sample = mlp(x_sample)
+            return_ids.append(patch_id)
+            x_sample = self.l2norm(x_sample)
+
+            if num_patches == 0:
+                x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W])
+            return_feats.append(x_sample)
+        return return_feats, return_ids
+
+
+
+# def calculate_NCE_loss(src, tgt):
+#     n_layers = len(self.nce_layers)
+#     feat_q = self.netG(tgt, self.nce_layers, encode_only=True)
+
+#     if self.opt.flip_equivariance and self.flipped_for_equivariance:
+#         feat_q = [torch.flip(fq, [3]) for fq in feat_q]
+
+#     feat_k = self.netG(src, self.nce_layers, encode_only=True)
+#     feat_k_pool, sample_ids = self.netF(feat_k, self.opt.num_patches, None)
+#     feat_q_pool, _ = self.netF(feat_q, self.opt.num_patches, sample_ids)
+
+#     total_nce_loss = 0.0
+#     for f_q, f_k, crit, nce_layer in zip(feat_q_pool, feat_k_pool, self.criterionNCE, self.nce_layers):
+#         loss = crit(f_q, f_k) * self.opt.lambda_NCE
+#         total_nce_loss += loss.mean()
+
+    # return total_nce_loss / n_layers
+
+class PatchNCELoss(nn.Module):
+    def __init__(self,batch_size, nce_T = 0.07):
+        super().__init__()
+        self.nce_T = nce_T
+        self.batch_size = batch_size
+        self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none')
+        self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool
+
+    def forward(self, feat_q, feat_k):
+        num_patches = feat_q.shape[0]
+        dim = feat_q.shape[1]
+        feat_k = feat_k.detach()
+
+        # pos logit
+        l_pos = torch.bmm(
+            feat_q.view(num_patches, 1, -1), feat_k.view(num_patches, -1, 1))
+        l_pos = l_pos.view(num_patches, 1)
+
+        # neg logit
+
+        # Should the negatives from the other samples of a minibatch be utilized?
+        # In CUT and FastCUT, we found that it's best to only include negatives
+        # from the same image. Therefore, we set
+        # --nce_includes_all_negatives_from_minibatch as False
+        # However, for single-image translation, the minibatch consists of
+        # crops from the "same" high-resolution image.
+        # Therefore, we will include the negatives from the entire minibatch.
+        # if self.opt.nce_includes_all_negatives_from_minibatch:
+        #     # reshape features as if they are all negatives of minibatch of size 1.
+        #     batch_dim_for_bmm = 1
+        # else:
+        batch_dim_for_bmm = self.batch_size
+
+        # reshape features to batch size
+        feat_q = feat_q.view(batch_dim_for_bmm, -1, dim)
+        feat_k = feat_k.view(batch_dim_for_bmm, -1, dim)
+        npatches = feat_q.size(1)
+        l_neg_curbatch = torch.bmm(feat_q, feat_k.transpose(2, 1))
+
+        # diagonal entries are similarity between same features, and hence meaningless.
+        # just fill the diagonal with very small number, which is exp(-10) and almost zero
+        diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :]
+        l_neg_curbatch.masked_fill_(diagonal, -10.0)
+        l_neg = l_neg_curbatch.view(-1, npatches)
+
+        out = torch.cat((l_pos, l_neg), dim=1) / self.nce_T
+
+        loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long,
+                                                        device=feat_q.device))
+
+        return loss
+
+if __name__ == "__main__":
+    batch = 16
+    nc = 256
+    num_patches = 64
+    netF = PatchSampleF(use_mlp=True, init_type='normal', init_gain=0.02, gpu_ids=["cuda:0"], nc=nc)
+
+    feat_q = [torch.ones((batch,nc,32,32)).cuda()]
+
+    # if self.opt.flip_equivariance and self.flipped_for_equivariance:
+    #     feat_q = [torch.flip(fq, [3]) for fq in feat_q]
+    crit = PatchNCELoss(batch)
+    feat_k = [torch.ones((batch,nc,32,32)).cuda()]
+    feat_k_pool, sample_ids = netF(feat_k, num_patches, None)
+    print(feat_k_pool[0].shape)
+    feat_q_pool, _ = netF(feat_q, num_patches, sample_ids)
+    print(feat_q_pool[0].shape)
+    loss = crit(feat_q_pool[0], feat_k_pool[0])
+    print(loss)

start_train.sh

@@ -0,0 +1,3 @@
+
+
+nohup python train.py > GramFM.log 2>&1 &

train.py

@@ -5,7 +5,7 @@
 # Created Date: Tuesday April 28th 2020
 # Author: Chen Xuanhong
 # Email: chenxuanhongzju@outlook.com
-# Last Modified:  Monday, 17th January 2022 4:42:53 pm
+# Last Modified:  Monday, 24th January 2022 3:32:47 pm
 # Modified By: Chen Xuanhong
 # Copyright (c) 2020 Shanghai Jiao Tong University
 #############################################################
@@ -31,9 +31,9 @@ def getParameters():
     
     parser = argparse.ArgumentParser()
     # general settings
-    parser.add_argument('-v', '--version', type=str, default='2layerFM',
+    parser.add_argument('-v', '--version', type=str, default='GramFM',
                                             help="version name for train, test, finetune")
-    parser.add_argument('-t', '--tag', type=str, default='Feature_match',
+    parser.add_argument('-t', '--tag', type=str, default='Gram_Feature_match',
                                             help="tag for current experiment")
 
     parser.add_argument('-p', '--phase', type=str, default="train",
@@ -46,9 +46,9 @@ def getParameters():
 
     # training
     parser.add_argument('--experiment_description', type=str,
-                                default="减小重建和feature match的权重，使用2和3的feature作为feature")
+                                default="使用3作为feature, 尝试使用gram矩阵来计算feature matching")
 
-    parser.add_argument('--train_yaml', type=str, default="train_2layer_FM.yaml")
+    parser.add_argument('--train_yaml', type=str, default="train_GramFM.yaml")
 
     # system logger
     parser.add_argument('--logger', type=str,

train_scripts/trainer_FM.py

@@ -5,7 +5,7 @@
 # Created Date: Sunday January 9th 2022
 # Author: Chen Xuanhong
 # Email: chenxuanhongzju@outlook.com
-# Last Modified:  Saturday, 22nd January 2022 12:45:09 pm
+# Last Modified:  Monday, 24th January 2022 6:56:17 pm
 # Modified By: Chen Xuanhong
 # Copyright (c) 2022 Shanghai Jiao Tong University
 #############################################################
@@ -29,6 +29,9 @@ class Trainer(TrainerBase):
                 reporter):
         super(Trainer, self).__init__(config, reporter)
 
+        import inspect
+        print("Current training script -----------> %s"%inspect.getfile(inspect.currentframe()))
+        
         self.img_std    = torch.Tensor([0.229, 0.224, 0.225]).view(3,1,1)
         self.img_mean   = torch.Tensor([0.485, 0.456, 0.406]).view(3,1,1)
     
@@ -276,25 +279,27 @@ class Trainer(TrainerBase):
                 elapsed = str(datetime.timedelta(seconds=elapsed))
                 
                 epochinformation="[{}], Elapsed [{}], Step [{}/{}], \
-                        G_loss: {:.4f}, Rec_loss: {:.4f}, Fm_loss: {:.4f}, \
+                        G_ID: {:.4f}, G_loss: {:.4f}, Rec_loss: {:.4f}, Fm_loss: {:.4f}, \
-                            D_loss: {:.4f}, D_fake: {:.4f}, D_real: {:.4f}". \
+                        D_loss: {:.4f}, D_fake: {:.4f}, D_real: {:.4f}". \
-                            format(self.config["version"], elapsed, step, total_step, \
+                        format(self.config["version"], elapsed, step, total_step, \
-                                loss_G.item(), loss_G_Rec.item(), feat_match_loss.item(), \
+                        loss_G_ID.item(), loss_G.item(), loss_G_Rec.item(), feat_match_loss.item(), \
-                                    loss_D.item(), loss_Dgen.item(), loss_Dreal.item())
+                        loss_D.item(), loss_Dgen.item(), loss_Dreal.item())
                 print(epochinformation)
                 self.reporter.writeInfo(epochinformation)
 
                 if self.config["logger"] == "tensorboard":
                     self.logger.add_scalar('G/G_loss', loss_G.item(), step)
-                    self.logger.add_scalar('G/Rec_loss', loss_G_Rec.item(), step)
+                    self.logger.add_scalar('G/G_Rec', loss_G_Rec.item(), step)
-                    self.logger.add_scalar('G/Fm_loss', feat_match_loss.item(), step)
+                    self.logger.add_scalar('G/G_feat_match', feat_match_loss.item(), step)
+                    self.logger.add_scalar('G/G_ID', loss_G_ID.item(), step)
                     self.logger.add_scalar('D/D_loss', loss_D.item(), step)
                     self.logger.add_scalar('D/D_fake', loss_Dgen.item(), step)
                     self.logger.add_scalar('D/D_real', loss_Dreal.item(), step)
                 elif self.config["logger"] == "wandb":
                     self.logger.log({"G_loss": loss_G.item()}, step = step)
-                    self.logger.log({"Rec_loss": loss_G_Rec.item()}, step = step)
+                    self.logger.log({"G_Rec": loss_G_Rec.item()}, step = step)
-                    self.logger.log({"Fm_loss": feat_match_loss.item()}, step = step)
+                    self.logger.log({"G_feat_match": feat_match_loss.item()}, step = step)
+                    self.logger.log({"G_ID": loss_G_ID.item()}, step = step)
                     self.logger.log({"D_loss": loss_D.item()}, step = step)
                     self.logger.log({"D_fake": loss_Dgen.item()}, step = step)
                     self.logger.log({"D_real": loss_Dreal.item()}, step = step)

train_scripts/trainer_GramFM.py

@@ -0,0 +1,351 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: trainer_naiv512.py
+# Created Date: Sunday January 9th 2022
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Monday, 24th January 2022 6:23:16 pm
+# Modified By: Chen Xuanhong
+# Copyright (c) 2022 Shanghai Jiao Tong University
+#############################################################
+
+import  os
+import  time
+import  random
+
+import  numpy as np
+
+import  torch
+import  torch.nn.functional as F
+from    utilities.plot import plot_batch
+
+from    train_scripts.trainer_base import TrainerBase
+
+from    utilities.utilities import Gram
+
+class Trainer(TrainerBase):
+
+    def __init__(self, 
+                config, 
+                reporter):
+        super(Trainer, self).__init__(config, reporter)
+        import inspect
+        print("Current training script -----------> %s"%inspect.getfile(inspect.currentframe()))
+        
+        self.img_std    = torch.Tensor([0.229, 0.224, 0.225]).view(3,1,1)
+        self.img_mean   = torch.Tensor([0.485, 0.456, 0.406]).view(3,1,1)
+    
+    # TODO modify this function to build your models
+    def init_framework(self):
+        '''
+            This function is designed to define the framework,
+            and print the framework information into the log file
+        '''
+        #===============build models================#
+        print("build models...")
+        # TODO [import models here]
+
+        model_config    = self.config["model_configs"]
+        
+        if self.config["phase"] == "train":
+            gscript_name     = "components." + model_config["g_model"]["script"]
+            dscript_name     = "components." + model_config["d_model"]["script"]
+            
+        elif self.config["phase"] == "finetune":
+            gscript_name     = self.config["com_base"] + model_config["g_model"]["script"]
+            dscript_name     = self.config["com_base"] + model_config["d_model"]["script"]
+        
+        class_name      = model_config["g_model"]["class_name"]
+        package         = __import__(gscript_name, fromlist=True)
+        gen_class       = getattr(package, class_name)
+        self.gen        = gen_class(**model_config["g_model"]["module_params"])
+        
+        # print and recorde model structure
+        self.reporter.writeInfo("Generator structure:")
+        self.reporter.writeModel(self.gen.__str__())
+
+        class_name      = model_config["d_model"]["class_name"]
+        package         = __import__(dscript_name, fromlist=True)
+        dis_class       = getattr(package, class_name)
+        self.dis        = dis_class(**model_config["d_model"]["module_params"])
+        self.dis.feature_network.requires_grad_(False)
+        
+        # print and recorde model structure
+        self.reporter.writeInfo("Discriminator structure:")
+        self.reporter.writeModel(self.dis.__str__())
+        arcface1        = torch.load(self.arcface_ckpt, map_location=torch.device("cpu"))
+        self.arcface    = arcface1['model'].module
+        
+        # train in GPU
+        if self.config["cuda"] >=0:
+            self.gen    = self.gen.cuda()
+            self.dis    = self.dis.cuda()
+            self.arcface= self.arcface.cuda()
+        
+        self.arcface.eval()
+        self.arcface.requires_grad_(False)
+
+        # if in finetune phase, load the pretrained checkpoint
+        if self.config["phase"] == "finetune":
+            model_path = os.path.join(self.config["project_checkpoints"],
+                                        "step%d_%s.pth"%(self.config["checkpoint_step"],
+                                        self.config["checkpoint_names"]["generator_name"]))
+            self.gen.load_state_dict(torch.load(model_path))
+
+            model_path = os.path.join(self.config["project_checkpoints"],
+                                        "step%d_%s.pth"%(self.config["checkpoint_step"],
+                                        self.config["checkpoint_names"]["discriminator_name"]))
+            self.dis.load_state_dict(torch.load(model_path))
+            
+            print('loaded trained backbone model step {}...!'.format(self.config["project_checkpoints"]))
+    
+    # TODO modify this function to configurate the optimizer of your pipeline
+    def __setup_optimizers__(self):
+        g_train_opt     = self.config['g_optim_config']
+        d_train_opt     = self.config['d_optim_config']
+
+        g_optim_params  = []
+        d_optim_params  = []
+        for k, v in self.gen.named_parameters():
+            if v.requires_grad:
+                g_optim_params.append(v)
+            else:
+                self.reporter.writeInfo(f'Params {k} will not be optimized.')
+                print(f'Params {k} will not be optimized.')
+        
+        for k, v in self.dis.named_parameters():
+            if v.requires_grad:
+                d_optim_params.append(v)
+            else:
+                self.reporter.writeInfo(f'Params {k} will not be optimized.')
+                print(f'Params {k} will not be optimized.')
+
+        optim_type = self.config['optim_type']
+        
+        if optim_type == 'Adam':
+            self.g_optimizer = torch.optim.Adam(g_optim_params,**g_train_opt)
+            self.d_optimizer = torch.optim.Adam(d_optim_params,**d_train_opt)
+        else:
+            raise NotImplementedError(
+                f'optimizer {optim_type} is not supperted yet.')
+        # self.optimizers.append(self.optimizer_g)
+        if self.config["phase"] == "finetune":
+            opt_path = os.path.join(self.config["project_checkpoints"],
+                                        "step%d_optim_%s.pth"%(self.config["checkpoint_step"],
+                                        self.config["optimizer_names"]["generator_name"]))
+            self.g_optimizer.load_state_dict(torch.load(opt_path))
+
+            opt_path = os.path.join(self.config["project_checkpoints"],
+                                        "step%d_optim_%s.pth"%(self.config["checkpoint_step"],
+                                        self.config["optimizer_names"]["discriminator_name"]))
+            self.d_optimizer.load_state_dict(torch.load(opt_path))
+            
+            print('loaded trained optimizer step {}...!'.format(self.config["project_checkpoints"]))
+    
+
+    # TODO modify this function to evaluate your model
+    # Evaluate the checkpoint
+    def __evaluation__(self,
+            step = 0,
+            **kwargs
+            ):
+        src_image1  = kwargs["src1"]
+        src_image2  = kwargs["src2"]
+        batch_size  = self.batch_size
+        self.gen.eval()
+        with torch.no_grad():
+            imgs        = []
+            zero_img    = (torch.zeros_like(src_image1[0,...]))
+            imgs.append(zero_img.cpu().numpy())
+            save_img    = ((src_image1.cpu())* self.img_std + self.img_mean).numpy()
+            for r in range(batch_size):
+                imgs.append(save_img[r,...])
+            arcface_112     = F.interpolate(src_image2,size=(112,112), mode='bicubic')
+            id_vector_src1  = self.arcface(arcface_112)
+            id_vector_src1  = F.normalize(id_vector_src1, p=2, dim=1)
+
+            for i in range(batch_size):
+                
+                imgs.append(save_img[i,...])
+                image_infer = src_image1[i, ...].repeat(batch_size, 1, 1, 1)
+                img_fake    = self.gen(image_infer, id_vector_src1).cpu()
+                
+                img_fake    = img_fake * self.img_std
+                img_fake    = img_fake + self.img_mean
+                img_fake    = img_fake.numpy()
+                for j in range(batch_size):
+                    imgs.append(img_fake[j,...])
+            print("Save test data")
+            imgs = np.stack(imgs, axis = 0).transpose(0,2,3,1)
+            plot_batch(imgs, os.path.join(self.sample_dir, 'step_'+str(step+1)+'.jpg'))
+
+    
+                
+
+    def train(self):
+
+        ckpt_dir    = self.config["project_checkpoints"]
+        log_freq    = self.config["log_step"]
+        model_freq  = self.config["model_save_step"]
+        sample_freq = self.config["sample_step"]
+        total_step  = self.config["total_step"]
+        random_seed = self.config["dataset_params"]["random_seed"]
+
+        self.batch_size  = self.config["batch_size"]
+        self.sample_dir  = self.config["project_samples"]
+        self.arcface_ckpt= self.config["arcface_ckpt"]
+        
+        
+        # prep_weights= self.config["layersWeight"]
+        id_w        = self.config["id_weight"]
+        rec_w       = self.config["reconstruct_weight"]
+        feat_w      = self.config["feature_match_weight"]
+        
+        
+        
+        super().train()
+        
+        #===============build losses===================#
+        # TODO replace below lines to build your losses
+        MSE_loss    = torch.nn.MSELoss()
+        l1_loss     = torch.nn.L1Loss()
+        cos_loss    = torch.nn.CosineSimilarity()
+
+            
+        start_time  = time.time()
+
+        # Caculate the epoch number
+        print("Total step = %d"%total_step)
+        random.seed(random_seed)
+        randindex = [i for i in range(self.batch_size)]
+        random.shuffle(randindex)
+        import datetime
+        for step in range(self.start, total_step):
+            self.gen.train()
+            self.dis.train()
+            for interval in range(2):
+                random.shuffle(randindex)
+                src_image1, src_image2  = self.train_loader.next()
+                
+                if step%2 == 0:
+                    img_id = src_image2
+                else:
+                    img_id = src_image2[randindex]
+
+                img_id_112      = F.interpolate(img_id,size=(112,112), mode='bicubic')
+                latent_id       = self.arcface(img_id_112)
+                latent_id       = F.normalize(latent_id, p=2, dim=1)
+                if interval:
+                    
+                    img_fake        = self.gen(src_image1, latent_id)
+                    gen_logits,_    = self.dis(img_fake.detach(), None)
+                    loss_Dgen       = (F.relu(torch.ones_like(gen_logits) + gen_logits)).mean()
+
+                    real_logits,_   = self.dis(src_image2,None)
+                    loss_Dreal      = (F.relu(torch.ones_like(real_logits) - real_logits)).mean()
+
+                    loss_D          = loss_Dgen + loss_Dreal
+                    self.d_optimizer.zero_grad()
+                    loss_D.backward()
+                    self.d_optimizer.step()
+                else:
+                    
+                    # model.netD.requires_grad_(True)
+                    img_fake        = self.gen(src_image1, latent_id)
+                    # G loss
+                    gen_logits,feat = self.dis(img_fake, None)
+                    
+                    loss_Gmain      = (-gen_logits).mean()
+                    img_fake_down   = F.interpolate(img_fake, size=(112,112), mode='bicubic')
+                    latent_fake     = self.arcface(img_fake_down)
+                    latent_fake     = F.normalize(latent_fake, p=2, dim=1)
+                    loss_G_ID       = (1 - cos_loss(latent_fake, latent_id)).mean()
+                    real_feat       = self.dis.get_feature(src_image1)
+                    feat_match_loss = l1_loss(Gram(feat["3"]), Gram(real_feat["3"])) + \
+                                        l1_loss(Gram(feat["2"]), Gram(real_feat["2"]))
+                    loss_G          = loss_Gmain + loss_G_ID * id_w + \
+                                                feat_match_loss * feat_w
+                    if step%2 == 0:
+                        #G_Rec
+                        loss_G_Rec  = l1_loss(img_fake, src_image1)
+                        loss_G      += loss_G_Rec * rec_w
+
+                    self.g_optimizer.zero_grad()
+                    loss_G.backward()
+                    self.g_optimizer.step()
+                
+            # Print out log info
+            if (step + 1) % log_freq == 0:
+                elapsed = time.time() - start_time
+                elapsed = str(datetime.timedelta(seconds=elapsed))
+                
+                epochinformation="[{}], Elapsed [{}], Step [{}/{}], \
+                        G_ID: {:.4f}, G_loss: {:.4f}, Rec_loss: {:.4f}, Fm_loss: {:.4f}, \
+                        D_loss: {:.4f}, D_fake: {:.4f}, D_real: {:.4f}". \
+                        format(self.config["version"], elapsed, step, total_step, \
+                        loss_G_ID.item(), loss_G.item(), loss_G_Rec.item(), feat_match_loss.item(), \
+                        loss_D.item(), loss_Dgen.item(), loss_Dreal.item())
+                print(epochinformation)
+                self.reporter.writeInfo(epochinformation)
+
+                if self.config["logger"] == "tensorboard":
+                    self.logger.add_scalar('G/G_loss', loss_G.item(), step)
+                    self.logger.add_scalar('G/G_Rec', loss_G_Rec.item(), step)
+                    self.logger.add_scalar('G/G_feat_match', feat_match_loss.item(), step)
+                    self.logger.add_scalar('G/G_ID', loss_G_ID.item(), step)
+                    self.logger.add_scalar('D/D_loss', loss_D.item(), step)
+                    self.logger.add_scalar('D/D_fake', loss_Dgen.item(), step)
+                    self.logger.add_scalar('D/D_real', loss_Dreal.item(), step)
+                elif self.config["logger"] == "wandb":
+                    self.logger.log({"G_loss": loss_G.item()}, step = step)
+                    self.logger.log({"G_Rec": loss_G_Rec.item()}, step = step)
+                    self.logger.log({"G_feat_match": feat_match_loss.item()}, step = step)
+                    self.logger.log({"G_ID": loss_G_ID.item()}, step = step)
+                    self.logger.log({"D_loss": loss_D.item()}, step = step)
+                    self.logger.log({"D_fake": loss_Dgen.item()}, step = step)
+                    self.logger.log({"D_real": loss_Dreal.item()}, step = step)
+            
+            if (step + 1) % sample_freq == 0:
+                self.__evaluation__(
+                    step = step,
+                    **{
+                    "src1": src_image1,
+                    "src2": src_image2
+                })
+                    
+                        
+                
+            #===============adjust learning rate============#
+            # if (epoch + 1) in self.config["lr_decay_step"] and self.config["lr_decay_enable"]:
+            #     print("Learning rate decay")
+            #     for p in self.optimizer.param_groups:
+            #         p['lr'] *= self.config["lr_decay"]
+            #         print("Current learning rate is %f"%p['lr'])
+
+            #===============save checkpoints================#
+            if (step+1) % model_freq==0:
+                
+                torch.save(self.gen.state_dict(),
+                        os.path.join(ckpt_dir, 'step{}_{}.pth'.format(step + 1, 
+                                    self.config["checkpoint_names"]["generator_name"])))
+                torch.save(self.dis.state_dict(),
+                        os.path.join(ckpt_dir, 'step{}_{}.pth'.format(step + 1, 
+                                    self.config["checkpoint_names"]["discriminator_name"])))
+                
+                torch.save(self.g_optimizer.state_dict(),
+                        os.path.join(ckpt_dir, 'step{}_optim_{}'.format(step + 1, 
+                                    self.config["checkpoint_names"]["generator_name"])))
+                
+                torch.save(self.d_optimizer.state_dict(),
+                        os.path.join(ckpt_dir, 'step{}_optim_{}'.format(step + 1, 
+                                    self.config["checkpoint_names"]["discriminator_name"])))
+                print("Save step %d model checkpoint!"%(step+1))
+                torch.cuda.empty_cache()
+
+                self.__evaluation__(
+                    step = step,
+                    **{
+                    "src1": src_image1,
+                    "src2": src_image2
+                })

train_yamls/train_512FM.yaml

@@ -50,7 +50,7 @@ d_optim_config:
   eps: !!float 1e-8
 
 id_weight: 20.0
-reconstruct_weight: 1.0
+reconstruct_weight: 10.0
 feature_match_weight: 10.0
 
 # Log 

train_yamls/train_512FM_Modulation.yaml

@@ -0,0 +1,63 @@
+# Related scripts
+train_script_name: FM
+
+# models' scripts
+model_configs:
+  g_model:
+    script: Generator_reduce
+    class_name: Generator
+    module_params:
+      g_conv_dim: 512
+      g_kernel_size: 3
+      res_num: 9
+
+  d_model:
+    script: projected_discriminator
+    class_name: ProjectedDiscriminator
+    module_params:
+      diffaug: False
+      interp224: False
+      backbone_kwargs: {}
+
+arcface_ckpt: arcface_ckpt/arcface_checkpoint.tar
+
+# Training information
+batch_size: 12
+
+# Dataset
+dataloader: VGGFace2HQ
+dataset_name: vggface2_hq
+dataset_params:
+  random_seed: 1234
+  dataloader_workers: 8
+
+eval_dataloader: DIV2K_hdf5
+eval_dataset_name: DF2K_H5_Eval
+eval_batch_size: 2
+
+# Dataset
+
+# Optimizer
+optim_type: Adam
+g_optim_config:
+  lr:  0.0004
+  betas: [ 0, 0.99]
+  eps: !!float 1e-8
+
+d_optim_config:
+  lr:  0.0004
+  betas: [ 0, 0.99]
+  eps: !!float 1e-8
+
+id_weight: 20.0
+reconstruct_weight: 10.0
+feature_match_weight: 10.0
+
+# Log 
+log_step: 300
+model_save_step: 10000
+sample_step: 1000
+total_step: 1000000
+checkpoint_names:
+  generator_name: Generator
+  discriminator_name: Discriminator

train_yamls/train_GramFM.yaml

@@ -0,0 +1,63 @@
+# Related scripts
+train_script_name: GramFM
+
+# models' scripts
+model_configs:
+  g_model:
+    script: Generator
+    class_name: Generator
+    module_params:
+      g_conv_dim: 512
+      g_kernel_size: 3
+      res_num: 9
+
+  d_model:
+    script: projected_discriminator
+    class_name: ProjectedDiscriminator
+    module_params:
+      diffaug: False
+      interp224: False
+      backbone_kwargs: {}
+
+arcface_ckpt: arcface_ckpt/arcface_checkpoint.tar
+
+# Training information
+batch_size: 12
+
+# Dataset
+dataloader: VGGFace2HQ
+dataset_name: vggface2_hq
+dataset_params:
+  random_seed: 1234
+  dataloader_workers: 8
+
+eval_dataloader: DIV2K_hdf5
+eval_dataset_name: DF2K_H5_Eval
+eval_batch_size: 2
+
+# Dataset
+
+# Optimizer
+optim_type: Adam
+g_optim_config:
+  lr:  0.0004
+  betas: [ 0, 0.99]
+  eps: !!float 1e-8
+
+d_optim_config:
+  lr:  0.0004
+  betas: [ 0, 0.99]
+  eps: !!float 1e-8
+
+id_weight: 20.0
+reconstruct_weight: 10.0
+feature_match_weight: 100.0
+
+# Log 
+log_step: 300
+model_save_step: 10000
+total_step: 1000000
+sample_step: 1000
+checkpoint_names:
+  generator_name: Generator
+  discriminator_name: Discriminator