move model.setup into train and test

Author: junyanz

models/cycle_gan_model.py

@@ -61,8 +61,6 @@ def initialize(self, opt):
             self.optimizers.append(self.optimizer_G)
             self.optimizers.append(self.optimizer_D)
 
-        self.setup(opt)
-
     def set_input(self, input):
         AtoB = self.opt.which_direction == 'AtoB'
         self.real_A = input['A' if AtoB else 'B'].to(self.device)

models/pix2pix_model.py

@@ -45,8 +45,6 @@ def initialize(self, opt):
             self.optimizers.append(self.optimizer_G)
             self.optimizers.append(self.optimizer_D)
 
-        self.setup(opt)
-
     def set_input(self, input):
         AtoB = self.opt.which_direction == 'AtoB'
         self.real_A = input['A' if AtoB else 'B'].to(self.device)

models/test_model.py

@@ -22,7 +22,6 @@ def initialize(self, opt):
                                       opt.norm, not opt.no_dropout,
                                       opt.init_type,
                                       self.gpu_ids)
-        self.setup(opt)
 
     def set_input(self, input):
         # we need to use single_dataset mode

test.py

@@ -16,6 +16,7 @@
     data_loader = CreateDataLoader(opt)
     dataset = data_loader.load_data()
     model = create_model(opt)
+    model.setup(opt)
     # create website
     web_dir = os.path.join(opt.results_dir, opt.name, '%s_%s' % (opt.phase, opt.which_epoch))
     webpage = html.HTML(web_dir, 'Experiment = %s, Phase = %s, Epoch = %s' % (opt.name, opt.phase, opt.which_epoch))

train.py

@@ -12,6 +12,7 @@
     print('#training images = %d' % dataset_size)
 
     model = create_model(opt)
+    model.setup(opt)
     visualizer = Visualizer(opt)
     total_steps = 0