Merge pull request #2 from CAAI/raphael-dev

merge new UserConfig class to main

Author: claesnl

rhtorch/callbacks/plotting.py

@@ -5,7 +5,7 @@
 import torch
 
 
-def plot_inline(d1, d2, d3, color_channel_axis=0):
+def plot_inline(d1, d2, d3, color_channel_axis=0, vmin=None, vmax=None):
     """
     Parameters
     ----------
@@ -19,28 +19,28 @@ def plot_inline(d1, d2, d3, color_channel_axis=0):
         Axis for color channel in the numpy array . 
         Default is 0 for Pytorch models (cc, dimx, dimy, dimz)
         Use 3 for TF models (dimx, dimy, dimz, cc)
+    vmin : Lower bound for color channel. Default (None) used to plot full range
+    vmax : Upper bound for color channel. Default (None) used to plot full range
 
     """
     # If input has more than 1 color channel, use only the first
     if d1.shape[color_channel_axis] > 1:
         d1 = d1[0,...] if color_channel_axis == 0 else d1[...,0]
-        d1 = torch.unsqueeze(d1,color_channel_axis)
-    d_arr = d_arr = np.concatenate((d1, d2, d3), color_channel_axis)
+        d1 = torch.unsqueeze(d1, color_channel_axis)
+    d_arr = np.concatenate((d1, d2, d3), color_channel_axis)
     num_dat = d_arr.shape[color_channel_axis]
 
     fig, ax = plt.subplots(1, num_dat, gridspec_kw={'wspace': 0, 'hspace': 0})
     slice_i = int(d1.size(1) / 2)
     orient = 0
     text_pos = d1.size(2) * 0.98
 
-    # make a list of subplot titles - may need several input subtitles
-    titles = [f"Input{i+1}" for i in range(d1.size(color_channel_axis))]
-    titles.extend(['Target', 'Prediction'])
+    titles = ['Input', 'Target', 'Prediction']
 
     for idx in range(num_dat):
         single_data = d_arr.take(indices=idx, axis=color_channel_axis) 
         ax[idx].imshow(single_data.take(indices=slice_i, axis=orient), 
-                       cmap='gray', vmin=0, vmax=1)
+                       cmap='gray', vmin=vmin, vmax=vmax)
         ax[idx].axis('off')
         ax[idx].text(3, text_pos, titles[idx], color='white', fontsize=12)
 
@@ -50,10 +50,12 @@ def plot_inline(d1, d2, d3, color_channel_axis=0):
     return wandb_im
 
 class ImagePredictionLogger(Callback):
-    def __init__(self, val_dataloader):
+    def __init__(self, val_dataloader, config=None):
         super().__init__()
         self.X, self.y = next(iter(val_dataloader))
 
+        # TODO: Read config file to parse vmin, vmax or e.g. custom titles
+        
     def on_validation_epoch_end(self, trainer, pl_module):
         # Dataloader loads on CPU --> pass to GPU
         X = self.X.to(device=pl_module.device)

rhtorch/config_utils.py

@@ -1,70 +1,78 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
 import ruamel.yaml as yaml
 from datetime import datetime
 from pathlib import Path
 import torch
 from rhtorch.version import __version__
 import socket
 
-loss_map = {'MeanAbsoluteError': 'mae',
-            'MeanSquaredError': 'mse',
-            'huber_loss': 'huber',
-            'BCEWithLogitsLoss': 'BCE'}
+class UserConfig:
+    def __init__(self, rootdir, arguments=None):
+        self.rootdir = rootdir
+        self.config_file = self.is_path(arguments.config)
+        self.args = arguments
+        
+        # load default configs
+        default_config_file = Path(__file__).parent.joinpath('default_config.yaml')
+        with open(default_config_file) as dcf:
+            self.default_params = yaml.load(dcf, Loader=yaml.Loader)
+            
+        # load user config file
+        with open(self.config_file) as cf:
+            self.hparams = yaml.load(cf, Loader=yaml.RoundTripLoader)
+        
+        # merge the two dicts
+        self.merge_dicts()
+        
+        # sanity check on data_folder provided by user
+        self.data_path = self.is_path(self.hparams['data_folder'])
+        
+        # make model name
+        self.fill_additional_info()
+        self.create_model_name()
+            
+    def is_path(self, path):
+        # check for path - assuming absolute path was given
+        filepath = Path(path)
+        if not filepath.exists():
+            # assuming path was given relative to rootdir
+            filepath = self.rootdir.joinpath(filepath)
+        if not filepath.exists():
+            raise FileNotFoundError(f"{path} not found. Define relative to project directory or as absolute path in config file/argument passing.")
+        
+        return filepath
 
+    def merge_dicts(self):
+        """ adds to the user_params dictionnary any missing key from the default params """
+        
+        for k, v in self.default_params.items():
+            if k not in self.hparams:
+                self.hparams[k] = v
+            ### TO DO - ENSURE NOT COPYING IRRELEVANT DATA e.g. GAN parameters if model is AE
 
