trainer.py

import torch
from abc import abstractmethod
from numpy import inf
# from logger import TensorboardWriter


import numpy as np
import torch


from tqdm import tqdm
import math
from torch.optim.lr_scheduler import CyclicLR, OneCycleLR, ReduceLROnPlateau, StepLR
import pandas as pd
from sklearn.metrics import roc_auc_score
from utils import mkdir, get_logger 
import os

def get_lr(optimizer):
    for param_group in optimizer.param_groups:
        return param_group['lr']


def train_step(epoch, model, criterion, optimizer, lr_scheduler, device, train_loader):
    logger = get_logger('train')
    desc = ("Training train set - epoch %i" % (epoch))
    logger.info(desc)

    # Turn on training mode which enables dropout.
    model.train()

    running_loss = 0
    lrs = []
    with tqdm(range(len(train_loader)), desc=desc, total=(len(train_loader))) as t:
        # do batches
        for batch_idx, batch in zip(t, train_loader):
            padded_sequence_tensors, labels, seq_lengths, samples_weights, samples_times, samples_hashes = batch  # batch return padded_sequence_tensor, syscalls_label, seq_length, sample_weight, self.times[item], self.hashes[item]
            labels = torch.Tensor(labels).to(device).long()
            # seq_lengths = torch.Tensor(seq_lengths).to(device).long()  # TODO: pytorch-1.7 (>1.4) bug - 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor
            padded_sequence_tensors = torch.stack(padded_sequence_tensors).to(device)
            samples_weights = torch.Tensor(samples_weights).to(device)

            # forward the packed sequences batch and get the predicted log probabilities
            predicted_log_probs = model(padded_sequence_tensors, seq_lengths)  # , hn)
            pred_loss = criterion(predicted_log_probs, labels)
            # if not math.isfinite(pred_loss):  # isfinite isnan
            #     logger.info("Loss is {}, resetting loss and skipping training iteration".format(pred_loss))
            #     logger.info('Loss values were: ', pred_loss)
            #     pred_loss = torch.tensor(0)  # pred_loss = {k: torch.tensor(0) for k, v in pred_loss.items()}
            # else:

            # apply current batch samples_weights
            samples_weights_vector = samples_weights.view(-1,1,1)
            pred_loss = (pred_loss*samples_weights_vector).mean()
            
            optimizer.zero_grad()
            # TODO: If ``False``, the graph used to compute
            # the grads will be freed. Note that in nearly all cases setting
            # this option to True is not needed and often can be worked around
            # in a much more efficient way. Defaults to the value of `create_graph`.
            pred_loss.backward(retain_graph=True)  # added retain_graph=True after BiGru
            optimizer.step()
            if isinstance(lr_scheduler, CyclicLR) or isinstance(lr_scheduler, OneCycleLR):
                # logger.info("LR scheduler is of type CyclicLR / OneCycleLR, doing step.")
                lr_scheduler.step()
                lr = get_lr(optimizer)
                lrs.append(lr)
                #logger.info("new lr is %.10f." %(lr))
            # if warmup_scheduler is not None:
                # warmup_scheduler.step()
            # TODO: should I limit and do this only once in X batches?
            scalar_loss = pred_loss.item()  # correct = predicted_log_probs.argmax(dim=1).eq(labels).sum().item()
            running_loss += scalar_loss
            avg_loss = running_loss / (batch_idx + 1)
            lr = get_lr(optimizer)
            t.set_postfix(loss=avg_loss, LR=lr)
            padded_sequence_tensors = labels = seq_lengths = predicted_log_probs = pred_loss = None
            # if batch % args.log_interval == 0 and batch > 0:
            #     cur_loss = total_loss / args.log_interval
            #     elapsed = time.time() - start_time
            #     print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.2f} | ms/batch {:5.2f} | '
            #             'loss {:5.2f} | ppl {:8.2f}'.format(
            #         epoch, batch, len(train_data) // args.bptt, lr,
            #         elapsed * 1000 / args.log_interval, cur_loss, math.exp(cur_loss)))
            #     total_loss = 0
            #     start_time = time.time()
    if len(lrs) > 1:
        logger.info("Finished training train set - epoch %i. \nnew LR is %.10f. LRs: first %.10f last %.10f min %.10f max %.10f." 
                    %(epoch, lr, lrs[0], lrs[-1], min(lrs), max(lrs)))



