trainer.py

"""Copyright © 2020-present, Swisscom (Schweiz) AG.
All rights reserved.

This class is used for training models and is the core of the framework.

With the help of this class, the user of the framework is able to train and
develop models. The framework gets all the relevant objects as an input, and
all the parameters from a YAML file or a dictionary with the parameters, it
instantiates all the relevant helper objects for training the model and does
the training.
"""

from copsolver.frank_wolfe_solver import FrankWolfeSolver
from copsolver.analytical_solver import AnalyticalSolver
from validator import Validator
from paretomanager.pareto_manager_class import ParetoManager
from commondescentvector.multi_objective_cdv import MultiObjectiveCDV
from metric.metric_at_k import MetricAtK
from loss.loss_class import Loss
import torch.nn as nn
from dataloader.mamo_data_handler import MamoDataHandler
import time
import numpy as np
import os
import torch
import torch.optim as optim
from torch.autograd import Variable
import yaml
import logging
logger = logging.getLogger('main')
logger.setLevel(logging.INFO)


class Trainer():
    """The trainer class, the core of the framework, used for training models.

    All the needed objects for this class have to be given through the constructor.
    Additionally, the other parameters needed by this trainer have to be supplied
    in a YAML file named 'trainer_params.yaml' or a dictionary containing the parameters.
    For more details about the parameters supplied in this YAML file, please refer to
    'Attributes from the YAML file' section below.

    Attributes:
        data_handler: A MamoDataHandler object which feeds the data set to the trainer.
        model: A torch.nn.Module object which is the model that is being trained.
        losses: A list of Loss objects which represent the losses/objectives that the
            model is trained on.
        validation_metrics: A list of MetricAtK objects which are used to evaluate
            the model while the training and validation process.
        save_to_path: A path to a directory where the trained models from the Pareto
            front will be saved during training.
        device: A variable indicating whether the model is trained on the gpu or on
            the cpu.
        _train_dataloader: A dataloader object used for feeding the data to the trainer.
        pareto_manager: A ParetoManager which is responsible for maintaining a pareto
            front of models and saving these models on permanent storage.
        validator: A Validator object which is used to evaluate the models on multiple
            objective and multiple losses.
        max_empirical_losses: A list of losses (float) which is the approximation of the
            maximum empirical losses the model will have.
        common_descent_vector: A MultiObjectiveCDV, is responsible for combining the multiple
            gradients from the multiple losses into a single gradient.
        optimizer: A pytorch optimizer which is used to train the model.

    Attributes from the YAML file:
        seed: An integer, used to initialize the numpy and pytorch random seeds, default = 42.
        normalize_gradients: A boolean value, indicating whether to normalize the gradients
            while training the model or not, default = True.
        learning_rate: A float value, the learning rate that is given to the pytorch
            optimizer, if the optimizer is not given in the constructor, default = 1e-3.
        batch_size_training: An integer value, representing the batch sizes in which the data is
            fed to the trainer, default = 500.
        shuffle_training: A boolean value, indicating if the training data should be shuffled,
            default = True.
        drop_last_batch_training: A boolean value, indicating to drop the last incomplete batch,
            if the training dataset size is not divisible by the batch size, default = True.
        batch_size_validation: An integer value, representing the batch sizes in which the data is
            fed to the validator, default = 500.
        shuffle_validation: A boolean value, indicating if the validation data should be shuffled,
            default = True.
        drop_last_batch_validation: A boolean value, indicating to drop the last incomplete batch,
            if the validation dataset size is not divisible by the batch size, default = False.
        number_of_epochs: An integer value, indicating for how many epochs should the model
            be trained, default = 50.
        frank_wolfe_max_iter: An integer value, indicating the maximum number of iterations
            to be used by the frank wolfe algorithm in the commonDescentVector object,
            default = 100.
        anneal: A boolean value, indicating if annealing should be used while training the
            model, default = True.
        beta_start: If the anneal is used, this will be the first value of the beta,
            default = 0.
        beta_cap: If the anneal is used, this will be the maximum value of the beta,
            default = 0.3.
        beta_step: If the anneal is used, this is the amount by which to increase the beta
            every batch, default = 0.3/10000.
    """

    def __init__(self, data_handler, model, losses, validation_metrics,
                 save_to_path, params='yaml_files/trainer_params.yaml',
                 optimizer=None):
        """The constructor which initializes a trainer object.

        Arguments:
            data_handler: A MamoDataHandler object which feeds the data set to the trainer.
            model: A torch.nn.Module object which is the model that is being trained.
            losses: A list of Loss objects which represent the losses/objectives that the
                model is trained on.
            validation_metrics: A list of MetricAtK objects which are used to evaluate
                the model while the training and validation process.
            save_to_path: A path to a directory where the trained models from the Pareto
                front will be saved during training.
            params: Path to the yaml file with the trainger parameters, or a dictionary
                containing the parameters.
            optimizer: A pytorch optimizer which is used to train the model, if it is None,
                a default Adam optimizer is created.
        Raises:
            TypeError: If any of the arguments passed are not an instance of the expected
                class or are None, a TypeError will be raised.
            ValueError: If the directory which save_to_path references is not empty, a
                ValueError will be raised.
        """
        logger.info('Trainer: Started with initializing trainer...')
        self._check_input_(data_handler, model, losses,
                           validation_metrics, save_to_path, optimizer)
        self._read_params(params)
        self.data_handler = data_handler
        self.model = model
        self.losses = losses
        logger.info('Trainer: Losses: ')
        logger.info('Trainer: '.join(['%s ' % loss.name for loss in self.losses]))
        self.validation_metrics = validation_metrics
        logger.info('Trainer: Validation metrics: ')
        logger.info('Trainer: '.join(['%s ' % m.get_name() for m in self.validation_metrics]))
        self.save_to_path = save_to_path
        logger.info('Trainer: Saving models to: %s' % self.save_to_path)
        self.optimizer = optimizer
        # set cuda if available
        self.device = torch.device(
            'cuda' if torch.cuda.is_available() else 'cpu')
        logger.info('Trainer: Training on device: %s' % self.device)
        self._init_objects()
        logger.info('Trainer: Initialization done.')

    def _check_input_(self, data_handler, model, losses, validation_metrics, save_to_path, optimizer):
        """A helper function for the __init__ to check the input of the constructor.
        """
        if not isinstance(data_handler, MamoDataHandler):
            raise TypeError(
                'Please check you are using the right data handler object, or the right order of the attributes!')
        if not isinstance(model, nn.Module):
            raise TypeError(
                'Please check you are using the right model object, or the right order of the attributes!')
        if not hasattr(model, 'initialize_model'):
            raise TypeError(
                'Please check if your models has initialize_model() method defined!')
        # check if losses is None
        if losses is None:
            raise ValueError(
                'The losses are None, please make sure to give valid losses!')
        if not all([isinstance(x, Loss) for x in losses]):
            raise TypeError(
                'Please check you are using the right loss objects, or the right order of the attributes!')
        # check if there are at least two losses
        if len(losses) < 2:
            raise ValueError(
                'Please check you have defined at least two losses,'
                + ' for training with one loss use the Single Objective Loss class!')
        # check if validation metrics is None
        if validation_metrics is None:
            raise ValueError(
                'The validation_metrics are None, please make sure to give valid validation_metrics!')
        if not all([isinstance(x, MetricAtK) for x in validation_metrics]):
            raise TypeError(
                'Please check you are using the right metric objects, or the right order of the attributes!')
        # check if length is at least 1
        if len(validation_metrics) == 0:
            raise ValueError(
                'Please check you have defined at least one validation metric!')
        if not os.path.exists(save_to_path):
            os.mkdir(save_to_path)
        # checking if the save_to_path directory is empty
        if os.listdir(save_to_path):
            raise ValueError(
                'Please make sure that the directory where you want to save the models is empty!')
        # if the optimizer is not None, than has to be pytorch optimizer object
        if optimizer is not None:
            if not isinstance(optimizer, optim.Optimizer):
                raise TypeError(
                    'Please make sure that the optimizer is a pytorch Optimizer object!')

    def _read_params(self, params):
        """A helper function for the __init__ to read the configuration yaml file.
        """
        logger.info('Trainer: Reading trainer parameters.')
        if type(params) is str:
            with open(params, 'r') as stream:
                params = yaml.safe_load(stream)

        self.seed = int(params.get('seed', 42))
        logger.info('Trainer: Random seed: %d' % self.seed)

        self.normalize_gradients = bool(
            params.get('normalize_gradients', True))
        logger.info('Trainer: Normalize gradients: %s' %
                    self.normalize_gradients)

        self.learning_rate = float(params.get('learning_rate', 1e-3))
        logger.info('Trainer: Learning rate: %f' % self.learning_rate)

        self.batch_size_training = int(params.get('batch_size_training', 500))
        logger.info('Trainer: Batch size training: %d' %
                    self.batch_size_training)

        self.shuffle_training = bool(params.get('shuffle_training', True))
        logger.info('Trainer: Shuffle training: %d' %
                    self.shuffle_training)

        self.drop_last_batch_training = bool(
            params.get('drop_last_batch_training', True))
        logger.info('Trainer: Drop last batch training: %d' %
                    self.drop_last_batch_training)

        self.batch_size_validation = int(
            params.get('batch_size_validation', 500))
        logger.info('Trainer: Batch size validation: %d' %
                    self.batch_size_validation)

        self.shuffle_validation = bool(params.get('shuffle_validation', True))
        logger.info('Trainer: Shuffle validation: %d' %
                    self.shuffle_validation)

        self.drop_last_batch_validation = bool(
            params.get('drop_last_batch_validation', False))
        logger.info('Trainer: Drop last batch validation: %d' %
                    self.drop_last_batch_validation)

        self.number_of_epochs = int(params.get('number_of_epochs', 50))
        logger.info('Trainer: Number of epochs: %f' % self.number_of_epochs)

        self.frank_wolfe_max_iter = int(
            params.get('frank_wolfe_max_iter', 100))
        logger.info('Trainer: Frank Wolfe max iterations: %d' %
                    self.frank_wolfe_max_iter)

        self.anneal = bool(params.get('anneal', True))
        logger.info('Trainer: Annealing: %s' % self.anneal)

        if self.anneal and ('beta_start' not in params or 'beta_cap' not in params or 'beta_step' not in params):
            raise ValueError(('Please make sure that if anneal is set to True, '
                              'the beta_start, beta_cap and beta_step are all '
                              'present in the parameters yaml file!'))
        if self.anneal:
            self.beta_start = float(params.get('beta_start', 0))
            logger.info('Trainer: Beta start: %f' % self.beta_start)
            self.beta_cap = float(params.get('beta_cap', 0.3))
            logger.info('Trainer: Beta cap: %f' % self.beta_cap)
            self.beta_step = float(eval(params.get('beta_step', '0.3/10000')))
            logger.info('Trainer: Beta step: %f' % self.beta_step)

    def _init_objects(self):
        """A helper function for the __init__ to initialize different objects.
        """
        logger.info('Trainer: Initializing helper trainer objects.')
        np.random.seed(self.seed)
        torch.manual_seed(self.seed)
        self.model.initialize_model()
        self.model.to(self.device)
        self._train_dataloader = self.data_handler.get_train_dataloader(batch_size=self.batch_size_training,
                                                                        shuffle=self.shuffle_training,
                                                                        drop_last=self.drop_last_batch_training)
        self.pareto_manager = ParetoManager(PATH=self.save_to_path)
        val_dataloader = self.data_handler.get_validation_dataloader(
            batch_size=self.batch_size_validation, shuffle=self.shuffle_validation,
            drop_last=self.drop_last_batch_validation)
        self.validator = Validator(
            self.model, val_dataloader, self.validation_metrics, self.losses)
        self.max_empirical_losses = None
        if self.normalize_gradients:
            self.max_empirical_losses = self._compute_max_empirical_losses()
            logger.info('Trainer: Max empirical losses: %s' %
                        self.max_empirical_losses)
        copsolver = None
        if len(self.losses) <= 2:
            copsolver = AnalyticalSolver()
        else:
            copsolver = FrankWolfeSolver(max_iter=self.frank_wolfe_max_iter)
        self.common_descent_vector = MultiObjectiveCDV(
            copsolver=copsolver, max_empirical_losses=self.max_empirical_losses,
            normalized=self.normalize_gradients)
        # create default optimizer
        if self.optimizer is None:
            self.optimizer = optim.Adam(
                self.model.parameters(), lr=self.learning_rate)

    def _compute_max_empirical_losses(self):
        """A helper function for approximating the maximum empirical loss the model
        could have. It is called by _init_objects function.
        """
        # approximate the max loss empirically
        max_losses = [0] * len(self.losses)
        cnt = 0

        for batch in self._train_dataloader:
            # fetch data
            x = batch[0]
            x = Variable(x).to(self.device)
            y = batch[1]
            y = Variable(y).to(self.device)
            cnt += 1

            # forward pass
            model_output = self.model(x)

            for i, loss in enumerate(self.losses):
                # if annealing is done, the KL divergence is ignored when computing
                # the max empirical loss, therefore the anneal is set to 0
                if self.anneal:
                    L = loss.compute_loss(y, model_output, anneal=0)
                else:
                    L = loss.compute_loss(y, model_output)
                # compute the moving average term
                max_losses[i] = (cnt - 1) / cnt * \
                    max_losses[i] + 1 / cnt * L.item()
        return max_losses

    def _get_gradient_np(self):
        """A helper function for obtaining the gradients of the model in a numpy
        array.

        Before the first backward call, all grad attributes are set to None, and
        that is when the exception is thrown, and the parameters are returned.
        After the first backward pass, the gradient values are available and are
        returned by this function.
        """
        gradient = []
        try:
            for p in self.model.parameters():
                gradient.append(p.grad.cpu().detach().numpy().ravel())

            return np.concatenate(gradient)
        except Exception:
            size = 0
            for p in self.model.parameters():
                size += len(p.cpu().detach().numpy().ravel())

            return np.zeros(shape=size)

    def train(self):
        """The main method of this class. By calling this method, the traning process
        starts.
        """
        # model training
        logger.info('Trainer: Started training...')

        if self.anneal:
            beta = self.beta_start

        for epoch in range(self.number_of_epochs):
            # start time for current epoch
            start_time = time.time()

            # statistics
            training_loss = 0
            average_alpha = [0] * len(self.losses)
            cnt = 0

            # set model in train mode
            self.model.train()

            # do training
            for batch in self._train_dataloader:
                # create x
                x = batch[0]
                x = Variable(x).to(self.device)
                y = batch[1]
                y = Variable(y).to(self.device)

                # anneal beta
                if self.anneal:
                    beta += self.beta_step
                    beta = beta if beta < self.beta_cap else self.beta_cap

                # calculate the gradients
                gradients = []
                for i, loss in enumerate(self.losses):
                    # forward pass
                    model_output = self.model(x)
                    # calculate loss
                    if self.anneal:
                        L = loss.compute_loss(y, model_output, anneal=beta)
                    else:
                        L = loss.compute_loss(y, model_output)
                    # zero gradient
                    self.optimizer.zero_grad()
                    # backward pass
                    L.backward()
                    # get gradient for correctness objective
                    gradients.append(self._get_gradient_np())

                # calculate the losses
                losses_computed = []
                # forward pass
                model_output = self.model(x)

                for i, loss in enumerate(self.losses):
                    if self.anneal:
                        L = loss.compute_loss(y, model_output, anneal=beta)
                    else:
                        L = loss.compute_loss(y, model_output)
                    losses_computed.append(L)

                # get the final loss to compute the common descent vector
                final_loss, alphas = self.common_descent_vector.get_descent_vector(
                    losses_computed, gradients)
                # zero gradient
                self.optimizer.zero_grad()
                # backward pass
                final_loss.backward()
                # update parameters
                self.optimizer.step()

                # statistics....
                cnt += 1
                # moving average loss
                training_loss = (cnt - 1) / cnt * \
                    training_loss + 1 / cnt * final_loss.item()
                # moving average alpha
                for i, alpha in enumerate(alphas):
                    average_alpha[i] = (cnt - 1) / cnt * \
                        average_alpha[i] + 1 / cnt * alpha

                # time in milliseconds for current batch
                batch_time = (time.time() - start_time) / cnt * 1000

                # log progress
                if cnt % 10 == 0:
                    average_alpha_string = ', '.join(
                        ['%.4f']*len(average_alpha)) % tuple(average_alpha)
                    logger.info('Trainer: Batch: %d/%d, Batch time: %.2fms,' %
                                (cnt,
                                 int(np.round(self.data_handler.get_traindata_len() / self.batch_size_training)),
                                 batch_time)
                                + ' Training loss: %.3f, Alphas: [%s]' %
                                (training_loss, average_alpha_string))

            # do validation
            val_metrics, val_objectives = self.validator.evaluate(
                disable_anneal=self.anneal, verbose=False)
            val_loss = self.validator.combine_objectives(
                val_objectives, alphas=average_alpha, max_normalization=self.max_empirical_losses)

            # add the solution to the pareto manager
            self.pareto_manager.add_solution(val_metrics, self.model)

            # calculate epoch time
            epoch_time = time.time() - start_time

            val_metrics_string = ', '.join(
                ['%.4f']*len(val_metrics)) % tuple(val_metrics)
            val_objectives_string = ', '.join(
                ['%.4f']*len(val_objectives)) % tuple(val_objectives)
            logger.info('Trainer: Epoch: %d, Epoch time: %.2fs, Training loss: %.3f,' %
                        (epoch + 1, epoch_time, training_loss)
                        + ' Validation loss: %.3f, Validation metrics: [%s], Validation losses: [%s]' %
                        (val_loss, val_metrics_string, val_objectives_string))
        return val_loss