Fixed speed issue with pytorch 0.4

Author: mingyuliutw

networks.py

@@ -305,7 +305,8 @@ def __init__(self, input_dim ,output_dim, kernel_size, stride,
         if norm == 'bn':
             self.norm = nn.BatchNorm2d(norm_dim)
         elif norm == 'in':
-            self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
+            #self.norm = nn.InstanceNorm2d(norm_dim, track_running_stats=True)
+            self.norm = nn.InstanceNorm2d(norm_dim)
         elif norm == 'ln':
             self.norm = LayerNorm(norm_dim)
         elif norm == 'adain':
@@ -474,6 +475,7 @@ def forward(self, x):
     def __repr__(self):
         return self.__class__.__name__ + '(' + str(self.num_features) + ')'
 
+
 class LayerNorm(nn.Module):
     def __init__(self, num_features, eps=1e-5, affine=True):
         super(LayerNorm, self).__init__()
@@ -487,11 +489,16 @@ def __init__(self, num_features, eps=1e-5, affine=True):
 
     def forward(self, x):
         shape = [-1] + [1] * (x.dim() - 1)
-        mean = x.view(x.size(0), -1).mean(1).view(*shape)
-        std = x.view(x.size(0), -1).std(1).view(*shape)
+        # print(x.size())
+        # mean = x.view(x.size(0), -1).mean(1).view(*shape)
+        # std = x.view(x.size(0), -1).std(1).view(*shape)
+        mean = x.view(-1).mean().view(*shape)
+        std = x.view(-1).std().view(*shape)
+
         x = (x - mean) / (std + self.eps)
 
         if self.affine:
             shape = [1, -1] + [1] * (x.dim() - 2)
             x = x * self.gamma.view(*shape) + self.beta.view(*shape)
         return x
+

train.py

@@ -2,7 +2,7 @@
 Copyright (C) 2018 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).
 """
-from utils import get_all_data_loaders, prepare_sub_folder, write_html, write_loss, get_config, write_2images
+from utils import get_all_data_loaders, prepare_sub_folder, write_html, write_loss, get_config, write_2images, Timer
 import argparse
 from torch.autograd import Variable
 from trainer import MUNIT_Trainer, UNIT_Trainer
@@ -41,10 +41,10 @@
     sys.exit("Only support MUNIT|UNIT")
 trainer.cuda()
 train_loader_a, train_loader_b, test_loader_a, test_loader_b = get_all_data_loaders(config)
-train_display_images_a = Variable(torch.stack([train_loader_a.dataset[i] for i in range(display_size)]).cuda())
-train_display_images_b = Variable(torch.stack([train_loader_b.dataset[i] for i in range(display_size)]).cuda())
-test_display_images_a = Variable(torch.stack([test_loader_a.dataset[i] for i in range(display_size)]).cuda())
-test_display_images_b = Variable(torch.stack([test_loader_b.dataset[i] for i in range(display_size)]).cuda())
+train_display_images_a = torch.stack([train_loader_a.dataset[i] for i in range(display_size)]).cuda()
+train_display_images_b = torch.stack([train_loader_b.dataset[i] for i in range(display_size)]).cuda()
+test_display_images_a = torch.stack([test_loader_a.dataset[i] for i in range(display_size)]).cuda()
+test_display_images_b = torch.stack([test_loader_b.dataset[i] for i in range(display_size)]).cuda()
 
 # Setup logger and output folders
 model_name = os.path.splitext(os.path.basename(opts.config))[0]
@@ -58,11 +58,13 @@
 while True:
     for it, (images_a, images_b) in enumerate(zip(train_loader_a, train_loader_b)):
         trainer.update_learning_rate()
-        images_a, images_b = Variable(images_a.cuda()), Variable(images_b.cuda())
+        images_a, images_b = images_a.cuda().detach(), images_b.cuda().detach()
 
-        # Main training code
-        trainer.dis_update(images_a, images_b, config)
-        trainer.gen_update(images_a, images_b, config)
+        with Timer("Elapsed time in update: %f"):
+            # Main training code
+            trainer.dis_update(images_a, images_b, config)
+            trainer.gen_update(images_a, images_b, config)
+            torch.cuda.synchronize()
 
         # Dump training stats in log file
         if (iterations + 1) % config['log_iter'] == 0:

utils.py

@@ -16,6 +16,7 @@
 import yaml
 import numpy as np
 import torch.nn.init as init
+import time
 # Methods
 # get_all_data_loaders      : primary data loader interface (load trainA, testA, trainB, testB)
 # get_data_loader_list      : list-based data loader
@@ -277,3 +278,16 @@ def init_fun(m):
                 init.constant_(m.bias.data, 0.0)
 
     return init_fun
+
+
+class Timer:
+    def __init__(self, msg):
+        self.msg = msg
+        self.start_time = None
+
+    def __enter__(self):
+        self.start_time = time.time()
+
+    def __exit__(self, exc_type, exc_value, exc_tb):
+        print(self.msg % (time.time() - self.start_time))
+