-def load_model_config(rootdir, arguments):
+    def fill_additional_info(self):
+        # additional info from args and miscellaneous to save in config
+        self.hparams['build_date'] = datetime.now().strftime("%Y%m%d-%H%M%S")
+        self.hparams['project_dir'] = str(self.rootdir)
+        self.hparams['data_folder'] = str(self.data_path)
+        self.hparams['config_file'] = str(self.config_file)
+        self.hparams['k_fold'] = self.args.kfold
+        self.hparams['GPUs'] = torch.cuda.device_count()
+        self.hparams['global_batch_size'] = self.hparams['batch_size'] * self.hparams['GPUs']
+        self.hparams['rhtorch_version'] = __version__
+        self.hparams['hostname'] = socket.gethostname()
 
-    # check for config_file
-    config_file = Path(arguments.config)
-    if not config_file.exists():
-        config_file = rootdir.joinpath(config_file)
-    if not config_file.exists():
-        raise FileNotFoundError("Config file not found. Define relative to project directory or as absolute path in config file")
-    
-    # read the config file
-    with open(config_file) as file:
-        config = yaml.load(file, Loader=yaml.RoundTripLoader)
+    def create_model_name(self):
 
-    data_shape = 'x'.join(map(str, config['data_shape']))
-    base_name = f"{config['module']}_{config['version_name']}_{config['data_generator']}"
-    dat_name = f"bz{config['batch_size']}_{data_shape}"
-    full_name = f"{base_name}_{dat_name}_k{arguments.kfold}_e{config['epoch']}"
-    
-    # check for data folder
-    data_folder = Path(config['data_folder'])
-    if not data_folder.exists():
-        # try relative to project dir - in this case overwrite config
-        data_folder = rootdir.joinpath(config['data_folder'])
-    if not data_folder.exists():
-        raise FileNotFoundError("Data path not found. Define relative to the project directory or as absolute path in config file")
+        data_shape = 'x'.join(map(str, self.hparams['data_shape']))
+        base_name = f"{self.hparams['module']}_{self.hparams['version_name']}_{self.hparams['data_generator']}"
+        dat_name = f"bz{self.hparams['batch_size']}_{data_shape}"
+        self.hparams['model_name'] = f"{base_name}_{dat_name}_k{self.args.kfold}_e{self.hparams['epoch']}"
 
-    # additional info from args and miscellaneous to save in config
-    config['build_date'] = datetime.now().strftime("%Y-%m-%d %H.%M.%S")
-    config['model_name'] = full_name
-    config['project_dir'] = str(rootdir)
-    config['data_folder'] = str(data_folder)
-    config['config_file'] = str(config_file)
-    config['k_fold'] = arguments.kfold
-    if 'precision' not in config:
-        config['precision'] = 32
-    config['GPUs'] = torch.cuda.device_count()
-    config['global_batch_size'] = config['batch_size'] * config['GPUs']
-    config['rhtorch_version'] = __version__
-    config['hostname'] = socket.gethostname()
-    if 'acc_grad_batches' not in config:
-        config['acc_grad_batches'] = 1
-
-    return config
-
-
-def copy_model_config(path, config, append_timestamp=False):
-    model_name = config['model_name']
-    if append_timestamp:
-        timestamp = config['build_date'].replace(' ','_')
-        config_file = path.joinpath(f"config_{model_name}_{timestamp}.yaml")
-    else:
-        config_file = path.joinpath(f"config_{model_name}.yaml") 
-    config.yaml_set_start_comment(f'Config file for {model_name}')
-    with open(config_file, 'w') as file:
-        yaml.dump(config, file, Dumper=yaml.RoundTripDumper)
+    def save_copy(self, output_dir, append_timestamp=False):
+        model_name = self.hparams['model_name']
+        timestamp = f"_{self.hparams['build_date']}" if append_timestamp else ""
+        save_config_file_name = f"config_{model_name}{timestamp}"        
+        config_file = output_dir.joinpath(save_config_file_name + ".yaml")
+        self.hparams.yaml_set_start_comment(f'Config file for {model_name}')
+        with open(config_file, 'w') as file:
+            yaml.dump(self.hparams, file, Dumper=yaml.RoundTripDumper)

rhtorch/default_config.yaml