def eval_step(epoch, model, criterion, optimizer, lr_scheduler, device, eval_loader, set_name):
    logger = get_logger('train')
    desc = ("Evaluating %s set - epoch %i" % (set_name, epoch))
    logger.info(desc)
    # Turn on eval mode which enables 
    model.eval()
    running_loss, running_log_preds, running_labels, running_times, running_hashes = 0, [], [], [], []
    with torch.no_grad():
        with tqdm(range(len(eval_loader)), desc=desc, total=(len(eval_loader))) as t:
            # do batches
            for batch_idx, batch in zip(t, eval_loader):
                # update prog bar
                padded_sequence_tensors, samples_labels, seq_lengths, _, samples_times, samples_hashes = batch  # batch return padded_sequence_tensor, syscalls_label, seq_length, sample_weight, self.times[item], self.hashes[item]
                labels = torch.Tensor(samples_labels).to(device).long()
                # seq_lengths = torch.Tensor(seq_lengths).to(device).long()  # shuould be on cpu from pytorch-1.7 (>1.4) bug - 'lengths' argument should be a 1D CPU int64 tensor, but got 1D cuda:0 Long tensor
                padded_sequence_tensors = torch.stack(padded_sequence_tensors).to(device)

                # forward the packed sequences batch and get the predicted log probabilities
                predicted_log_probs = model(padded_sequence_tensors, seq_lengths)  # , hn)
                pred_loss = criterion(predicted_log_probs, labels)
                
                # TODO: do this in eval as well?
                # if not math.isfinite(pred_loss):  # isfinite isnan
                #     print("Loss is {}, resetting loss and skipping eval iteration".format(pred_loss))
                #     print('Loss values were: ', pred_loss)
                #     pred_loss = torch.tensor(0)  # pred_loss = {k: torch.tensor(0) for k, v in pred_loss.items()}
                # else:
                #     # TODO: If ``False``, the graph used to compute
                #     # the grads will be freed. Note that in nearly all cases setting
                #     # this option to True is not needed and often can be worked around
                #     # in a much more efficient way. Defaults to the value of `create_graph`.
                #     pred_loss.backward(retain_graph=True)  # added retain_graph=True after BiGru
                #     optimizer.step()

                scalar_loss = pred_loss.item()  # correct = predicted_log_probs.argmax(dim=1).eq(labels).sum().item()
                running_loss += scalar_loss
                avg_loss = running_loss / (batch_idx + 1)
                # acc = metric(F.softmax(predicted_log_probs, dim=1), labels)

                # save batch stats
                running_log_preds+=predicted_log_probs.tolist()
                running_labels+=samples_labels
                running_times+=samples_times
                running_hashes+=samples_hashes
                
                t.set_postfix(loss=avg_loss)
                padded_sequence_tensors = labels = seq_lengths = predicted_log_probs = pred_loss = None

                # if batch % args.log_interval == 0 and batch > 0:
                #     cur_loss = total_loss / args.log_interval
                #     elapsed = time.time() - start_time
                #     print('| epoch {:3d} | {:5d}/{:5d} batches | lr {:02.2f} | ms/batch {:5.2f} | '
                #             'loss {:5.2f} | ppl {:8.2f}'.format(
                #         epoch, batch, len(train_data) // args.bptt, lr,
                #         elapsed * 1000 / args.log_interval, cur_loss, math.exp(cur_loss)))
                #     total_loss = 0
                #     start_time = time.time()
    epoch_stats_df = pd.DataFrame()
    # running_preds = np.exp(running_log_preds.cpu()) 
    # epoch_stats_df['predicted_probability'] = running_preds
    epoch_stats_df['predicted_log_probability'] = running_log_preds
    epoch_stats_df['label'] = running_labels
    epoch_stats_df['time'] = running_times
    epoch_stats_df['hash'] = running_hashes
    return avg_loss, epoch_stats_df

