run.py

"""Brain Mapper Run File

Description:

    This file contains all the relevant functions for running BrainMapper.
    The network can be ran in one of these modes:
        - train
        - evaluate path
        - evaluate whole

    TODO: Might be worth adding some information on uncertaintiy estimation, later down the line

Usage:

    In order to run the network, in the terminal, the user needs to pass it relevant arguments:

        $ ./setup.sh
        $ source env/bin/activate
        $ python run.py --mode ...

    The arguments for mode are the following:

        mode=train # For training the model
        mode=evaluate-score # For evaluating the model score
        mode=evaluate-mapping # For evaluating the model mapping
        # For clearning the experiments and logs directories of the last experiment
        mode=clear-experiment
        mode=clear-all # For clearing all the files from the experiments and logs directories/

"""

import os
import shutil
import argparse
import logging
from settings import Settings

import torch
import torch.utils.data as data
import numpy as np

from solver import Solver
# from BrainMapperUNet import BrainMapperUNet3D, BrainMapperResUNet3D, BrainMapperResUNet3Dshallow, BrainMapperCompResUNet3D
from BrainMapperAE import BrainMapperAE3D
from utils.data_utils import get_datasets, data_test_train_validation_split, update_shuffling_flag, create_folder
import utils.data_evaluation_utils as evaluations
from utils.data_logging_utils import LogWriter

# Set the default floating point tensor type to FloatTensor

torch.set_default_tensor_type(torch.FloatTensor)


def load_data(data_parameters):
    """Dataset Loader

    This function loads the training and validation datasets.

    Args:
        data_parameters (dict): Dictionary containing relevant information for the datafiles.

    Returns:
        train_data (dataset object): Pytorch map-style dataset object, mapping indices to training data samples.
        validation_data (dataset object): Pytorch map-style dataset object, mapping indices to testing data samples.

    """
    print("Data is loading...")
    train_data, validation_data = get_datasets(data_parameters)
    print("Data has loaded!")
    print("Training dataset size is {}".format(len(train_data)))
    print("Validation dataset size is {}".format(len(validation_data)))

    return train_data, validation_data


