trainer.py

# from utils import config
import itertools
import os
import torch
import numpy as np 
import utils 
from sklearn import metrics
import random
import pickle 
from tqdm import tqdm 
from matplotlib import pyplot as plt
import torch.nn as nn
import shutil 
from sklearn.utils import resample  

# from apex import amp 

class Trainer(object):
    def __init__(self, seed, model, criterion, optimizer, params, loaders, loader_names,  device, checkpoint, config_str, batch_size=None,save_best_num = 1,save_every=None,cuda=True, verbose = True, eval_train = True, savename = ""):
        self.num_best_saved = 0

        self.best_checkpoint_to_score = {}
        self.save_best_num = save_best_num
        self.eval_train = eval_train
        self.verbose = verbose
        self.seed = seed
        self.model = model
        self.criterion = criterion
        self.optimizer = optimizer
        self.batch_size = batch_size
        self.device = device
        self.n_classes = config(config_str + ".num_classes")
        self.loaders = loaders
        self.loader_names = loader_names
        self.params = params
        self.n_iters = 0
        self.epoch = 0
        self.best_iter = 0
        self.config_str = config_str
        self.save_every = save_every 
        self.checkpoint = checkpoint
        split = self.checkpoint.split('/')
        split[-1] = 'best_' + split[-1]
        self.best_checkpoint = '/'.join(split)     
        self.stop = False
        self._iter = 0
        self.cuda = cuda and torch.cuda.is_available()
        self.num_train = 500 
        self.eval_all = True

        
        self.loader_pts = {}
        
        for loader_name, loader in list(zip(self.loader_names, self.loaders)):
            self.loader_pts[loader_name] = [] if loader_name != "train" else False 
            pos = []
            if (loader_name != "train"):
                for X, y, pt_id in loader:
                    for idx,p in enumerate(pt_id.numpy()):
                        self.loader_pts[loader_name].append(p)

        if (os.path.exists(self.checkpoint)):
            self._load_best()
        else:
            self.reset_logs()

        

    def _eval_all(self):
        for loader_name, loader in list(zip(self.loader_names, self.loaders)):
            if (loader_name == "valid" or (loader_name == "train" and self.eval_train)):
                self.latest_losses[loader_name], self.latest_scores[loader_name] = self._eval(loader, self.loader_pts[loader_name])
                self.losses[loader_name].append(self.latest_losses[loader_name])
                self.scores[loader_name].append(self.latest_scores[loader_name])

            else: 
                self.losses[loader_name].append(0)
                self.scores[loader_name].append(0) 
            
    def _print_status(self, print_all = False):
        if (print_all):
            for loader_name in self.loader_names:
                print(loader_name, ", loss: ", self.latest_losses[loader_name], ", score: ", self.latest_scores[loader_name])
        else:
                print("valid", ", loss: ", self.latest_losses["valid"], ", score: ", self.latest_scores["valid"])

    def reset_logs(self):
        
        self.best_iter = 0

        self.losses = {}
        self.scores = {}
        
        self.latest_losses = {}
        self.latest_scores = {}
        for loader_name in self.loader_names:
            self.losses[loader_name] = []
            self.scores[loader_name] = []
            self.latest_losses[loader_name] = []
            self.latest_scores[loader_name] = []
              
        self._eval_all()
        self.best_loss, self.best_score = self._eval(self.loaders[1], self.loader_pts["valid"])
        self._save(self.best_score)
        
    def fit(self):
        self.stop = self._check_early_stopping()                          
        if (self.stop): 
            self._save_final()
        while not self.stop:
            if (self.verbose):
                print("Epoch:", self.epoch)
            for i, (X, y, pt_id) in tqdm(enumerate(self.loaders[0]), desc='epoch '+str(self.epoch) + ', iter ' + str(self._iter), leave=False):
                loss, pred = self._train_batch(X, y)
                if self.save_every and self._iter % self.save_every == 0:
                    print("Iteration:", self._iter, "saving")
                    self._eval_all()
                    self._save(self.latest_scores["valid"])
                    self.stop = self._check_early_stopping()
                self._iter += 1
               
            self._eval_all()
            self.epoch += 1 
            self._print_status()
            self._save(self.latest_scores["valid"])
            self.stop = self._check_early_stopping()
            
            if (self.stop):
                self._save_final()
                break
            
                

    def _train_batch(self, X, y):
        self.model.train()
        X,y = X.to(self.device), y.to(self.device)
        self.optimizer.zero_grad()
        output = self.model(X)
            
        if (config(self.config_str + ".mask")):
            output[y<0] = 0
            y[y<0] = 0
        
        loss = self.criterion(torch.squeeze(output), torch.squeeze(y))
        loss.backward()
        self.optimizer.step()
        

        return loss, (output.data.detach().cpu(), y.data.detach().cpu())
  
    def get_roc_CI(self, y_true, y_score):
        lower = []
        upper = []
        if (self.n_classes > 1):
            for i in range(self.n_classes):
                true = y_true[:,i]
                score = y_score[:,i]
                roc_curves, auc_scores, aupr_scores = [], [], []
                if (config(self.config_str + ".mask")):
                    score[true<0] = 0
                    true[true<0] = 0
                for j in range(1000):
                    yte_true_b, yte_pred_b = resample(true, score, replace=True, random_state=j)

                    roc_curve = metrics.roc_curve(yte_true_b, yte_pred_b)
                    auc_score = metrics.roc_auc_score(yte_true_b, yte_pred_b)
                    aupr_score = metrics.auc(*metrics.precision_recall_curve(yte_true_b, yte_pred_b)[1::-1])
                    roc_curves.append(roc_curve)
                    auc_scores.append(auc_score)
                    aupr_scores.append(aupr_score)
                lower.append(np.percentile(auc_scores, 2.5))
                upper.append(np.percentile(auc_scores, 97.5))
                
        else:
            roc_curves, auc_scores, aupr_scores = [], [], []
            if (config(self.config_str + ".mask")):
                y_score[y_true<0] = 0
                y_true[y_true<0] = 0
            for j in range(1000):
                try:
                    yte_true_b, yte_pred_b = resample(y_true, y_score, replace=True, random_state=j)


                    roc_curve = metrics.roc_curve(yte_true_b, yte_pred_b)
                    auc_score = metrics.roc_auc_score(yte_true_b, yte_pred_b)
                    aupr_score = metrics.auc(*metrics.precision_recall_curve(yte_true_b, yte_pred_b)[1::-1])

                    roc_curves.append(roc_curve)
                    auc_scores.append(auc_score)
                    aupr_scores.append(aupr_score)
                except:
                    j -= 1
                    
            lower = np.percentile(auc_scores, 2.5)
            upper = np.percentile(auc_scores, 97.5)
    
        return lower, upper 

    def _eval(self, data_loader, pts, train = False, get_CI = False):
        self.model.eval()
        running_loss = []
        running_pred = []
        num_processed = 0
        with torch.no_grad():
            for X, y, pt_id in data_loader:
                X,y = X.to(self.device), y.to(self.device)
                output = self.model(X).float()

                predicted = predictions(output.data)
                running_pred.append((predicted.data.detach().cpu().numpy(), y.data.detach().cpu().numpy(), pt_id.data.detach().cpu().numpy()))
                
                 # mask loss if needed 
                if (config(self.config_str + ".mask")):
                    output[y<0] = 0
                    y[y<0] = 0
                    
                loss = self.criterion(torch.squeeze(output), torch.squeeze(y))
                running_loss.append(loss.data.detach().cpu())
                num_processed += len(y)

        return np.mean(running_loss), self._get_score(running_pred, pts, get_CI)
    # check for early stopping 
    def _check_early_stopping(self):
        if (not config(self.config_str + ".early_stop")):
            return self.epoch > 3
        stop = False 
        if(len(self.losses['valid']) > 5):
            min_epoch = np.argmin(self.losses['valid'])
            if  min_epoch < len(self.losses['valid']) - 5 or np.abs(self.losses['valid'][-1] - self.losses['valid'][-2]) < 0.001:
                stop = True  
        return stop 
   
    def _get_score(self, running_pred, pt_ids, get_CI = False):
        y_pred, y_true, pt_ids = zip(*running_pred)
    
        y_pred = np.concatenate(y_pred)
        y_true = np.concatenate(y_true)
        pt_ids = np.concatenate(pt_ids)

        assert(len(y_pred) == len(y_true) == len(pt_ids))
        
        unique_pt_ids = np.unique(pt_ids)
        unique_predictions = []
        unique_truth_values = []
        for pt_id in unique_pt_ids:
            indices = np.where(pt_ids == pt_id)[0]
            if(len(y_pred[indices]) == 1):
                unique_predictions.append(y_pred[indices][0])
                unique_truth_values.append(y_true[indices][0])
            else:
                unique_predictions.append(np.average(y_pred[indices], axis = 0))
                unique_truth_values.append(np.average(y_true[indices], axis = 0))
            
        unique_predictions = np.squeeze(np.array(unique_predictions))

        unique_truth_values = np.squeeze(np.array(unique_truth_values))

        score = []
        for n in range(self.n_classes):
            if (self.n_classes == 1):
                unique_truth = unique_truth_values
                unique_pred = unique_predictions 
            else:
                unique_truth = unique_truth_values[:,n]
                unique_pred = unique_predictions[:,n] 
            try:
                if (config(self.config_str + ".mask")):
                    predictions = unique_pred[unique_truth > -1]
                    truths = unique_truth[unique_truth > -1]
                    score.append(metrics.roc_auc_score(truths, predictions))
                else:
                    score.append(metrics.roc_auc_score(unique_truth, unique_pred))
            except:
                print("Only 1 class present")
                score.append(0.5)
        assert(len(unique_predictions) == len(unique_pt_ids))
        
        
        if (get_CI):
            lower, upper = self.get_roc_CI(unique_truth_values, unique_predictions)
            return score, lower, upper 
        return score
    
    def _save_final(self):
