Updated train scripts

Brain_study/MultiTrain_config.py

@@ -0,0 +1,78 @@
+import os, sys
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+
+import argparse
+from configparser import ConfigParser
+from datetime import datetime
+
+import DeepDeformationMapRegistration.utils.constants as C
+
+TRAIN_DATASET = '/mnt/EncryptedData1/Users/javier/ext_datasets/IXI_dataset/T1/training'
+
+err = list()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--ini', help='Configuration file')
+    args = parser.parse_args()
+
+    configFile = ConfigParser()
+    configFile.read(args.ini)
+    print('Loaded configuration file: ' + args.ini)
+    print({section: dict(configFile[section]) for section in configFile.sections()})
+    print('\n\n')
+
+    trainConfig = configFile['TRAIN']
+    lossesConfig = configFile['LOSSES']
+    datasetConfig = configFile['DATASETS']
+    othersConfig = configFile['OTHERS']
+    augmentationConfig = configFile['AUGMENTATION']
+
+    simil = lossesConfig['similarity'].split(',')
+    segm = lossesConfig['segmentation'].split(',')
+    if trainConfig['name'].lower() == 'uw':
+        from Brain_study.Train_UncertaintyWeighted import launch_train
+        loss_config = {'simil': simil, 'segm': segm}
+    elif trainConfig['name'].lower() == 'segguided':
+        from Brain_study.Train_SegmentationGuided import launch_train
+        loss_config = {'simil': simil[0], 'segm': segm[0]}
+    else:
+        from Brain_study.Train_Baseline import launch_train
+        loss_config = {'simil': simil[0]}
+
+    output_folder = os.path.join(othersConfig['outputFolder'],
+                                 '{}_Lsim_{}__Lseg_{}'.format(trainConfig['name'], '_'.join(simil), '_'.join(segm)))
+    output_folder = output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y")
+
+    print('TRAIN ' + datasetConfig['train'])
+
+    if augmentationConfig:
+        C.GAMMA_AUGMENTATION = augmentationConfig['gamma'].lower() == 'true'
+        C.BRIGHTNESS_AUGMENTATION = augmentationConfig['brightness'].lower() == 'true'
+
+    try:
+        unet = [int(x) for x in trainConfig['unet'].split(',')]
+    except KeyError as e:
+        unet = [16, 32, 64, 128, 256]
+
+    try:
+        head = [int(x) for x in trainConfig['head'].split(',')]
+    except KeyError as e:
+        head = [16, 16]
+
+    launch_train(dataset_folder=datasetConfig['train'],
+                 validation_folder=datasetConfig['validation'],
+                 output_folder=output_folder,
+                 gpu_num=eval(trainConfig['gpu']),
+                 lr=eval(trainConfig['learningRate']),
+                 rw=eval(trainConfig['regularizationWeight']),
+                 acc_gradients=eval(trainConfig['accumulativeGradients']),
+                 batch_size=eval(trainConfig['batchSize']),
+                 max_epochs=eval(trainConfig['epochs']),
+                 image_size=eval(trainConfig['imageSize']),
+                 early_stop_patience=eval(trainConfig['earlyStopPatience']),
+                 unet=unet,
+                 head=head,
+                 **loss_config)

Brain_study/Train_Baseline.py

@@ -1,4 +1,6 @@
 import os, sys
+import warnings
+
 currentdir = os.path.dirname(os.path.realpath(__file__))
 parentdir = os.path.dirname(currentdir)
 sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
@@ -32,25 +34,30 @@ from tqdm import tqdm
 from datetime import datetime
 
 
-def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim'):
+def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim',
+                 acc_gradients=16, batch_size=1, max_epochs=10000, early_stop_patience=1000, image_size=64,
+                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
     os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_num) # Check availability before running using 'nvidia-smi'
     C.GPU_NUM = str(gpu_num)
 
+    if batch_size != 1 and acc_gradients != 1:
+        warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
+
     output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
     os.makedirs(output_folder, exist_ok=True)
     # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
     log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
     C.TRAINING_DATASET = dataset_folder #dataset_copy.copy_dataset()
     C.VALIDATION_DATASET = validation_folder