def train(data_parameters, training_parameters, network_parameters, misc_parameters):
    """Training Function

    This function trains a given model using the provided training data.
    Currently, the data loaded is set to have multiple sub-processes.
    A high enough number of workers assures that CPU computations are efficiently managed, i.e. that the bottleneck is indeed the neural network's forward and backward operations on the GPU (and not data generation)
    Loader memory is also pinned, to speed up data transfer from CPU to GPU  by using the page-locked memory.
    Train data is also re-shuffled at each training epoch.

    Args:
        data_parameters (dict): Dictionary containing relevant information for the datafiles.

        training_parameters(dict): Dictionary containing relevant hyperparameters for training the network.
        training_parameters = {
            'training_batch_size': 5
            'validation_batch_size: 5
            'use_pre_trained': False
            'pre_trained_path': 'pre_trained/path'
            'experiment_name': 'experiment_name'
            'learning_rate': 1e-4
            'optimizer_beta': (0.9, 0.999)
            'optimizer_epsilon': 1e-8
            'optimizer_weigth_decay': 1e-5
            'number_of_epochs': 10
            'loss_log_period': 50
            'learning_rate_scheduler_step_size': 3
            'learning_rate_scheduler_gamma': 1e-1
            'use_last_checkpoint': True
            'final_model_output_file': 'path/to/model'
        }

        network_parameters (dict): Contains information relevant parameters

        misc_parameters (dict): Dictionary of aditional hyperparameters
        misc_parameters = {
            'save_model_directory': 'directory_name'
            'model_name': 'BrainMapper'
            'logs_directory': 'log-directory'
            'device': 1
            'experiments_directory': 'experiments-directory'
        }
    """

    def _train_runner(data_parameters, training_parameters, network_parameters, misc_parameters):
        """Wrapper for the training operation

        This function wraps the training operation for the network

        Args:
            data_parameters (dict): Dictionary containing relevant information for the datafiles.
            training_parameters(dict): Dictionary containing relevant hyperparameters for training the network.
            network_parameters (dict): Contains information relevant parameters
            misc_parameters (dict): Dictionary of aditional hyperparameters

        """
        train_data, validation_data = load_data(data_parameters)

        train_loader = data.DataLoader(
            dataset=train_data,
            batch_size=training_parameters['training_batch_size'],
            shuffle=True,
            num_workers=4,
            pin_memory=True
        )

        validation_loader = data.DataLoader(
            dataset=validation_data,
            batch_size=training_parameters['validation_batch_size'],
            shuffle=False,
            num_workers=4,
            pin_memory=True
        )

        if training_parameters['use_pre_trained']:
            BrainMapperModel = torch.load(
                training_parameters['pre_trained_path'])
        else:
            # BrainMapperModel = BrainMapperUNet3D(network_parameters)
            # BrainMapperModel = BrainMapperResUNet3D(network_parameters)
            # BrainMapperModel = BrainMapperResUNet3Dshallow(network_parameters)
            # BrainMapperModel = BrainMapperCompResUNet3D(network_parameters)
            
            BrainMapperModel = BrainMapperAE3D(network_parameters)

        BrainMapperModel.reset_parameters()

        optimizer = torch.optim.Adam

        solver = Solver(model=BrainMapperModel,
                        device=misc_parameters['device'],
                        number_of_classes=network_parameters['number_of_classes'],
                        experiment_name=training_parameters['experiment_name'],
                        optimizer=optimizer,
                        optimizer_arguments={'lr': training_parameters['learning_rate'],
                                             'betas': training_parameters['optimizer_beta'],
                                             'eps': training_parameters['optimizer_epsilon'],
                                             'weight_decay': training_parameters['optimizer_weigth_decay']
                                             },
                        model_name=training_parameters['experiment_name'],
                        number_epochs=training_parameters['number_of_epochs'],
                        loss_log_period=training_parameters['loss_log_period'],
                        learning_rate_scheduler_step_size=training_parameters[
                            'learning_rate_scheduler_step_size'],
                        learning_rate_scheduler_gamma=training_parameters['learning_rate_scheduler_gamma'],
                        use_last_checkpoint=training_parameters['use_last_checkpoint'],
                        experiment_directory=misc_parameters['experiments_directory'],
                        logs_directory=misc_parameters['logs_directory'],
                        checkpoint_directory=misc_parameters['checkpoint_directory'],
                        save_model_directory=misc_parameters['save_model_directory'],
                        final_model_output_file=training_parameters['final_model_output_file']
                        )

        validation_loss = solver.train(train_loader, validation_loader)

        del train_data, validation_data, train_loader, validation_loader, BrainMapperModel, solver, optimizer
        torch.cuda.empty_cache()

        return validation_loss

    if data_parameters['k_fold'] is None:

        _ = _train_runner(data_parameters, training_parameters,
                          network_parameters, misc_parameters)

    else:
        print("Training initiated using K-fold Cross Validation!")
        k_fold_losses = []

        for k in range(data_parameters['k_fold']):

            print("K-fold Number: {}".format(k+1))

            data_parameters['train_list'] = os.path.join(
                data_parameters['data_folder_name'], 'train' + str(k+1)+'.txt')
            data_parameters['validation_list'] = os.path.join(
                data_parameters['data_folder_name'], 'validation' + str(k+1)+'.txt')
            training_parameters['final_model_output_file'] = training_parameters['final_model_output_file'].replace(
                ".pth.tar", str(k+1)+".pth.tar")

            validation_loss = _train_runner(
                data_parameters, training_parameters, network_parameters, misc_parameters)

            k_fold_losses.append(validation_loss)

        for k in range(data_parameters['k_fold']):
            print("K-fold Number: {} Loss: {}".format(k+1, k_fold_losses[k]))
        print("K-fold Cross Validation Avearge Loss: {}".format(np.mean(k_fold_losses)))


