add equlaized-lr feature, add smoothed genetor

Author: min-nashory

config.py

@@ -6,7 +6,7 @@
 parser = argparse.ArgumentParser('PGGAN')
 
 ## general settings.
-parser.add_argument('--train_data_root', type=str, default='/home1/irteam/nashory/data/CelebA/Img')
+parser.add_argument('--train_data_root', type=str, default='/home1/work/nashory/data/CelebA/Img')
 parser.add_argument('--random_seed', type=int, default=int(time.time()))
 parser.add_argument('--n_gpu', type=int, default=1)         # for Multi-GPU training.
 
@@ -16,7 +16,8 @@
 
 
 ## training parameters.
-parser.add_argument('--lr', type=float, default=0.0002)
+parser.add_argument('--lr', type=float, default=0.001)
+parser.add_argument('--smoothing', type=float, default=0.997)
 parser.add_argument('--nc', type=int, default=3)
 parser.add_argument('--nz', type=int, default=512)
 parser.add_argument('--ngf', type=int, default=512)
@@ -32,7 +33,7 @@
 parser.add_argument('--flag_bn', type=bool, default=False)
 parser.add_argument('--flag_pixelwise', type=bool, default=True)
 parser.add_argument('--flag_leaky', type=bool, default=True)
-parser.add_argument('--flag_tanh', type=bool, default=False)
+parser.add_argument('--flag_tanh', type=bool, default=True)
 parser.add_argument('--flag_sigmoid', type=bool, default=True)
 
 

custom_layers.py

@@ -66,3 +66,38 @@ def forward(self, x):
         return torch.addcdiv(t, 1, x, norm)
 
 
+# for equaliaeed-learning rate.
+class equalized_conv2d(nn.Module):
+    def __init__(self, c_in, c_out, k_size, stride, pad, initializer='kaiming'):
+        super(equalized_conv2d, self).__init__()
+        self.conv = nn.Conv2d(c_in, c_out, k_size, stride, pad)
+        if initializer == 'kaiming':    torch.nn.init.kaiming_normal(self.conv.weight)
+        elif initializer == 'xavier':   torch.nn.init.xavier_normal(self.conv.weight)
+        self.inv_c = np.sqrt(2.0/(c_in*k_size**2))
+
+    def forward(self, x):
+        return self.conv(x.mul(self.inv_c))
+        
+ 
+class equalized_deconv2d(nn.Module):
+    def __init__(self, c_in, c_out, k_size, stride, pad, initializer='kaiming'):
+        super(equalized_deconv2d, self).__init__()
+        self.deconv = nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad)
+        if initializer == 'kaiming':    torch.nn.init.kaiming_normal(self.deconv.weight)
+        elif initializer == 'xavier':   torch.nn.init.xavier_normal(self.deconv.weight)
+        self.inv_c = np.sqrt(2.0/(c_in*k_size**2))
+
+    def forward(self, x):
+        return self.deconv(x.mul(self.inv_c))
+
+
+class equalized_linear(nn.Module):
+    def __init__(self, c_in, c_out, initializer='kaiming'):
+        super(equalized_linear, self).__init__()
+        self.linear = nn.Linear(c_in, c_out)
+        if initializer == 'kaiming':    torch.nn.init.kaiming_normal(self.linear.weight)
+        elif initializer == 'xavier':   torch.nn.init.xavier_normal(self.linear.weight)
+        self.inv_c = np.sqrt(2.0/(c_in))
+
+    def forward(self, x):
+        return self.linear(x.mul(self.inv_c))

network.py

@@ -3,13 +3,13 @@
 import torch.nn.functional as F
 import numpy as np
 from torch.autograd import Variable
-from custom_layers import fadein_layer, ConcatTable, minibatch_std_concat_layer, Flatten, pixelwise_norm_layer
+from custom_layers import fadein_layer, ConcatTable, minibatch_std_concat_layer, Flatten, pixelwise_norm_layer, equalized_conv2d, equalized_deconv2d, equalized_linear
 import copy
 
 
 # defined for code simplicity.
 def deconv(layers, c_in, c_out, k_size, stride=1, pad=0, leaky=True, bn=False, wn=False, pixel=False):
-    if wn:  layers.append(nn.utils.weight_norm(nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad), name='weight'))
+    if wn:  layers.append(equalized_deconv2d(c_in, c_out, k_size, stride, pad))
     else:   layers.append(nn.ConvTranspose2d(c_in, c_out, k_size, stride, pad))
     if leaky:   layers.append(nn.LeakyReLU(0.2))
     else:       layers.append(nn.ReLU())