-    C.ACCUM_GRADIENT_STEP = 16
-    C.BATCH_SIZE = 16 if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
-    C.EARLY_STOP_PATIENCE = 5 * (C.ACCUM_GRADIENT_STEP / 2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.ACCUM_GRADIENT_STEP = acc_gradients
+    C.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else 1
+    C.EARLY_STOP_PATIENCE = early_stop_patience
     C.LEARNING_RATE = lr
     C.LIMIT_NUM_SAMPLES = None
-    C.EPOCHS = 10000
+    C.EPOCHS = max_epochs
 
     aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
           "GPU: {}\n" \
@@ -84,7 +91,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     validation_generator = data_generator.get_validation_generator()
 
     image_input_shape = train_generator.get_data_shape()[-1][:-1]
-    image_output_shape = [64] * 3
+    image_output_shape = [image_size] * 3
 
     # Config the training sessions
     config = tf.compat.v1.ConfigProto()  # device_count={'GPU':0})
@@ -125,13 +132,15 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     augm_model_valid = Model(inputs=input_layer_valid, outputs=augm_layer_valid(input_layer_valid))
 
     # Build model
-    enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
-    dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    # enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+    # dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    enc_features = unet     # const.ENCODER_FILTERS
+    dec_features = enc_features[::-1] + head   # const.ENCODER_FILTERS[::-1]
     nb_features = [enc_features, dec_features]
     network = vxm.networks.VxmDense(inshape=image_output_shape,
                                     nb_unet_features=nb_features,
                                     int_steps=0)
-
+    network.summary(line_length=150)
     # Losses and loss weights
     SSIM_KER_SIZE = 5
     MS_SSIM_WEIGHTS = _MSSSIM_WEIGHTS[:3]

Brain_study/Train_SegmentationGuided.py

@@ -29,25 +29,32 @@ from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
 
 import time
+import warnings
 
-def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim', segm='hd'):
+
+def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr=1e-4, rw=5e-3, simil='ssim', segm='hd',
+                 acc_gradients=16, batch_size=1, max_epochs=10000, early_stop_patience=1000, image_size=64,
+                 unet=[16, 32, 64, 128, 256], head=[16, 16]):
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
     os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_num)  # Check availability before running using 'nvidia-smi'
     C.GPU_NUM = str(gpu_num)
 
+    if batch_size != 1 and acc_gradients != 1:
+        warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
+
     output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
     os.makedirs(output_folder, exist_ok=True)
     log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
     C.TRAINING_DATASET = dataset_folder
     C.VALIDATION_DATASET = validation_folder
-    C.ACCUM_GRADIENT_STEP = 16
-    C.BATCH_SIZE = 2 if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
-    C.EARLY_STOP_PATIENCE = 10 * (C.ACCUM_GRADIENT_STEP / 2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.ACCUM_GRADIENT_STEP = acc_gradients
+    C.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else 1
+    C.EARLY_STOP_PATIENCE = early_stop_patience
     C.LEARNING_RATE = lr
     C.LIMIT_NUM_SAMPLES = None
-    C.EPOCHS = 10000
+    C.EPOCHS = max_epochs
 
     aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
           "GPU: {}\n" \
@@ -81,7 +88,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
     validation_generator = data_generator.get_validation_generator()
 
     image_input_shape = train_generator.get_data_shape()[1][:-1]
-    image_output_shape = [64] * 3
+    image_output_shape = [image_size] * 3
 
     nb_labels = len(train_generator.get_segmentation_labels())
 
@@ -109,8 +116,10 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                                          trainable=False)
     augm_model = Model(inputs=input_layer_augm, outputs=augm_layer(input_layer_augm))
 
-    enc_features = [16, 32, 32, 32]# const.ENCODER_FILTERS
-    dec_features = [32, 32, 32, 32, 32, 16, 16]# const.ENCODER_FILTERS[::-1]
+    # enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+    # dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    enc_features = unet     # const.ENCODER_FILTERS
+    dec_features = enc_features[::-1] + head   # const.ENCODER_FILTERS[::-1]
     nb_features = [enc_features, dec_features]
 
     network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
@@ -138,7 +147,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
         @function_decorator('MS_SSIM_MSE__loss')
         def loss_simil(y_true, y_pred):
             return MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).loss(y_true, y_pred) + vxm.losses.MSE().loss(y_true, y_pred)
-    elif simil=='ssim__ncc' or simil=='ncc__ssim' :
+    elif simil=='ssim__ncc' or simil=='ncc__ssim':
         @function_decorator('SSIM_NCC__loss')
         def loss_simil(y_true, y_pred):
             return StructuralSimilarity_simplified(patch_size=SSIM_KER_SIZE, dim=3, dynamic_range=1.).loss(y_true, y_pred) + NCC(image_input_shape).loss(y_true, y_pred)
@@ -153,6 +162,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
         loss_segm = HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).loss
     elif segm == 'dice':
         loss_segm = GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss
+    elif segm == 'dice_macro':
+        loss_segm = GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss_macro
     else:
         raise ValueError('No valid value for segm')
 
