first commit

1 files changed, 86 insertions, 0 deletions

diff --git a/models/base_model.py b/models/base_model.py
new file mode 100755
index 0000000..d3879d0
--- /dev/null
+++ b/models/base_model.py
@@ -0,0 +1,86 @@
+### Copyright (C) 2017 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).
+import os
+import torch
+
+class BaseModel(torch.nn.Module):
+    def name(self):
+        return 'BaseModel'
+
+    def initialize(self, opt):
+        self.opt = opt
+        self.gpu_ids = opt.gpu_ids
+        self.isTrain = opt.isTrain
+        self.Tensor = torch.cuda.FloatTensor if self.gpu_ids else torch.Tensor
+        self.save_dir = os.path.join(opt.checkpoints_dir, opt.name)
+
+    def set_input(self, input):
+        self.input = input
+
+    def forward(self):
+        pass
+
+    # used in test time, no backprop
+    def test(self):
+        pass
+
+    def get_image_paths(self):
+        pass
+
+    def optimize_parameters(self):
+        pass
+
+    def get_current_visuals(self):
+        return self.input
+
+    def get_current_errors(self):
+        return {}
+
+    def save(self, label):
+        pass
+
+    # helper saving function that can be used by subclasses
+    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 len(gpu_ids) and torch.cuda.is_available():
+            network.cuda()
+
+    # helper loading function that can be used by subclasses
+    def load_network(self, network, network_label, epoch_label, save_dir=''):        
+        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
+        if not save_dir:
+            save_dir = self.save_dir
+        save_path = os.path.join(save_dir, save_filename)        
+        if not os.path.isfile(save_path):
+            print('%s not exists yet!' % save_path)
+            if network_label == 'G':
+                raise('Generator must exist!')
+        else:
+            #network.load_state_dict(torch.load(save_path))
+            try:
+                network.load_state_dict(torch.load(save_path))
+            except:   
+                pretrained_dict = torch.load(save_path)                
+                model_dict = network.state_dict()
+                try:
+                    pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}                    
+                    network.load_state_dict(pretrained_dict)
+                    print('Pretrained network %s has excessive layers; Only loading layers that are used' % network_label)
+                except:
+                    print('Pretrained network %s has fewer layers; The following are not initialized:' % network_label)
+                    from sets import Set
+                    not_initialized = Set()
+                    for k, v in pretrained_dict.items():                      
+                        if v.size() == model_dict[k].size():
+                            model_dict[k] = v
+
+                    for k, v in model_dict.items():
+                        if k not in pretrained_dict or v.size() != pretrained_dict[k].size():
+                            not_initialized.add(k.split('.')[0])                            
+                    print(sorted(not_initialized))
+                    network.load_state_dict(model_dict)                  
+
+    def update_learning_rate():
+        pass