Merge branch 'master' of https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

5 files changed, 16 insertions, 17 deletions

diff --git a/.gitignore b/.gitignore
index 4a26633..f99f990 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,3 +1,4 @@
+debug*
 datasets/
 checkpoints/
 results/
diff --git a/models/base_model.py b/models/base_model.py
index 0ea83d8..ce18635 100644
--- a/models/base_model.py
+++ b/models/base_model.py
@@ -39,12 +39,12 @@ class BaseModel():
         pass
 
     # helper saving function that can be used by subclasses
-    def save_network(self, network, network_label, epoch_label, use_gpu):
+    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 use_gpu and torch.cuda.is_available():
-            network.cuda()
+        if len(gpu_ids) and torch.cuda.is_available():
+            network.cuda(device_id=gpu_ids[0])
 
     # helper loading function that can be used by subclasses
     def load_network(self, network, network_label, epoch_label):
diff --git a/models/cycle_gan_model.py b/models/cycle_gan_model.py
index c3b5b72..eb1c443 100644
--- a/models/cycle_gan_model.py
+++ b/models/cycle_gan_model.py
@@ -202,11 +202,10 @@ class CycleGANModel(BaseModel):
                                 ('real_B', real_B), ('fake_A', fake_A), ('rec_B', rec_B)])
 
     def save(self, label):
-        use_gpu = self.gpu_ids is not None
-        self.save_network(self.netG_A, 'G_A', label, use_gpu)
-        self.save_network(self.netD_A, 'D_A', label, use_gpu)
-        self.save_network(self.netG_B, 'G_B', label, use_gpu)
-        self.save_network(self.netD_B, 'D_B', label, use_gpu)
+        self.save_network(self.netG_A, 'G_A', label, self.gpu_ids)
+        self.save_network(self.netD_A, 'D_A', label, self.gpu_ids)
+        self.save_network(self.netG_B, 'G_B', label, self.gpu_ids)
+        self.save_network(self.netD_B, 'D_B', label, self.gpu_ids)
 
     def update_learning_rate(self):
         lrd = self.opt.lr / self.opt.niter_decay
diff --git a/models/networks.py b/models/networks.py
index 5624a7b..6777359 100644
--- a/models/networks.py
+++ b/models/networks.py
@@ -32,7 +32,7 @@ def define_G(input_nc, output_nc, ngf, which_model_netG, norm, gpu_ids=[]):
     if which_model_netG == 'resnet_9blocks':
         netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, n_blocks=9, gpu_ids=gpu_ids)
     elif which_model_netG == 'resnet_6blocks':
-        netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, 6, gpu_ids=gpu_ids)
+        netG = ResnetGenerator(input_nc, output_nc, ngf, norm_layer, n_blocks=6, gpu_ids=gpu_ids)
     elif which_model_netG == 'unet_128':
         netG = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer, gpu_ids=gpu_ids)
     elif which_model_netG == 'unet_256':
@@ -40,7 +40,7 @@ def define_G(input_nc, output_nc, ngf, which_model_netG, norm, gpu_ids=[]):
     else:
         print('Generator model name [%s] is not recognized' % which_model_netG)
     if len(gpu_ids) > 0:
-        netG.cuda()
+        netG.cuda(device_id=gpu_ids[0])
     netG.apply(weights_init)
     return netG
 
@@ -59,7 +59,7 @@ def define_D(input_nc, ndf, which_model_netD,
         print('Discriminator model name [%s] is not recognized' %
               which_model_netD)
     if use_gpu:
-        netD.cuda()
+        netD.cuda(device_id=gpu_ids[0])
     netD.apply(weights_init)
     return netD
 
@@ -213,8 +213,7 @@ class UnetGenerator(nn.Module):
         unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, unet_block)
         unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, unet_block)
         unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, unet_block)
-        unet_block = UnetSkipConnectionBlock(input_nc, ngf, unet_block,
-                                             outermost=True)
+        unet_block = UnetSkipConnectionBlock(output_nc, ngf, unet_block, outermost=True)
 
         self.model = unet_block
 
diff --git a/models/pix2pix_model.py b/models/pix2pix_model.py
index 1d89b29..3bdd237 100644
--- a/models/pix2pix_model.py
+++ b/models/pix2pix_model.py
@@ -8,7 +8,7 @@ import util.util as util
 from util.image_pool import ImagePool
 from .base_model import BaseModel
 from . import networks
-
+from pdb import set_trace as st
 
 class Pix2PixModel(BaseModel):
     def name(self):
@@ -108,6 +108,7 @@ class Pix2PixModel(BaseModel):
         self.loss_G.backward()
 
     def optimize_parameters(self):
+        # st()
         self.forward()
 
         self.optimizer_D.zero_grad()
@@ -132,9 +133,8 @@ class Pix2PixModel(BaseModel):
         return OrderedDict([('real_A', real_A), ('fake_B', fake_B), ('real_B', real_B)])
 
     def save(self, label):
-        use_gpu = self.gpu_ids is not None
-        self.save_network(self.netG, 'G', label, use_gpu)
-        self.save_network(self.netD, 'D', label, use_gpu)
+        self.save_network(self.netG, 'G', label, self.gpu_ids)
+        self.save_network(self.netD, 'D', label, self.gpu_ids)
 
     def update_learning_rate(self):
         lrd = self.opt.lr / self.opt.niter_decay