#         load best checkpoint 
        print("finished training, loading best checkpoint for test set")
        self._load_best(best = True)
#             get test set ROC_CI 
        self.best_lower = {}
        self.best_upper = {}
        self.best_test = {}
        self.best_loss = {}
        for loader_name in self.loader_names:
            if loader_name != "train" or (loader_name == "train" and self.eval_train):
                self.best_loss[loader_name], (self.best_test[loader_name], self.best_lower[loader_name], self.best_upper[loader_name]) = self._eval(self.loaders[self.loader_names.index(loader_name)], self.loader_pts[loader_name], train = loader_name == "train", get_CI = True)


    def _save(self, new_score):
        if (self.save_best_num > 1):
            highest_best_score = 0
#             check if we have 10 anyways 
            if (len(self.best_checkpoint_to_score) == self.save_best_num):
                sorted_checkpoints = sorted(self.best_checkpoint_to_score.items(), key = 
                             lambda kv:(np.mean(kv[1]), kv[0]))
    
                lowest_best_score = sorted_checkpoints[0][1]
                lowest_best_checkpoint = sorted_checkpoints[0][0]
                highest_best_score = sorted_checkpoints[-1][1]
                highest_best_checkpoint = sorted_checkpoints[-1][0]
                is_best = bool(np.mean(new_score) >= np.mean(lowest_best_score))
            else:
                is_best = True 
        else:
            is_best = bool(np.mean(new_score) >= np.mean(self.best_score))
        if is_best:
            self.best_score = new_score
            