@@ -18,18 +18,19 @@ def deconv(layers, c_in, c_out, k_size, stride=1, pad=0, leaky=True, bn=False, w
     return layers
 
 def conv(layers, c_in, c_out, k_size, stride=1, pad=0, leaky=True, bn=False, wn=False, pixel=False):
-    if wn:  layers.append(nn.utils.weight_norm(nn.Conv2d(c_in, c_out, k_size, stride, pad), name='weight'))
+    if wn:  layers.append(equalized_conv2d(c_in, c_out, k_size, stride, pad, initializer='kaiming'))
     else:   layers.append(nn.Conv2d(c_in, c_out, k_size, stride, pad))
     if leaky:   layers.append(nn.LeakyReLU(0.2))
     else:       layers.append(nn.ReLU())
     if bn:      layers.append(nn.BatchNorm2d(c_out))
     if pixel:   layers.append(pixelwise_norm_layer())
     return layers
 
-def linear(layers, c_in, c_out, sigmoid=True):
+def linear(layers, c_in, c_out, sig=True, wn=False):
     layers.append(Flatten())
-    layers.append(nn.Linear(c_in, c_out))
-    if sigmoid: layers.append(nn.Sigmoid())
+    if wn:      layers.append(equalized_linear(c_in, c_out))
+    else:       layers.append(nn.Linear(c_in, c_out))
+    if sig:     layers.append(nn.Sigmoid())
     return layers
 
 
@@ -41,6 +42,13 @@ def deepcopy_module(module, target):
             new_module[-1].load_state_dict(m.state_dict())         # copy weights
     return new_module
 
+def soft_copy_param(target_link, source_link, tau):
+    ''' soft-copy parameters of a link to another link. '''
+    target_params = dict(target_link.named_parameters())
+    for param_name, param in source_link.named_parameters():
+        target_params[param_name].data = target_params[param_name].data.mul(1.0-tau)
+        target_params[param_name].data = target_params[param_name].data.add(param.data.mul(tau))
+
 def get_module_names(model):
     names = []
     for key, val in model.state_dict().iteritems():
@@ -199,9 +207,9 @@ def last_block(self):
         ndim = self.ndf
         layers = []
         layers.append(minibatch_std_concat_layer())
-        layers = conv(layers, ndim+1, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
-        layers = conv(layers, ndim, ndim, 4, 1, 0, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
-        layers = linear(layers, ndim, 1, self.flag_sigmoid)
+        layers = conv(layers, ndim+1, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
+        layers = conv(layers, ndim, ndim, 4, 1, 0, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
+        layers = linear(layers, ndim, 1, sig=self.flag_sigmoid, wn=self.flag_wn)
         return  nn.Sequential(*layers), ndim
 
     def intermediate_block(self, resl):
@@ -217,18 +225,18 @@ def intermediate_block(self, resl):
                 ndim = ndim/2
         layers = []
         if halving:
-            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
-            layers = deconv(layers, ndim, ndim*2, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
+            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
+            layers = deconv(layers, ndim, ndim*2, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
         else:
-            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
-            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
+            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
+            layers = deconv(layers, ndim, ndim, 3, 1, 1, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
 
         layers.append(nn.AvgPool2d(kernel_size=2))       # scale up by factor of 2.0
         return  nn.Sequential(*layers), ndim, layer_name
 
     def from_rgb_block(self, ndim):
         layers = []
-        layers = conv(layers, self.nc, ndim, 1, 1, 0, self.flag_leaky, self.flag_bn, self.flag_wn, self.flag_pixelwise)
+        layers = conv(layers, self.nc, ndim, 1, 1, 0, self.flag_leaky, self.flag_bn, self.flag_wn, pixel=False)
         return  nn.Sequential(*layers)
 
     def get_init_dis(self):

requirements.txt

@@ -22,7 +22,5 @@ subprocess32==3.2.7
 tensorboardX==0.8
 tensorflow==1.3.0
 tensorflow-tensorboard==0.1.8
-torch==0.2.0.post3
-torchvision==0.1.9
 tqdm==4.19.4
 Werkzeug==0.12.2

trainer.py

@@ -26,6 +26,7 @@ def __init__(self, config):
         self.optimizer = config.optimizer
 
         self.resl = 2           # we start from 2^2 = 4
+        self.smoothing = config.smoothing
         self.max_resl = config.max_resl
         self.trns_tick = config.trns_tick
         self.stab_tick = config.stab_tick
@@ -45,6 +46,20 @@ def __init__(self, config):
 
         # network and cirterion
         self.G = net.Generator(config)
+        self.Gs = net.Generator(config)
+        
+        # shallow copy test.
+        #net.soft_copy_param(self.Gs, self.G, self.smoothing)
+        #for param in self.G.parameters():
+        #    print(param.data.mean())
+        #print('------------------')
+        #for param in self.Gs.parameters():
+        #    print(param.data.mean())
+        #print('------------------')
+
+
+
+
         self.D = net.Discriminator(config)
         print ('Generator structure: ')
         print(self.G.model)
@@ -123,8 +138,8 @@ def resl_scheduler(self):
                     self.fadein['gen'].update_alpha(d_alpha)
                     self.complete['gen'] = self.fadein['gen'].alpha*100
                 self.flag_flush_gen = False
-                self.G.module.flush_network()   # flush and,
-                #self.G.module.freeze_layers()   # freeze.
+                self.G.module.flush_network()   # flush G
+                self.Gs.flush_network()         # flush Gs
                 self.fadein['gen'] = None
                 self.complete['gen'] = 0.0
                 self.phase = 'dtrns'
@@ -134,18 +149,14 @@ def resl_scheduler(self):
                     self.complete['dis'] = self.fadein['dis'].alpha*100
                 self.flag_flush_dis = False
                 self.D.module.flush_network()   # flush and,
-                #self.D.module.freeze_layers()   # freeze.
                 self.fadein['dis'] = None
                 self.complete['dis'] = 0.0
                 self.phase = 'gtrns'
 
             # grow network.
             if floor(self.resl) != prev_resl:
-                #if prev_resl==2:
-                #    self.G.module.freeze_layers()   # freeze.
-                #    self.D.module.freeze_layers()   # freeze.
-                    
                 self.G.module.grow_network(floor(self.resl))
+                self.Gs.grow_network(floor(self.resl))
                 self.D.module.grow_network(floor(self.resl))
                 self.renew_everything()
                 self.fadein['gen'] = self.G.module.model.fadein_block
@@ -248,6 +259,10 @@ def train(self):
                 loss_g.backward()
                 self.opt_g.step()
 
+                # generator smoothing
+                net.soft_copy_param(self.Gs, self.G.module, self.smoothing)
+
+
                 # logging.
                 log_msg = ' [E:{0}][T:{1}][{2:6}/{3:6}]  errD: {4:.4f} | errG: {5:.4f} | [cur:{6:.3f}][resl:{7:4}][{8}][{9:.1f}%][{10:.1f}%]'.format(self.epoch, self.globalTick, self.stack, len(self.loader.dataset), loss_d.data[0], loss_g.data[0], self.resl, int(pow(2,floor(self.resl))), self.phase, self.complete['gen'], self.complete['dis'])
                 tqdm.write(log_msg)
@@ -267,12 +282,15 @@ def train(self):
                 # tensorboard visualization.
                 if self.use_tb:
                     x_test = self.G(self.z_test)
+                    x_test_s = self.Gs(self.z_test)
                     self.tb.add_scalar('data/loss_g', loss_g.data[0], self.globalIter)
                     self.tb.add_scalar('data/loss_d', loss_d.data[0], self.globalIter)
                     self.tb.add_scalar('tick/globalTick', int(self.globalTick), self.globalIter)
-                    self.tb.add_image_grid('grid/x_test', 4, x_test.data.float(), self.globalIter)
-                    self.tb.add_image_grid('grid/x_tilde', 4, self.x_tilde.data.float(), self.globalIter)
-                    self.tb.add_image_grid('grid/x_intp', 1, self.x.data.float(), self.globalIter)
+                    self.tb.add_scalar('tick/cur_resl', int(pow(2,floor(self.resl))), self.globalIter)
+                    self.tb.add_image_grid('grid/x_test', 4, utils.adjust_dyn_range(x_test.data.float(), [-1,1], [0,1]), self.globalIter)
+                    self.tb.add_image_grid('grid/x_test_s', 4, utils.adjust_dyn_range(x_test_s.data.float(), [-1,1], [0,1]), self.globalIter)
+                    self.tb.add_image_grid('grid/x_tilde', 4, utils.adjust_dyn_range(self.x_tilde.data.float(), [-1,1], [0,1]), self.globalIter)
+                    self.tb.add_image_grid('grid/x_intp', 1, utils.adjust_dyn_range(self.x.data.float(), [-1,1], [0,1]), self.globalIter)
 
 
     def snapshot(self, path):

utils.py

@@ -10,6 +10,13 @@
 import time
 
 
+def adjust_dyn_range(x, drange_in, drange_out):
+    if not drange_in == drange_out:
+        scale = float(drange_out[1]-drange_out[0])/float(drange_in[1]-drange_in[0])
+        bias = drange_out[0]-drange_in[0]*scale
+        x = x.mul(scale).add(bias)
+    return x
+
 
 def resize(x, size):
     transform = transforms.Compose([