@@ -0,0 +1,49 @@
+# main config to train models
+project_name: YOUR_PROJECT_NAME       # Used for WANDB
+version_name: v0                      # Make run unique by changing this counter
+
+# main model:
+precision: 32
+epoch: 100
+batch_size: 1
+acc_grad_batches: 1
+module: LightningAE
+
+# generator
+generator: UNet3DFullConv
+# depth: 4
+# initial_num_filters: 64
+g_activation: ReLU
+g_optimizer: Adam
+g_lr: 1e-4
+#lr_scheduler: 'exponential_decay_0.01'
+g_loss: MeanAbsoluteError
+
+# transfer learning
+pretrained_generator: null           # absolute path to .pt or .ckpt
+freeze_encoder: False
+
+# discriminator - will be used if model is GAN
+discriminator: ConvNetDiscriminator
+d_optimizer: Adam
+d_lr: 2e-4
+d_loss: BCEWithLogitsLoss
+
+# data:
+data_split_pkl: data_split.pickle       # inside data folder .json or .pickle
+data_generator: DefaultDataLoader
+data_folder: Data/data_noblur_25_64x64x64    # inside project dir
+pet_normalization_constant: 32676
+augment: True
+data_shape: [64,64,64]
+color_channels_in: 1
+repeat_patient_list: 1
+
+# to implement
+# full_data_shape: [64,64,64]
+# loading_style: 'volume', 'patch'
+# input_data_shape: [8,8,8] for patch or [16,64,64] for slice
+# for plotting during training
+callback_image2image: ImagePredictionLogger
+
+# model-specific info (self-generated) - do not write anything beyond here

rhtorch/torch_training.py

@@ -12,7 +12,7 @@
 # library package imports
 from rhtorch.models import modules
 from rhtorch.callbacks import plotting
-from rhtorch.config_utils import load_model_config, copy_model_config
+from rhtorch.config_utils import UserConfig
 
 def main():
     import argparse
@@ -30,8 +30,9 @@ def main():
     is_test = args.test
 
     # load configs from file + additional info from args
-    configs = load_model_config(project_dir, args)
-
+    user_configs = UserConfig(project_dir, args)
+    configs = user_configs.hparams   ### WARNING TO CHECK IF THIS is 2 names for the same memory address or 2 distinct memory addresses (matters when saving copy in the end)
+    
     # Set local data_generator
     sys.path.insert(1, args.input)
     import data_generator
@@ -61,19 +62,21 @@ def main():
     model = module(configs, shape_in)
 
     # transfer learning setup
-    if 'pretrained_generator' in configs:
+    if configs['pretrained_generator']:
         print("Setting up transfer learning")
         pretrained_model_path = Path(configs['pretrained_generator'])
         if pretrained_model_path.exists():
             if pretrained_model_path.name.endswith(".ckpt"):
                 # important to pass in new configs here as we want to load the weights but config may differ from pretrained model
                 model = module.load_from_checkpoint(pretrained_model_path, hparams=configs, in_shape=shape_in, strict=False)
-            elif pretrained_model_path.endswith(".pt"):
+            elif pretrained_model_path.name.endswith(".pt"):
                 # this works for both .pt and .ckpt actually
                 # WARNING I don't know which of the above or below method is the correct way to load ckpt
                 # this below method only load the weights. Above also load state of optimizer, etc...
                 ckpt = torch.load(pretrained_model_path)
-                pretrained_model = ckpt['state_dict']
+                # OBS, the 'state_dict' is not set during save? 
+                # What if we are to save multiple models used later for pretrain? (e.g. a GAN with 3 networks?)
+                pretrained_model = ckpt['state_dict'] if 'state_dict' in ckpt.keys() else ckpt
                 model.load_state_dict(pretrained_model, strict=False)
             else:
                 raise ValueError("Expected model format: '.pt' or '.ckpt'.")
@@ -117,15 +120,16 @@ def main():
 
     # Save the config prior to training the model - one for each time the script is started
     if not is_test:        
-        copy_model_config(model_path, configs, append_timestamp=True)
+        user_configs.save_copy(model_path, append_timestamp=True)
         print("Saved config prior to model training")
 
     # set the trainer and fit
+    accelerator = 'ddp' if configs['GPUs'] > 1 else None
     trainer = pl.Trainer(max_epochs=configs['epoch'], 
                          logger=wandb_logger, 
                          callbacks=callbacks, 
                          gpus=-1, 
-                         accelerator='ddp',
+                         accelerator=accelerator,
                          resume_from_checkpoint=existing_checkpoint,
                          auto_select_gpus=True,
                          accumulate_grad_batches=configs['acc_grad_batches'],
@@ -136,12 +140,12 @@ def main():
     trainer.fit(model, train_dataloader, valid_dataloader)
 
     # add useful info to saved configs
-    configs['best_model'] = checkpoint_callback.best_model_path
+    user_configs.hparams['best_model'] = checkpoint_callback.best_model_path
 
     # save the model
     output_file = model_path.joinpath(f"{configs['model_name']}.pt")
     torch.save(model.state_dict(), output_file)
-    copy_model_config(model_path, configs)
+    user_configs.save_copy(model_path)
     print("Saved model and config file to disk")
 
 

rhtorch/version.py

@@ -13,6 +13,5 @@
 VERSIONING (UPDATED WHEN PR ARE MERGED INTO MASTER BRANCH)
  0.0.1 # Added repository (CL 18-05-2021)
  0.0.2 # Cleaned up main, moved to torchmetrics in modules (RD 20-05-2021)
- 0.0.3 # Added version control to config-logfiles
-
-"""
+ 0.0.3 # Added version control to config-logfiles, and default config yaml settings (CL, RD 23-05-2021)
+"""