def evaluate_score(training_parameters, network_parameters, misc_parameters, evaluation_parameters):
    """Mapping Score Evaluator

    This function evaluates a given trained model by calculating the it's dice score prediction.

    Args:

        training_parameters(dict): Dictionary containing relevant hyperparameters for training the network.
        training_parameters = {
            'experiment_name': 'experiment_name'
        }

        network_parameters (dict): Contains information relevant parameters
        network_parameters= {
            'number_of_classes': 1
        }

        misc_parameters (dict): Dictionary of aditional hyperparameters
        misc_parameters = {
            'logs_directory': 'log-directory'
            'device': 1
            'experiments_directory': 'experiments-directory'
        }

        evaluation_parameters (dict): Dictionary of parameters useful during evaluation.
        evaluation_parameters = {
            'trained_model_path': 'path/to/model'
            'data_directory': 'path/to/data'
            'targets_directory': 'path/to/targets'
            'data_list': 'path/to/datalist.txt/
            'orientation': 'coronal'
            'saved_predictions_directory': 'directory-of-saved-predictions'
        }
    """

    # TODO - NEED TO UPDATE THE DATA FUNCTIONS!

    logWriter = LogWriter(number_of_classes=network_parameters['number_of_classes'],
                          logs_directory=misc_parameters['logs_directory'],
                          experiment_name=training_parameters['experiment_name']
                          )

    prediction_output_path = os.path.join(misc_parameters['experiments_directory'],
                                          training_parameters['experiment_name'],
                                          evaluation_parameters['saved_predictions_directory']
                                          )

    _ = evaluations.evaluate_dice_score(trained_model_path=evaluation_parameters['trained_model_path'],
                                        number_of_classes=network_parameters['number_of_classes'],
                                        data_directory=evaluation_parameters['data_directory'],
                                        targets_directory=evaluation_parameters['targets_directory'],
                                        data_list=evaluation_parameters['data_list'],
                                        orientation=evaluation_parameters['orientation'],
                                        prediction_output_path=prediction_output_path,
                                        device=misc_parameters['device'],
                                        LogWriter=logWriter
                                        )

    logWriter.close()


def evaluate_mapping(mapping_evaluation_parameters):
    """Mapping Evaluator

    This function passes through the network an input and generates the rsfMRI outputs.

    Args:
        mapping_evaluation_parameters (dict): Dictionary of parameters useful during mapping evaluation.
        mapping_evaluation_parameters = {
            'trained_model_path': 'path/to/model'
            'data_directory': 'path/to/data'
            'data_list': 'path/to/datalist.txt/
            'prediction_output_path': 'directory-of-saved-predictions'
            'batch_size': 2
            'device': 0
            'exit_on_error': True
        }

    """
    trained_model_path = mapping_evaluation_parameters['trained_model_path']
    data_directory = mapping_evaluation_parameters['data_directory']
    mapping_data_file = mapping_evaluation_parameters['mapping_data_file']
    data_list = mapping_evaluation_parameters['data_list']
    prediction_output_path = mapping_evaluation_parameters['prediction_output_path']
    device = mapping_evaluation_parameters['device']
    exit_on_error = mapping_evaluation_parameters['exit_on_error']
    brain_mask_path = mapping_evaluation_parameters['brain_mask_path']
    mean_mask_path = mapping_evaluation_parameters['mean_mask_path']
    mean_reduction = mapping_evaluation_parameters['mean_reduction']
    scaling_factors = mapping_evaluation_parameters['scaling_factors']

    evaluations.evaluate_mapping(trained_model_path,
                                 data_directory,
                                 mapping_data_file,
                                 data_list,
                                 prediction_output_path,
                                 brain_mask_path,
                                 mean_mask_path,
                                 mean_reduction,
                                 scaling_factors,
                                 device=device,
                                 exit_on_error=exit_on_error)