def load_opt_sched(filename):
    checkpoint = torch.load(filename) 
    return checkpoint['optimizer_state_dict'], checkpoint['lr_scheduler_state_dict']
    # load architecture params from checkpoint.
    # if checkpoint['arch'] != self.config['arch']:
        # self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
                            # "checkpoint. This may yield an exception while state_dict is being loaded.")
    # load architecture params from checkpoint.
    # if checkpoint['arch'] != self.config['arch']:
        # self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
    #                         # "checkpoint. This may yield an exception while state_dict is being loaded.")
    # # load optimizer state from checkpoint only when optimizer type is not changed.
    # if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
    #     self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
    #                         "Optimizer parameters not being resumed.")
    
def save_checkpoint(output_dir, save_checkpoints, epoch, architecture, model, optimizer, lr_scheduler):
    checkpoints_dir = os.path.join(output_dir, 'checkpoints')
    mkdir(checkpoints_dir)
    state = {}
    if save_checkpoints == 'all' and model and optimizer and lr_scheduler:
        state = {
            'architecture': architecture,
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'lr_scheduler_state_dict': lr_scheduler.state_dict()        
            # 'monitor_best': self.mnt_best,
            # 'config': self.config
        }
    elif save_checkpoints == 'lr' and optimizer and lr_scheduler:
        state = {
            'architecture': architecture,
            'epoch': epoch,
            'optimizer_state_dict': optimizer.state_dict(),
            'lr_scheduler_state_dict': lr_scheduler.state_dict()        
            # 'monitor_best': self.mnt_best,
            # 'config': self.config
        }
    else:
        return
    filename = os.path.join(checkpoints_dir, 'checkpoint-epoch-'+str(epoch)+'.pt')
    torch.save(state, filename)
    # self.logger.info("Saving checkpoint: {} ...".format(filename))
    # if save_best:
    #     best_path = str(self.checkpoint_dir / 'model_best.pt')
    #     torch.save(state, best_path)
    #     self.logger.info("Saving current best: model_best.pt ...")

def save_epoch_stats(output_dir, set_name, epoch, epoch_stats_df):
    epochs_stats_dir = os.path.join(output_dir, 'epochs_stats', set_name)
    mkdir(epochs_stats_dir)
    epoch_stats_df.to_pickle(os.path.join(epochs_stats_dir, 'epoch_'+str(epoch)+'_stats_pickle.gz'))
    # epoch_stats_df.to_pickle(os.path.join(epochs_stats_dir, 'epoch'+str(epoch)+'_stats_df.gzip'), compression='gzip')
    # epoch_stats_df.to_pickle(os.path.join(epochs_stats_dir, 'epoch'+str(epoch)+'_stats_df.bz2'), compression='bz2')
    # epoch_stats_df.to_pickle(os.path.join(epochs_stats_dir, 'epoch'+str(epoch)+'_stats_df.zip'), compression='zip')
    # epoch_stats_df.to_pickle(os.path.join(epochs_stats_dir, 'epoch'+str(epoch)+'_stats_df.xz'), compression='xz')

def get_accuracy(output, target):
    with torch.no_grad():
        output = torch.tensor(output)
        target = torch.tensor(target, dtype=torch.long)
        pred = torch.argmax(output, dim=1)
        assert pred.shape[0] == len(target)
        correct = 0
        correct += torch.sum(pred == target).item()
    return correct / len(target)


