Smt

datasets/synthetic.py

@@ -1,15 +1,16 @@
 import os
 import sys
+from os.path import dirname as up
+
+ROOT_DIR = up(up(up(os.path.realpath(__file__))))
+if ROOT_DIR not in sys.path: sys.path.append(ROOT_DIR)
 
 import numpy as np
 from scipy.stats import norm
 from PIL import Image
 import torch
 import torchvision
 
-ROOT_DIR = os.path.dirname(os.getcwd())
-if ROOT_DIR not in sys.path: sys.path.append(ROOT_DIR)
-
 import DeepSparseCoding.utils.data_processing as dp
 
 class SyntheticImages(torchvision.datasets.vision.VisionDataset):

models/base_model.py

@@ -1,23 +1,28 @@
 import os
+import subprocess
+import pprint
 
 import numpy as np
 import torch
 
+from DeepSparseCoding.utils.file_utils import summary_string
 from DeepSparseCoding.utils.file_utils import Logger
+from DeepSparseCoding.utils.run_utils import compute_conv_output_shape
+import DeepSparseCoding.utils.loaders as loaders
 
 
 class BaseModel(object):
     def setup(self, params, logger=None):
         """
         Setup required model components
-        #TODO: log system info, including git commit hash
         """
         self.load_params(params)
         self.check_params()
         self.make_dirs()
         if logger is None:
             self.init_logging()
             self.log_params()
+            self.logger.log_info(self.get_env_details())
         else:
             self.logger = logger
 
@@ -92,24 +97,106 @@ def log_params(self, params=None):
             dump_obj = self.params.__dict__
         self.logger.log_params(dump_obj)
 
-    def log_info(self, string):
-        """Log input string"""
-        self.logger.log_info(string)
+    def get_train_stats(self, batch_step=None):
+        """
+        Get default statistics about current training run
+
+        Keyword arguments:
+            batch_step: [int] current batch iteration. The default assumes that training has finished.
+        """
+        if batch_step is None:
+            batch_step = self.params.num_batches
+        epoch = batch_step / self.params.batches_per_epoch
+        stat_dict = {
+            'epoch':int(epoch),
+            'batch_step':batch_step,
+            'train_progress':np.round(batch_step/self.params.num_batches, 3),
+        }
+        return stat_dict
 