def delete_files(folder):
    """ Clear Folder Contents

    Function which clears contents (like experiments or logs)

    Args:
        folder (str): Name of folders whose conents is to be deleted

    """

    for object_name in os.listdir(folder):
        file_path = os.path.join(folder, object_name)
        try:
            if os.path.isfile(file_path):
                os.unlink(file_path)
            elif os.path.isdir(file_path):
                shutil.rmtree(file_path)
        except Exception as exception:
            print(exception)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--mode', '-m', required=True,
                        help='run mode, valid values are train or evaluate')
    parser.add_argument('--settings_path', '-sp', required=False,
                        help='optional argument, set path to settings_evaluation.ini')

    arguments = parser.parse_args()

    settings = Settings('settings.ini')
    data_parameters = settings['DATA']
    training_parameters = settings['TRAINING']
    network_parameters = settings['NETWORK']
    misc_parameters = settings['MISC']
    evaluation_parameters = settings['EVALUATION']

    # Here we shuffle the data!

    if data_parameters['data_split_flag'] == True:

        print('Data is shuffling... This could take a few minutes!')

    if data_parameters['data_split_flag'] == True:
        if data_parameters['use_data_file'] == True:
            data_test_train_validation_split(data_parameters['data_folder_name'],
                                             data_parameters['test_percentage'],
                                             data_parameters['subject_number'],
                                             data_directory=data_parameters['data_directory'],
                                             train_inputs=data_parameters['train_data_file'],
                                             train_targets=data_parameters['train_output_targets'],
                                             mean_mask_path=data_parameters['mean_mask_path'],
                                             data_file=data_parameters['data_file'],
                                             K_fold=data_parameters['k_fold']
                                             )
        else:
            data_test_train_validation_split(data_parameters['data_folder_name'],
                                             data_parameters['test_percentage'],
                                             data_parameters['subject_number'],
                                             data_directory=data_parameters['data_directory'],
                                             train_inputs=data_parameters['train_data_file'],
                                             train_targets=data_parameters['train_output_targets'],
                                             mean_mask_path=data_parameters['mean_mask_path'],
                                             K_fold=data_parameters['k_fold']
                                             )
        update_shuffling_flag('settings.ini')

        print('Data is shuffling... Complete!')

    if arguments.mode == 'train':
        train(data_parameters, training_parameters,
              network_parameters, misc_parameters)

    # NOTE: THE EVAL FUNCTIONS HAVE NOT YET BEEN DEBUGGED (16/04/20)
    # NOTE: THE EVAL-MAPPING FUNCTION HAS BEEN DEBUGGED (28/04/20)

    elif arguments.mode == 'evaluate-score':
        evaluate_score(training_parameters,
                       network_parameters, misc_parameters, evaluation_parameters)
    elif arguments.mode == 'evaluate-mapping':
        logging.basicConfig(filename='evaluate-mapping-error.log')
        if arguments.settings_path is not None:
            settings_evaluation = Settings(arguments.settings_path)
        else:
            settings_evaluation = Settings('settings_evaluation.ini')
        mapping_evaluation_parameters = settings_evaluation['MAPPING']
        evaluate_mapping(mapping_evaluation_parameters)
    elif arguments.mode == 'clear-experiments':
        shutil.rmtree(os.path.join(
            misc_parameters['experiments_directory'], training_parameters['experiment_name']))
        shutil.rmtree(os.path.join(
            misc_parameters['logs_directory'], training_parameters['experiment_name']))
        print('Cleared the current experiments and logs directory successfully!')
    elif arguments.mode == 'clear-everything':
        delete_files(misc_parameters['experiments_directory'])
        delete_files(misc_parameters['logs_directory'])
        print('Cleared the current experiments and logs directory successfully!')
    elif arguments.mode == 'train-and-evaluate-mapping':
        if arguments.settings_path is not None:
            settings_evaluation = Settings(arguments.settings_path)
        else:
            settings_evaluation = Settings('settings_evaluation.ini')
        mapping_evaluation_parameters = settings_evaluation['MAPPING']
        train(data_parameters, training_parameters,
              network_parameters, misc_parameters)
        logging.basicConfig(filename='evaluate-mapping-error.log')
        evaluate_mapping(mapping_evaluation_parameters)
    else:
        raise ValueError(
            'Invalid mode value! Only supports: train, evaluate-score, evaluate-mapping, train-and-evaluate-mapping, clear-experiments and clear-everything')