"""
Full training logic
"""
# Full training logic
def train(args, model, criterion, optimizer, lr_scheduler, device, train_loader, val_loader=None, test_loader=None, future_loader=None, len_epoch=None):
    logger = get_logger('train')

    # not_improved_count = 0

    # Training logic for an epoch
    for epoch in range(1, args.epochs + 1):
        # do train step for this epoch
        train_step(epoch, model, criterion, optimizer, lr_scheduler, device, train_loader)
        
        # do evaluation of each set    # if epoch % args.eval_every == (args.eval_every - 1):
        for eval_loader, set_name in zip([train_loader, val_loader, test_loader, future_loader], ['train', 'valid', 'test', 'future']):
            # skip None sets for retrain/double-retrain procedures
            if eval_loader:
                eval_loss, epoch_stats_df = eval_step(epoch, model, criterion, optimizer, lr_scheduler, device, eval_loader, set_name)
                save_epoch_stats(args.output_dir, set_name, epoch, epoch_stats_df)
                auc = roc_auc_score(epoch_stats_df['label'], epoch_stats_df.predicted_log_probability.map(lambda x: x[1]))
                accuracy = get_accuracy(epoch_stats_df.predicted_log_probability, epoch_stats_df['label'])
                logger.info("finished eval of epoch-%d of %s-set \t eval_loss=%.4f \t accuracy=%.4f \t auc=%.4f" %(epoch, set_name, eval_loss, accuracy, auc))
                # Do LR-step (from type ReduceLROnPlateau) and early-stopping according to validation-set.
                if set_name == 'valid':
                    # double-retrain needs to do set the same LR as retrain checkpoints. 
                    if isinstance(lr_scheduler, ReduceLROnPlateau):
                        logger.info("LR scheduler is of type ReduceLROnPlateau, doing step(auc).")
                        lr_scheduler.step(auc)
                        lr = get_lr(optimizer)
                        logger.info("new lr is %.10f." %(lr))
                    if isinstance(lr_scheduler, StepLR):
                        logger.info("LR scheduler is of type StepLR, doing step.")
                        lr_scheduler.step()
                        lr = get_lr(optimizer)
                        logger.info("new lr is %.10f." %(lr))                    
                # TODO: if none train set for eval then switch to 'future' 
                if set_name == 'train' and args.train_procedure == 'double_retrain':
                    filename = os.path.join(args.output_dir, "..", "retrain", 'checkpoints', 'checkpoint-epoch-'+str(epoch)+'.pt')
                    if os.path.exists(filename):
                        loaded = load_opt_sched(filename)
                        if loaded:
                            optimizer.load_state_dict(loaded[0])
                            lr_scheduler.load_state_dict(loaded[1])

                    # best_loss, stop_step, stop = early_stopping(eval_loss, best_loss, stop_step, args.patience)
                    # if stop:
                    #     break
                    # if (stop_step == 0) & (args.save_results):
                    #     best_epoch = epoch
                    #     torch.save(model.state_dict(), model_weights / (model_name + ".pt"))
        
        # save opt sched model
        if args.save_checkpoints != 'none':
            save_checkpoint(args.output_dir, args.save_checkpoints, epoch, args.architecture, model, optimizer, lr_scheduler)
            


# def _resume_checkpoint(self, resume_path):
#     """
#     Resume from saved checkpoints

#     :param resume_path: Checkpoint path to be resumed
#     """
#     resume_path = str(resume_path)
#     self.logger.info("Loading checkpoint: {} ...".format(resume_path))
#     checkpoint = torch.load(resume_path)
#     self.start_epoch = checkpoint['epoch'] + 1
#     self.mnt_best = checkpoint['monitor_best']

#     # load architecture params from checkpoint.
#     if checkpoint['config']['arch'] != self.config['arch']:
#         self.logger.warning("Warning: Architecture configuration given in config file is different from that of "
#                             "checkpoint. This may yield an exception while state_dict is being loaded.")
#     self.model.load_state_dict(checkpoint['state_dict'])

#     # load optimizer state from checkpoint only when optimizer type is not changed.
#     if checkpoint['config']['optimizer']['type'] != self.config['optimizer']['type']:
#         self.logger.warning("Warning: Optimizer type given in config file is different from that of checkpoint. "
#                             "Optimizer parameters not being resumed.")
#     else:
#         self.optimizer.load_state_dict(checkpoint['optimizer'])

#     self.logger.info("Checkpoint loaded. Resume training from epoch {}".format(self.start_epoch))