add strict_load option

WuyiWong · Jun 9, 2019 · 876f282 · 876f282
1 parent 1c2f6bc
commit 876f282
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 5 deletions.
diff --git a/codes/models/SRGAN_model.py b/codes/models/SRGAN_model.py
@@ -232,11 +232,11 @@ def load(self):
         load_path_G = self.opt['path']['pretrain_model_G']
         if load_path_G is not None:
             logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
-            self.load_network(load_path_G, self.netG)
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
         load_path_D = self.opt['path']['pretrain_model_D']
         if self.opt['is_train'] and load_path_D is not None:
             logger.info('Loading pretrained model for D [{:s}] ...'.format(load_path_D))
-            self.load_network(load_path_D, self.netD)
+            self.load_network(load_path_D, self.netD, self.opt['path']['strict_load'])
 
     def save(self, iter_step):
         self.save_network(self.netG, 'G', iter_step)

diff --git a/codes/models/SRRaGAN_model.py b/codes/models/SRRaGAN_model.py
@@ -243,11 +243,11 @@ def load(self):
         load_path_G = self.opt['path']['pretrain_model_G']
         if load_path_G is not None:
             logger.info('Loading pretrained model for G [{:s}] ...'.format(load_path_G))
-            self.load_network(load_path_G, self.netG)
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
         load_path_D = self.opt['path']['pretrain_model_D']
         if self.opt['is_train'] and load_path_D is not None:
             logger.info('Loading pretrained model for D [{:s}] ...'.format(load_path_D))
-            self.load_network(load_path_D, self.netD)
+            self.load_network(load_path_D, self.netD, self.opt['path']['strict_load'])
 
     def save(self, iter_step):
         self.save_network(self.netG, 'G', iter_step)

diff --git a/codes/models/SR_model.py b/codes/models/SR_model.py
@@ -168,7 +168,7 @@ def load(self):
         load_path_G = self.opt['path']['pretrain_model_G']
         if load_path_G is not None:
             logger.info('Loading model for G [{:s}] ...'.format(load_path_G))
-            self.load_network(load_path_G, self.netG)
+            self.load_network(load_path_G, self.netG, self.opt['path']['strict_load'])
 
     def save(self, iter_label):
         self.save_network(self.netG, 'G', iter_label)
diff --git a/codes/options/train/train_SRResNet.yml b/codes/options/train/train_SRResNet.yml
@@ -42,6 +42,7 @@ network_G:
 #### path
 path:
   pretrain_model_G: ~
+  strict_load: true
   resume_state: ~
 
 #### training settings: learning rate scheme, loss