Skip to content
Snippets Groups Projects
Select Git revision
  • db8197a25f72cebe5cb01e27d7dcc98c286b4830
  • master default protected
2 results

index.html

Blame
  • Code owners
    Assign users and groups as approvers for specific file changes. Learn more.
    ANN_Training.py 7.22 KiB
    # -*- coding: utf-8 -*-
    """
    @author: Laura C. Kühle, Soraya Terrab (sorayaterrab)
    
    TODO: Give option to compare multiple models
    TODO: Add more evaluation measures (AUROC, ROC, F1, training accuracy, boxplot over CVF, etc.)
    TODO: Add log to pipeline
    TODO: Remove object set-up
    TODO: Optimize Snakefile-vs-config relation
    TODO: Improve maximum selection runtime
    
    """
    import numpy as np
    import matplotlib.pyplot as plt
    import os
    import torch
    from torch.utils.data import TensorDataset, DataLoader, random_split
    from sklearn.model_selection import KFold
    # from sklearn.metrics import accuracy_score
    from sklearn.metrics import accuracy_score, precision_recall_fscore_support, precision_score
    
    import ANN_Model
    from Plotting import plot_classification_accuracy
    
    
    class ModelTrainer(object):
        def __init__(self, config):
            self._reset(config)
    
        def _reset(self, config):
            self._dir = config.pop('dir', 'test_data')
            self._model_name = config.pop('model_name', '0')
            self._read_training_data()
    
            self._batch_size = config.pop('batch_size', min(len(self._training_data)//2, 500))
            self._num_epochs = config.pop('num_epochs', 1000)
            self._threshold = config.pop('threshold', 1e-5)
            self._model = config.pop('model', 'ThreeLayerReLu')
            self._model_config = config.pop('model_config', {})
            self._loss_function = config.pop('loss_function', 'BCELoss')
            self._loss_config = config.pop('loss_config', {})
            self._optimizer = config.pop('optimizer', 'Adam')
            self._optimizer_config = config.pop('optimizer_config', {})
    
            # Set learning rate
            self._learning_rate = config.pop('learning_rate', 1e-2)
            self._optimizer_config['lr'] = self._learning_rate
    
            if not hasattr(ANN_Model, self._model):
                raise ValueError('Invalid model: "%s"' % self._model)
            if not hasattr(torch.nn.modules.loss, self._loss_function):
                raise ValueError('Invalid loss function: "%s"' % self._loss_function)
            if not hasattr(torch.optim, self._optimizer):
                raise ValueError('Invalid optimizer: "%s"' % self._optimizer)
    
            self._model = getattr(ANN_Model, self._model)(self._model_config)
            self._loss_function = getattr(torch.nn.modules.loss, self._loss_function)(
                **self._loss_config)
            self._optimizer = getattr(torch.optim, self._optimizer)(
                self._model.parameters(), **self._optimizer_config)
            self._validation_loss = torch.zeros(self._num_epochs//10)
    
        def _read_training_data(self):
            # Get training dataset from saved file and map to Torch tensor and dataset
            input_file = self._dir + '/input_data.npy'
            output_file = self._dir + '/output_data.npy'
            self._training_data = TensorDataset(*map(torch.tensor, (np.load(input_file),
                                                                    np.load(output_file))))
    
        def epoch_training(self, dataset=None, num_epochs=None):
            # Split data into training and validation set
            if dataset is None:
                dataset = self._training_data
            if num_epochs is None:
                num_epochs = self._num_epochs
            num_samples = len(dataset)
            train_ds, valid_ds = random_split(dataset, [round(num_samples*0.8), round(num_samples*0.2)])
    
            # Load sets
            train_dl = DataLoader(train_ds, batch_size=self._batch_size, shuffle=True)
            valid_dl = DataLoader(valid_ds, batch_size=self._batch_size * 2)
    
            # Training with Validation
            for epoch in range(num_epochs):
                self._model.train()
                for x_batch, y_batch in train_dl:
                    pred = self._model(x_batch.float())
                    loss = self._loss_function(pred, y_batch.float()).mean()
    
                    # Run back propagation, update the weights, and zero gradients for next epoch
                    loss.backward()
                    self._optimizer.step()
                    self._optimizer.zero_grad()
    
                self._model.eval()
                with torch.no_grad():
                    valid_loss = sum(
                        self._loss_function(self._model(x_batch_valid.float()), y_batch_valid.float())
                        for x_batch_valid, y_batch_valid in valid_dl)
    
                    if (epoch+1) % 100 == 0:
                        self._validation_loss[int((epoch+1) / 100)-1] = valid_loss / len(valid_dl)
                        print(epoch+1, valid_loss / len(valid_dl))
    
                    if valid_loss / len(valid_dl) < self._threshold:
                        break
    
        def test_model(self, num_iterations=100):
            classification_stats = []
            for iteration in range(num_iterations):
                dataset = self._training_data
                for train_index, test_index in KFold(n_splits=5, shuffle=True).split(dataset):
                    # print("TRAIN:", train_index, "TEST:", test_index)
                    training_set = TensorDataset(*dataset[train_index])
                    test_set = dataset[test_index]
    
                    classification_stats.append(self._test_fold(training_set, test_set))
    
            # print(classification_stats)
            # print(np.array(classification_stats).mean(axis=0))
            classification_stats = np.array(classification_stats).mean(axis=0)
    
            plot_classification_accuracy([self._model_name], *classification_stats)
    
            # Set paths for plot files if not existing already
            plot_dir = self._dir + '/model evaluation'
            if not os.path.exists(plot_dir):
                os.makedirs(plot_dir)
    
            # Save plots
            for identifier in plt.get_figlabels():
                # Set path for figure directory if not existing already
                if not os.path.exists(plot_dir + '/' + identifier):
                    os.makedirs(plot_dir + '/' + identifier)
    
                plt.figure(identifier)
                plt.savefig(plot_dir + '/' + identifier + '/' + self._model_name + '.pdf')
    
        def _test_fold(self, training_set, test_set):
            self.epoch_training(training_set, num_epochs=100)
            self._model.eval()
    
            x_test, y_test = test_set
            # print(self._model(x_test.float()))
            model_output = torch.tensor([[1.0, 0.0] if value == 0 else [0.0, 1.0]
                                         for value in torch.max(self._model(x_test.float()), 1)[1]])
    
            y_true = y_test.detach().numpy()
            y_pred = model_output.detach().numpy()
            accuracy = accuracy_score(y_true, y_pred)
            # print('sklearn', accuracy)
            precision, recall, f_score, support = precision_recall_fscore_support(y_true, y_pred)
            # print(precision, recall)
            # print()
    
            return [precision[0], recall[0], accuracy]
    
        def save_model(self):
            # Saving Model
            name = self._model_name
    
            # Set paths for plot files if not existing already
            model_dir = self._dir + '/trained models'
            if not os.path.exists(model_dir):
                os.makedirs(model_dir)
    
            torch.save(self._model.state_dict(), model_dir + '/model__' + name + '.pt')
            torch.save(self._validation_loss, model_dir + '/loss__' + name + '.pt')
    
        def _classify(self):
            pass
    
    
    # Loss Functions: BCELoss, BCEWithLogitsLoss,
    # CrossEntropyLoss (not working), MSELoss (with reduction='sum')
    # Optimizer: Adam, SGD
    # trainer = ModelTrainer({'num_epochs': 1000})
    # trainer.epoch_training()
    # trainer.test_model()
    # trainer.save_model()