-    def write_checkpoint(self):
-        """Write checkpoints"""
-        torch.save(self.state_dict(), self.params.cp_latest_filename)
-        self.log_info('Full model saved in file %s'%self.params.cp_latest_filename)
+    def get_env_details(self):
+        env = {}
+        for k in ['SYSTEMROOT', 'PATH']:
+            v = os.environ.get(k)
+            if v is not None:
+                env[k] = v
+        commit_cmd = ['git', 'rev-parse', 'HEAD']
+        commit = subprocess.Popen(commit_cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
+        commit = commit.strip().decode('ascii')
+        branch_cmd = ['git', 'rev-parse', '--abbrev-ref', 'HEAD']
+        branch = subprocess.Popen(branch_cmd, stdout=subprocess.PIPE, env=env).communicate()[0]
+        branch = branch.strip().decode('ascii')
+        system_details = os.uname()
+        out_dict = {
+            'current_branch':branch,
+            'current_commit_hash':commit,
+            'sysname':system_details.sysname,
+            'release':system_details.release,
+            'machine':system_details.machine
+        }
+        if torch.cuda.is_available():
+            out_dict['gpu_device'] = torch.cuda.get_device_name(0)
+        return out_dict
 
-    def load_checkpoint(self, cp_file=None):
+    def log_architecture_details(self):
+        """
+        Log model architecture with computed output sizes and number of parameters for each layer
+        """
+        architecture_string = '<architecture>\n'+summary_string(
+            self,
+            input_size=tuple(self.params.data_shape),
+            batch_size=self.params.batch_size,
+            device=self.params.device,
+            dtype=torch.FloatTensor
+        )[0]
+        architecture_string  += '\n</architecture>'
+        self.logger.log_string(architecture_string)
+
+    def write_checkpoint(self, batch_step=None):
+        """
+        Write checkpoints
+
+        Keyword arguments:
+            batch_step: [int] current batch iteration. The default assumes that training has finished.
+        """
+        output_dict = {}
+        if(self.params.model_type.lower() == 'ensemble'):
+            for module in self:
+                module_name = module.params.submodule_name
+                output_dict[module_name+'_module_state_dict'] = module.state_dict()
+                output_dict[module_name+'_optimizer_state_dict'] = module.optimizer.state_dict()
+        else:
+            output_dict['model_state_dict'] =  self.state_dict()
+            module_state_dict_name = 'optimizer_state_dict'
+            output_dict[module_state_dict_name] = self.optimizer.state_dict(),
+            ## TODO: Save scheduler state dict as well
+        training_stats = self.get_train_stats(batch_step)
+        output_dict.update(training_stats)
+        torch.save(output_dict, self.params.cp_latest_filename)
+        self.logger.log_string('Full model saved in file %s'%self.params.cp_latest_filename)
+
+    def get_checkpoint_from_log(self, logfile):
+        model_params = loaders.load_params_from_log(logfile)
+        checkpoint = torch.load(model_params.cp_latest_filename)
+        return checkpoint
+
+    def load_checkpoint(self, cp_file=None, load_optimizer=False):
         """
         Load checkpoint
-        Inputs:
-          model_dir: String specifying the path to the checkpoint
+        Keyword arguments:
+          model_dir: [str] specifying the path to the checkpoint
         """
         if cp_file is None:
             cp_file = self.params.cp_latest_filename
-        return self.load_state_dict(torch.load(cp_file))
+        checkpoint = torch.load(cp_file)
+        if load_optimizer:
+            self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+        self.load_state_dict(checkpoint['model_state_dict'])
+        _ = checkpoint.pop('optimizer_state_dict', None)
+        _ = checkpoint.pop('model_state_dict', None)
+        training_status = pprint.pformat(checkpoint, compact=True)#, sort_dicts=True #TODO: Python 3.8 adds the sort_dicts parameter
+        out_str = f'Loaded checkpoint from {cp_file} with the following stats:\n{training_status}'
+        return out_str
 
     def get_optimizer(self, optimizer_params, trainable_variables):
         optimizer_name = optimizer_params.optimizer.name
@@ -129,8 +216,8 @@ def get_optimizer(self, optimizer_params, trainable_variables):
 
     def setup_optimizer(self):
         self.optimizer = self.get_optimizer(
-                optimizer_params=self.params,
-                trainable_variables=self.parameters())
+            optimizer_params=self.params,
+            trainable_variables=self.parameters())
         self.scheduler = torch.optim.lr_scheduler.MultiStepLR(
             self.optimizer,
             milestones=self.params.optimizer.milestones,
@@ -143,21 +230,18 @@ def print_update(self, input_data, input_labels=None, batch_step=0):
           input_data: data object containing the current image batch
           input_labels: data object containing the current label batch
           batch_step: current batch number within the schedule
-        NOTE: For the analysis code to parse update statistics, the self.js_dumpstring() call
-          must receive a dict object. Additionally, the self.js_dumpstring() output must be
-          logged with <stats> </stats> tags.
-          For example: logging.info('<stats>'+self.js_dumpstring(output_dictionary)+'</stats>')
+        NOTE: For the analysis code to parse update statistics,
+        the logger.log_stats() function must be used
         """
         update_dict = self.generate_update_dict(input_data, input_labels, batch_step)
-        js_str = self.js_dumpstring(update_dict)
-        self.log_info('<stats>'+js_str+'</stats>')
+        self.logger.log_stats(update_dict)
 
     def generate_update_dict(self, input_data, input_labels=None, batch_step=0, update_dict=None):
         """
         Generates a dictionary to be logged in the print_update function
         """
         if update_dict is None:
-            update_dict = dict()
+            update_dict = self.get_train_stats(batch_step)
         for param_name, param_var in self.named_parameters():
             grad = param_var.grad
             update_dict[param_name+'_grad_max_mean_min'] = [