trainer.py

import json
import os
import math
import torch

from tqdm import tqdm
from sklearn.metrics import f1_score, roc_auc_score, precision_recall_fscore_support, classification_report, accuracy_score, hamming_loss
import logging
from torch import nn



from utils import save_model

def Trainer(args, model, train_loader, val_loader, test_loader, opt, schdlr, seed, num_labels, output_dir, run):
    torch.manual_seed(seed)

    step = 0
    epochs = args.epochs + 1
    metric_name = args.best_metric_name
    best_metric_value = None
    early_stopping_counter = 0
    early_stopping_patience = args.early_stopping_patience

    torch.autograd.set_detect_anomaly(True)

    for epoch in range(1, epochs):

        ###################################################
        # training
        ###################################################

        logging.info(f"Epoch: {epoch}")

        loss_epoch, loss_step = [], 0
        distance_epoch, distance_step = [], 0
        bce_loss_epoch, bce_loss_step = [], 0
        pc_bce_loss_epoch, pc_bce_loss_step = [], 0

        model.train()

        for batch in tqdm(train_loader):
            labels = batch[-1] 

            if args.use_poincare_loss or args.use_euclidean_loss or args.use_cosine_loss or args.use_contrastive_loss:  
                result = model(batch)
                loss = result['loss']
                bce_loss = result['bce_loss']
                distance = result['distance']

                distance = distance / args.accumulate_grad
                bce_loss = bce_loss / args.accumulate_grad

                bce_loss_step += bce_loss.item()
                distance_step += distance.item()
                
            else:
                result = model(batch)
                loss = result['loss']
            

            loss = loss / args.accumulate_grad
            loss.backward()  # update gradients
            loss_step += loss.item()

            if (step + 1) % args.accumulate_grad == 0:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)

                opt.step()
                opt.zero_grad()
                schdlr.step()

                loss_epoch.append(loss_step)
                loss_step = 0

                if args.use_poincare_loss or args.use_euclidean_loss or args.use_cosine_loss or args.use_contrastive_loss:
                    bce_loss_epoch.append(bce_loss.item())
                    distance_epoch.append(distance.item())

                    distance_step = 0
                    bce_loss_step = 0

            step += 1

        train_results = {"epoch": epoch, "train/loss": sum(loss_epoch) / len(loss_epoch), "learning_rate": schdlr.get_last_lr()[0]}
        if args.use_poincare_loss or args.use_euclidean_loss or args.use_cosine_loss or args.use_contrastive_loss:
            train_results['train/bce_loss'] = sum(bce_loss_epoch) / len(bce_loss_epoch)
            train_results['train/distance'] = sum(distance_epoch) / len(distance_epoch)
        
        # log results
        logging.info(f"Loss: {sum(loss_epoch)/len(loss_epoch)}, learning_rate: {schdlr.get_last_lr()[0]}")

        # save training results
        with open(os.path.join(output_dir, 'train_results.json'), 'w', encoding='utf8') as f:
            json.dump(train_results, f)

        # log training results to wandb
        run.log(train_results)


        ###################################################
        # validation
        ###################################################

        all_labels_val, all_preds_val = [], []
        loss_epoch = []
        distance_epoch = []
        bce_loss_epoch = []
        pc_bce_loss_epoch = [] 

        logging.info("Validation ...")

        model.eval()

        for batch in tqdm(val_loader):
            labels = batch[-2]
            with torch.no_grad():

                result = model(batch)
                loss = result['loss']
                logits = result['logits']

            loss_epoch.append(loss.detach().item())


            # process labels
            all_labels_val.append(labels) 

            ######### multi-label ########
            if args.dataset != 'WOS':
                sig = logits.sigmoid()
                # if args.do_distance_based_classification:
                #   final_pred = x | y for x, y in zip(bert_sig, distance_sig)
                if len(sig.shape) == 1:
                    sig = torch.unsqueeze(sig, 0)

                all_preds_val.append(sig)

            ######### multi-class ########
            else:
                sm = torch.max(logits, dim=1)[1]
                all_preds_val.append(sm)

        all_labels_val = torch.cat(all_labels_val, dim=0).cpu().numpy()
        all_preds_val = torch.cat(all_preds_val, dim=0).cpu().numpy()

        # any label that is predicted with a probability higher than the threshold (default = 0.5) is assigned
        all_preds_val[all_preds_val < args.threshold] = 0
        all_preds_val[all_preds_val >= args.threshold] = 1

        acc_val = accuracy_score(all_labels_val, all_preds_val)  # exact match ratio in case of multi-label
        prec_val, rec_val, _, _ = precision_recall_fscore_support(all_labels_val, all_preds_val, average='macro')
        hl_val = hamming_loss(all_labels_val, all_preds_val)
      
        #auc_roc_val = roc_auc_score(all_labels_val, all_preds_val)  # AUC-ROC

        f1_macro_val = f1_score(all_labels_val, all_preds_val, average='macro')  # F1-macro
        f1_micro_val = f1_score(all_labels_val, all_preds_val, average='micro')  # F1-micro
        eval_loss = sum(loss_epoch) / len(loss_epoch)

        # log results
        logging.info(f"Validation Loss: {eval_loss}")
        logging.info(f"Validation F1-macro: {f1_macro_val}")
        logging.info(f"Validation F1-micro: {f1_micro_val}")
        logging.info(f"Validation Accuracy: {acc_val}")
        logging.info(f"Hamming Loss: {hl_val}")
        #logging.info(f"Validation ROC-AUC: {auc_roc_val}")

        logging.info(classification_report(all_labels_val, all_preds_val))

        eval_results = {
                        "epoch": epoch, "eval/loss": eval_loss,
                        "eval/accuracy": acc_val, "eval/f1_macro": f1_macro_val, "eval/f1_micro": f1_micro_val,
                        "eval/hamming_loss": hl_val, "eval/recall": rec_val, "eval/precision": prec_val,
                        #"eval/auc_roc": auc_roc_val
                        }
        
        # save results
        with open(os.path.join(output_dir, 'eval_results.json'), 'w', encoding='utf8') as f:
            json.dump(eval_results, f)
        
        # log evaluation results to wandb
        run.log(eval_results)

        # Check for early stopping
        if args.early_stopping_patience:

            if best_metric_value == None: 
                    best_metric_value = eval_results[metric_name]
                    save_model(args, model, output_dir, epoch)

            if metric_name == 'eval/loss':
                if eval_results[metric_name] >= best_metric_value:
                    early_stopping_counter += 1
            
                elif eval_results[metric_name] < best_metric_value:
                    early_stopping_counter = 0
                    best_metric_value = eval_results[metric_name]
                    save_model(args, model, output_dir, epoch)


            else:
                if eval_results[metric_name] <= best_metric_value:
                    early_stopping_counter += 1
                
                elif eval_results[metric_name] > best_metric_value:
                    early_stopping_counter = 0
                    best_metric_value = eval_results[metric_name]
                    save_model(args, model, output_dir, epoch)

            if early_stopping_counter >= early_stopping_patience:
                logging.info(f"Early stopping at epoch {epoch}")
                break