@@ -163,8 +174,8 @@ def launch_train(dataset_folder, validation_folder, output_folder, gpu_num=0, lr
                     'seg_transformer': 1.,
                     'flow': 5e-3}
     metrics = {'transformer': [vxm.losses.MSE().loss, NCC(image_input_shape).metric, MultiScaleStructuralSimilarity(max_val=1., filter_size=SSIM_KER_SIZE, power_factors=MS_SSIM_WEIGHTS).metric],
-               'seg_transformer': [GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric,
-                                     HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [train_generator.get_data_shape()[2][-1]]).metric
+               'seg_transformer': [GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).metric_macro,
+                                     #HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [train_generator.get_data_shape()[2][-1]]).metric
                                      ]}
     metrics_weights = {'transformer': 1,
                        'seg_transformer': 1,

Brain_study/Train_UncertaintyWeighted.py

@@ -26,28 +26,33 @@ from DeepDeformationMapRegistration.utils.acummulated_optimizer import AdamAccum
 
 from Brain_study.data_generator import BatchGenerator
 from Brain_study.utils import SummaryDictionary, named_logs
+import warnings
 
 
-def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5e-3, lr=1e-4,
-                 gpu_num=0, simil=['mse'], segm=['dice']):
+def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5e-3, lr=1e-4, rw=5e-3,
+                 gpu_num=0, simil=['mse'], segm=['dice'], acc_gradients=16, batch_size=1, max_epochs=10000,
+                 early_stop_patience=1000, image_size=64, unet=[16, 32, 64, 128, 256], head=[16, 16]):
     assert dataset_folder is not None and output_folder is not None
 
     os.environ['CUDA_DEVICE_ORDER'] = C.DEV_ORDER
     os.environ['CUDA_VISIBLE_DEVICES'] = str(gpu_num) # Check availability before running using 'nvidia-smi'
     C.GPU_NUM = str(gpu_num)
 
+    if batch_size != 1 and acc_gradients != 1:
+        warnings.warn('WARNING: Batch size and Accumulative gradient step are set!')
+
     output_folder = os.path.join(output_folder + '_' + datetime.now().strftime("%H%M%S-%d%m%Y"))
     # dataset_copy = DatasetCopy(dataset_folder, os.path.join(output_folder, 'temp'))
     os.makedirs(output_folder, exist_ok=True)
     log_file = open(os.path.join(output_folder, 'log.txt'), 'w')
     C.TRAINING_DATASET = dataset_folder #dataset_copy.copy_dataset()
     C.VALIDATION_DATASET = validation_folder
-    C.ACCUM_GRADIENT_STEP = 16
-    C.BATCH_SIZE = 2 if C.ACCUM_GRADIENT_STEP == 1 else C.ACCUM_GRADIENT_STEP
-    C.EARLY_STOP_PATIENCE = 10 * (C.ACCUM_GRADIENT_STEP/2 if C.ACCUM_GRADIENT_STEP != 1 else 1)
+    C.ACCUM_GRADIENT_STEP = acc_gradients
+    C.BATCH_SIZE = batch_size if C.ACCUM_GRADIENT_STEP == 1 else 1
+    C.EARLY_STOP_PATIENCE = early_stop_patience
     C.LEARNING_RATE = lr
     C.LIMIT_NUM_SAMPLES = None
-    C.EPOCHS = 10000
+    C.EPOCHS = max_epochs
 
     aux = "[{}]\tINFO:\nTRAIN DATASET: {}\nVALIDATION DATASET: {}\n" \
           "GPU: {}\n" \
@@ -81,7 +86,7 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
     validation_generator = data_generator.get_validation_generator()
 
     image_input_shape = train_generator.get_data_shape()[-1][:-1]
-    image_output_shape = [64] * 3
+    image_output_shape = [image_size] * 3
 
     nb_labels = len(train_generator.get_segmentation_labels())
 
@@ -119,13 +124,16 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
     loss_segm = []
     for s in segm:
         if s=='dice':
-            loss_segm.append(GeneralizedDICEScore(image_output_shape + [train_generator.get_data_shape()[2][-1]], num_labels=nb_labels).loss)
+            loss_segm.append(GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss)
             prior_loss_w.append(1.)
         elif s=='hd':
-            loss_segm.append(HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [train_generator.get_data_shape()[2][-1]]).loss)
+            loss_segm.append(HausdorffDistanceErosion(3, 10, im_shape=image_output_shape + [nb_labels]).loss)
+            prior_loss_w.append(1.)
+        elif s == 'dice_macro':
+            loss_segm.append(GeneralizedDICEScore(image_output_shape + [nb_labels], num_labels=nb_labels).loss_macro)
             prior_loss_w.append(1.)
         else:
-            raise ValueError('Unknown similarity function: ', s)
+            raise ValueError('Unknown similarity function: ' + s)
 
     # Build augmentation layer model
     input_layer_augm = Input(shape=train_generator.get_data_shape()[0], name='input_augmentation')
@@ -142,8 +150,10 @@ def launch_train(dataset_folder, validation_folder, output_folder, prior_reg_w=5
                                    trainable=False)
     augmentation_model = Model(inputs=input_layer_augm, outputs=augm_layer(input_layer_augm))
 
-    enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
-    dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    # enc_features = [16, 32, 32, 32]     # const.ENCODER_FILTERS
+    # dec_features = [32, 32, 32, 32, 32, 16, 16]     # const.ENCODER_FILTERS[::-1]
+    enc_features = unet     # const.ENCODER_FILTERS
+    dec_features = enc_features[::-1] + head   # const.ENCODER_FILTERS[::-1]
     nb_features = [enc_features, dec_features]
     network = vxm.networks.VxmDenseSemiSupervisedSeg(inshape=image_output_shape,
                                                      nb_labels=nb_labels,