Merge pull request #6 from CAAI/new_unet3d

New unet3d

Author: RaphaelSura

.gitignore

@@ -3,4 +3,5 @@ __pycache__
 wandb/
 rhtorch.egg-info
 .vscode/
-rhtorch/models/dev
+rhtorch/models/dev/*
+!rhtorch/models/dev/__init__.py

example_project/LowdosePETv1_config.yaml

@@ -1,6 +1,8 @@
-# main config to train models
+# main config to train your model
+# check default.config for all available options
+
 project_name: YOUR_PROJECT_NAME       # Used for WANDB
-version_name: RESUNETv1               # Make run unique by changing this counter
+version_name: UNETv1               # Make run unique by changing this counter
 
 # main model:
 precision: 32
@@ -10,29 +12,24 @@ acc_grad_batches: 4
 module: LightningAE
 
 # generator
-generator: Res3DUnet
+generator: UNet3D
+g_pooling_type: full_conv
+g_filters: [64, 128, 256, 512, 1024]
 # g_activation: ReLU
 g_optimizer: Adam
 g_lr: 1e-4
 #lr_scheduler: 'exponential_decay_0.01'
 g_loss: MeanAbsoluteError
-# pretrained_generator: absolute path to .pt or .ckpt
-
-# 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/data_5fold.pickle       # inside data folder .json or .pickle
 data_generator: CustomDataLoader
 data_folder: Data/data_noblur_25_64x64x64    # inside project dir
 pet_normalization_constant: 32676
 augment: True
-data_shape: [64,64,64]
+data_shape: [128, 128, 128]
 color_channels_in: 1
-repeat_patient_list: 500
+repeat_patient_list: 100
 
 # for plotting during training
 callback_image2image: 'ImagePredictionLogger'

rhtorch/callbacks/__init__.py

@@ -5,50 +5,51 @@
 from rhtorch.version import __version__
 import socket
 
+
 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')
+        default_config_file = Path(__file__).parent.joinpath('default.config')
         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.")
-        
+            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 key, value in self.default_params.items():
             # copy from default if value is not None/0/False and key not already in user config
             if value and key not in self.hparams:
                 self.hparams[key] = value
-            
 
     def fill_additional_info(self):
         # additional info from args and miscellaneous to save in config
@@ -58,21 +59,22 @@ def fill_additional_info(self):
         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['global_batch_size'] = self.hparams['batch_size'] * \
+            self.hparams['GPUs']
         self.hparams['rhtorch_version'] = __version__
         self.hparams['hostname'] = socket.gethostname()
-    
+
     def create_model_name(self):
-        
+
         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']}"
-    
+
     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}"        
+        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:

rhtorch/config_utils.py

@@ -10,10 +10,10 @@ acc_grad_batches: 1
 module: LightningAE
 
 # generator
-generator: UNet3DFullConv
-# depth: 4
-# initial_num_filters: 64
-g_activation: ReLU
+generator: UNet3D
+#g_pooling_type: full_conv                                       # full_conv or max_pool
+#g_filters: [64, 128, 256, 512, 1024]
+#g_activation: ReLU
 g_optimizer: Adam
 g_lr: 1e-4
 #lr_scheduler: 'exponential_decay_0.01'
@@ -33,9 +33,8 @@ d_loss: null
 data_split_pkl: data_split_file_inside_data_folder.pickle       # inside data folder .json or .pickle
 data_generator: DefaultDataLoader
 data_folder: data_folder_inside_project_dir                     # inside project dir
-pet_normalization_constant: 32676
 augment: True
-data_shape: [64,64,64]
+data_shape: [128, 128, 128]
 color_channels_in: 1
 repeat_patient_list: 1