requires_grad for speedup

Author: junyanz

models/base_model.py

@@ -108,14 +108,15 @@ def __patch_instance_norm_state_dict(self, state_dict, module, keys, i=0):
     def load_networks(self, which_epoch):
         for name in self.model_names:
             if isinstance(name, str):
-                save_filename = '%s_net_%s.pth' % (which_epoch, name)
-                save_path = os.path.join(self.save_dir, save_filename)
+                load_filename = '%s_net_%s.pth' % (which_epoch, name)
+                load_path = os.path.join(self.save_dir, load_filename)
                 net = getattr(self, 'net' + name)
                 if isinstance(net, torch.nn.DataParallel):
                     net = net.module
+                print('loading the model from %s' % load_path)
                 # if you are using PyTorch newer than 0.4 (e.g., built from
                 # GitHub source), you can remove str() on self.device
-                state_dict = torch.load(save_path, map_location=str(self.device))
+                state_dict = torch.load(load_path, map_location=str(self.device))
                 # patch InstanceNorm checkpoints prior to 0.4
                 for key in list(state_dict.keys()):  # need to copy keys here because we mutate in loop
                     self.__patch_instance_norm_state_dict(state_dict, net, key.split('.'))
@@ -134,3 +135,12 @@ def print_networks(self, verbose):
                     print(net)
                 print('[Network %s] Total number of parameters : %.3f M' % (name, num_params / 1e6))
         print('-----------------------------------------------')
+
+    # set requies_grad=Fasle to avoid computation
+    def set_requires_grad(self, nets, requires_grad=False):
+        if not isinstance(nets, list):
+            nets = [nets]
+        for net in nets:
+            if net is not None:
+                for param in net.parameters():
+                    param.requires_grad = requires_grad

models/cycle_gan_model.py

@@ -113,13 +113,10 @@ def backward_G(self):
 
         # GAN loss D_A(G_A(A))
         self.loss_G_A = self.criterionGAN(self.netD_A(self.fake_B), True)
-
         # GAN loss D_B(G_B(B))
         self.loss_G_B = self.criterionGAN(self.netD_B(self.fake_A), True)
-
         # Forward cycle loss
         self.loss_cycle_A = self.criterionCycle(self.rec_A, self.real_A) * lambda_A
-
         # Backward cycle loss
         self.loss_cycle_B = self.criterionCycle(self.rec_B, self.real_B) * lambda_B
         # combined loss
@@ -130,10 +127,12 @@ def optimize_parameters(self):
         # forward
         self.forward()
         # G_A and G_B
+        self.set_requires_grad([self.netD_A, self.netD_B], False)
         self.optimizer_G.zero_grad()
         self.backward_G()
         self.optimizer_G.step()
         # D_A and D_B
+        self.set_requires_grad([self.netD_A, self.netD_B], True)
         self.optimizer_D.zero_grad()
         self.backward_D_A()
         self.backward_D_B()

models/pix2pix_model.py

@@ -86,11 +86,14 @@ def backward_G(self):
 
     def optimize_parameters(self):
         self.forward()
-
+        # update D
+        self.set_requires_grad(self.netD, True)
         self.optimizer_D.zero_grad()
         self.backward_D()
         self.optimizer_D.step()
 
+        # update G
+        self.set_requires_grad(self.netD, False)
         self.optimizer_G.zero_grad()
         self.backward_G()
         self.optimizer_G.step()