#             avoids duplicate maxes 
            if self.epoch and (self.save_best_num == 1):
                occurences = np.where(np.array(self.scores['valid']) == max(self.scores['valid']))[0][-1]
                self.best_iter = occurences
            else:
                self.best_iter = self._iter
            split = self.best_checkpoint.find("checkpoint.")
            new_best_checkpoint = self.best_checkpoint[0:split] + "_" + str(self.num_best_saved) + "_" + self.best_checkpoint[split:]
            
            

 
        self.best_lower = {}
        self.best_upper = {}
        self.best_test = {}
        state = {
            'best_iter' : self.best_iter,
            'best_score' : self.best_score, 
            'best_loss' : self.best_loss, 
            '_epoch': self.epoch,
            '_iter' : self._iter,
            'batch_size': self.batch_size,
            'state_dict': self.model.state_dict(),
            'arch': str(type(self.model)),
            'optimizer': self.optimizer.state_dict(),
            'losses' : self.losses, 
            'scores' : self.scores, 
            'latest_losses' : self.latest_losses, 
            'latest_scores' : self.latest_scores, 
            'params' : self.params,
            'stop' : self.stop, 
            'best_lower': self.best_lower, 
            'best_upper': self.best_upper, 
            'best_test': self.best_test,
            'best_checkpoint_to_score' : self.best_checkpoint_to_score
        }
        torch.save(state, self.checkpoint)
        if is_best: 
            if (self.save_best_num == 1):
                shutil.copyfile(self.checkpoint, self.best_checkpoint)
            else:
                self.best_checkpoint_to_score[new_best_checkpoint]= new_score
                if (np.mean(new_score) >= np.mean(highest_best_score)):
                    self.best_checkpoint = new_best_checkpoint
                if (len(self.best_checkpoint_to_score) > self.save_best_num):
                    try:
                        os.remove(lowest_best_checkpoint)

                        del self.best_checkpoint_to_score[lowest_best_checkpoint]
                    except: 
                        pass

                torch.save(state, self.checkpoint)
                self.num_best_saved += 1
                shutil.copyfile(self.checkpoint, new_best_checkpoint)
        torch.save(state, self.checkpoint)
                         
                    
    def _load_best(self, best = False):
        if (best):
            if (self.save_best_num > 1):
                sorted_checkpoints = sorted(self.best_checkpoint_to_score.items(), key = 
                             lambda kv:(np.mean(kv[1]), kv[0]))

                highest_best_score = sorted_checkpoints[-1][1]
                highest_best_checkpoint = sorted_checkpoints[-1][0]
                checkpoint = torch.load(highest_best_checkpoint)
            else:
                checkpoint = torch.load(self.best_checkpoint)
            self.best_checkpoint_to_score = torch.load(self.checkpoint)["best_checkpoint_to_score"]

        else:
            checkpoint = torch.load(self.checkpoint)
            self.best_checkpoint_to_score = checkpoint["best_checkpoint_to_score"]
        
        self.epoch = checkpoint["_epoch"]
        self._iter = checkpoint["_iter"]
        self.batch_size = checkpoint["batch_size"]
        self.model.load_state_dict(checkpoint['state_dict'])
        self.optimizer.load_state_dict(checkpoint["optimizer"])
        self.latest_losses = checkpoint['latest_losses']
        self.losses = checkpoint['losses']
        
        self.latest_scores = checkpoint['latest_scores']
        self.scores = checkpoint['scores']
        
        self.best_lower = checkpoint['best_lower']
        self.best_upper = checkpoint['best_upper']
        self.best_test = checkpoint['best_test']

        self.params = checkpoint['params']
        self.stop = checkpoint['stop']
        
        
        if (not len(self.scores['valid'])):
            self._eval_all()

        self.best_iter = checkpoint['best_iter']
        self.best_loss = checkpoint['best_loss']
        self.best_score = checkpoint['best_score']
        if (self.save_best_num == 1):
            assert(self.epoch + 1 == len(self.scores["valid"]))

    

def count_parameters(model):
    return sum(p.numel() for p in model.parameters() if p.requires_grad)

def predictions(logits):
    """
    Given the network output, determines the predicted class index

    Returns:
        the predicted class output as a PyTorch Tensor
    """
    sig = torch.nn.Sigmoid()
    return sig(logits)
    #

    
def config(attr):
    """
    Retrieves the queried attribute value from the config file. Loads the
    config file on first call.
    """
    if not hasattr(config, 'config'):
        with open('config.json') as f:
            config.config = eval(f.read())
    node = config.config
    for part in attr.split('.'):
        node = node[part]
    return node