train_ranker.py

import argparse
import os
os.environ["OMP_NUM_THREADS"] = "1"
import os.path as osp
import time

import torch
import torch.distributed as dist
import torch.multiprocessing as mp

from tqdm import tqdm
import wandb

import config
from dataset2 import load_data
from models.utils import load_config, load_tokenizer, load_model
from models.ranker import pad_fn
from logger import FileLogger
from utils import *


def main(args):
    if args.seed is not None:
        set_seed(args.seed)

    if args.dist_url == "env://" and args.world_size == -1:
        args.world_size = int(os.environ["WORLD_SIZE"])

    args.distributed = args.world_size > 1 or args.multiprocessing_distributed

    ngpus_per_node = torch.cuda.device_count()
    if args.multiprocessing_distributed:
        # Since we have ngpus_per_node processes per node, the total world_size
        # needs to be adjusted accordingly
        args.world_size = ngpus_per_node * args.world_size
        # Use torch.multiprocessing.spawn to launch distributed processes: the
        # main_worker process function
        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
    else:
        # Simply call main_worker function
        main_worker(args.gpu, ngpus_per_node, args)


def main_worker(gpu, ngpus_per_node, args):
    args.gpu = gpu
    args.n_gpu = ngpus_per_node
    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    is_rank0 = args.rank == 0

    if args.distributed:
        if args.dist_url == "env://" and args.rank == -1:
            args.rank = int(os.environ["RANK"])
        if args.multiprocessing_distributed:
            # For multiprocessing distributed training, rank needs to be the
            # global rank among all the processes
            args.rank = args.rank * ngpus_per_node + gpu
            args.wandb_on = args.wandb_on if args.rank == 0 else False
        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
                                world_size=args.world_size, rank=args.rank)

    is_master = (not args.distributed) or (args.rank == 0)
    args.one2many = True if args.paradigm == "one2seq" else False
    os.makedirs(args.train_output_dir, exist_ok=True)
    os.makedirs(args.cache_dir, exist_ok=True)

    global log
    log = FileLogger(args.train_output_dir, is_master=is_master, is_rank0=is_rank0, log_to_file=args.log_to_file)
    log.console(args)
    if args.wandb_on:
        wandb.init(project="SimCKP-S2", name="/".join(args.train_output_dir.split("/")[1:]))

    trainer = Trainer(args)
    start_time = time.time()
    trainer.train()
    log.console(f"Time for training: {time.time() - start_time:.1f} seconds")


class Trainer:

    def __init__(self, args):
        self.args = args

        ### Load config / tokenizer / model ###
        self.config = load_config(args)
        self.tokenizer = load_tokenizer(args)

        ### Load data ###
        self.train_loader, self.train_sampler = load_data(args, self.config, self.tokenizer, split="train")
        self.valid_loader, _ = load_data(args, self.config, self.tokenizer, split="valid")

        self.model = load_model(args, self.config, self.tokenizer)

        ### Calculate steps ###
        args.total_steps = int(len(self.train_loader) * args.epochs // args.gradient_accumulation_steps)
        args.warmup_steps = int(args.total_steps * args.warmup_ratio)
        log.console(f"warmup steps: {args.warmup_steps}, total steps: {args.total_steps}")

        ### scaler / optimizer / scheduler ###
        self.scaler = init_scaler(args)
        self.optimizer = init_optimizer(args, self.model)
        self.scheduler = init_scheduler(args, self.optimizer)

        self.best_valid_loss = float("inf")
        self.best_valid_f1_at_k = float("-inf")
        self.best_valid_f1_at_m = float("-inf")
        self.start_epoch = 0
        self.tolerance = 0
        self.global_step = 0

        ### Resume training ###
        ckpt_model_path = osp.join(args.train_output_dir, "best_valid_f1_at_m.pt")
        if args.resume and osp.exists(ckpt_model_path):
            log.console(f"Resuming {args.paradigm} model checkpoint from {ckpt_model_path}...")
            ckpt = torch.load(ckpt_model_path)
            self.best_valid_loss = ckpt["loss"]
            self.best_valid_f1_at_k = ckpt["f1_at_k"]
            self.best_valid_f1_at_m = ckpt["f1_at_m"]
            self.start_epoch = ckpt["epoch"]
            self.global_step = ckpt["steps"]
            self.model.load_state_dict(ckpt['model_state_dict'])
            self.optimizer.load_state_dict(ckpt['optimizer_state_dict'])
            self.scheduler.load_state_dict(ckpt['scheduler_state_dict'])
            log.console(f"Validation loss was {ckpt['loss']:.4f}")
            log.console(f"Validation avg theta was {ckpt['theta']:.4f}")
            log.console(f"Validation avg topk was {ckpt['topk']:.4f}")
            log.console(f"Validation F1@5 was {ckpt['f1_at_k']:.4f}")
            log.console(f"Validation F1@M was {ckpt['f1_at_m']:.4f}")
        else:
            log.console(f"Training {args.paradigm} model from scratch")


    def train(self):
        for epoch in range(self.start_epoch, self.args.epochs):
            if self.args.distributed:
                self.train_sampler.set_epoch(epoch)
            avg_train_loss = self.__epoch_train(epoch)
            # avg_valid_loss, _ = self.__epoch_valid()
            if self.tolerance == self.args.max_tolerance: break
            log.console(f"epoch: {epoch+1}, " +
                        f"steps: {self.global_step}, " +
                        f"current lr: {self.optimizer.param_groups[0]['lr']:.8f}, " +
                        f"train loss: {avg_train_loss:.4f}")


    def __epoch_train(self, epoch):
        self.model.train()
        train_loss = 0.
        total = len(self.train_loader)

        with tqdm(desc="Training", total=total, ncols=100, disable=self.args.hide_tqdm) as pbar:
            for step, inputs in enumerate(self.train_loader, 1):
                for k, v in inputs.items():
                    inputs[k] = v.cuda(self.args.gpu, non_blocking=True)

                ### Forward pass ###
                with torch.cuda.amp.autocast(enabled=self.args.use_amp):
                    loss, _, _ = self.model(**inputs)
                    if self.args.gradient_accumulation_steps > 1:
                        loss = loss / self.args.gradient_accumulation_steps

                train_loss += loss.item()

                ### Backward pass ###
                if step % self.args.gradient_accumulation_steps == 0:
                    self.optimizer.zero_grad()
                    self.scaler.scale(loss).backward()
                    self.scaler.unscale_(self.optimizer)
                    if self.args.max_grad_norm > 0:
                        torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.args.max_grad_norm)
                    self.scaler.step(self.optimizer)
                    self.scaler.update()
                    self.scheduler.step()
                    self.global_step += 1
                
                    if self.global_step == 1 or self.global_step % self.args.logging_steps == 0:
                        log.console(f"current lr: {self.optimizer.param_groups[0]['lr']:.8f}, " +
                                    f"steps: {self.global_step}, " +
                                    f"train loss: {(train_loss / (step / self.args.gradient_accumulation_steps)):.4f}")
                        if self.args.wandb_on:
                            wandb.log({"Train Loss": train_loss / (step / self.args.gradient_accumulation_steps)}, step=self.global_step)

                    if self.global_step % self.args.evaluation_steps == 0:
                        avg_valid_loss, valid_score_dict = self.__epoch_valid()
                        if avg_valid_loss < self.best_valid_loss:
                            self.tolerance = 0
                            self.best_valid_loss = avg_valid_loss

                            if not self.args.multiprocessing_distributed or (self.args.multiprocessing_distributed
                                and self.args.rank % self.args.n_gpu == 0):
                                log.console(f"Saving lowest valid loss checkpoint to {self.args.train_output_dir}...")
                                torch.save({'epoch': epoch,
                                            'steps': self.global_step,
                                            'loss': avg_valid_loss,
                                            'theta': valid_score_dict['theta'],
                                            'topk': valid_score_dict['topk'],
                                            'f1_at_k': valid_score_dict['F1@5'],
                                            'f1_at_m': valid_score_dict['F1@M'],
                                            'model_state_dict': self.model.state_dict(),
                                            'optimizer_state_dict': self.optimizer.state_dict(),
                                            'scheduler_state_dict': self.scheduler.state_dict()
                                            }, osp.join(self.args.train_output_dir, "lowest_valid_loss.pt"))
                                
                                with open(osp.join(args.train_output_dir, "best_loss_results.txt"), "w") as f:
                                    f.write(f"Epoch: {epoch}\n" +
                                            f"Total Steps: {self.global_step}\n" +
                                            f"Valid Loss: {avg_valid_loss}\n" +
                                            f"Theta: {valid_score_dict['theta']}\n" +
                                            f"Top K: {valid_score_dict['topk']}\n" +
                                            f"F1@5: {valid_score_dict['F1@5']}\n" +
                                            f"F1@M: {valid_score_dict['F1@M']}")
                        else:
                            self.tolerance += 1
                            log.console(f"Valid loss does not drop, patience: {self.tolerance}/{self.args.max_tolerance}")
                            # self.scheduler.step(avg_valid_loss)

                        if valid_score_dict['F1@M'] > self.best_valid_f1_at_m:
                            self.tolerance = 0
                            self.best_valid_f1_at_m = valid_score_dict['F1@M']
                            if not self.args.multiprocessing_distributed or (self.args.multiprocessing_distributed
                                and self.args.rank % self.args.n_gpu == 0):
                                log.console(f"Saving best valid F1@M checkpoint to {self.args.train_output_dir}...")
                                torch.save({'epoch': epoch,
                                            'steps': self.global_step,
                                            'loss': avg_valid_loss,
                                            'theta': valid_score_dict['theta'],
                                            'topk': valid_score_dict['topk'],
                                            'f1_at_k': valid_score_dict['F1@5'],
                                            'f1_at_m': valid_score_dict['F1@M'],
                                            'model_state_dict': self.model.state_dict(),
                                            'optimizer_state_dict': self.optimizer.state_dict(),
                                            'scheduler_state_dict': self.scheduler.state_dict()
                                            }, osp.join(self.args.train_output_dir, "best_valid_f1_at_m.pt"))

                                with open(osp.join(args.train_output_dir, "best_f1_results.txt"), "w") as f:
                                    f.write(f"Epoch: {epoch}\n" +
                                            f"Total Steps: {self.global_step}\n" +
                                            f"Valid Loss: {avg_valid_loss}\n" +
                                            f"Theta: {valid_score_dict['theta']}\n" +
                                            f"Top K: {valid_score_dict['topk']}\n" +
                                            f"F1@5: {valid_score_dict['F1@5']}\n" +
                                            f"F1@M: {valid_score_dict['F1@M']}")

                        # Switch back to train mode!
                        self.model.train()

                if self.tolerance == self.args.max_tolerance:
                    log.console(f"Has not increased for {self.tolerance} checkpoints, early stop training.")
                    break

                pbar.update(1)
                del loss

        return train_loss / total


    @torch.no_grad()
    def __epoch_valid(self):
        self.model.eval()
        valid_loss = 0.
        valid_logits, valid_labels = [], []
        score_dict = {}
        total = len(self.valid_loader)

        with tqdm(desc="Validating", total=total, ncols=100, disable=self.args.hide_tqdm) as pbar:
            for step, inputs in enumerate(self.valid_loader, 1):
                for k, v in inputs.items():
                    inputs[k] = v.cuda(self.args.gpu, non_blocking=True)

                ### Forward pass ###
                loss, logits, labels = self.model(**inputs)
                valid_logits.append(logits)
                valid_labels.append(labels)

                valid_loss += loss.item()

                pbar.update(1)
                del loss, logits, labels
        
        log.console(f"steps: {self.global_step}, " +
                    f"valid loss: {(valid_loss / total):.4f}, " +
                    f"best valid loss: {self.best_valid_loss:.4f}")

        valid_logits = pad_fn(valid_logits, padding=float("-inf"))
        valid_labels = pad_fn(valid_labels, padding=-100)
        score_dict = self.calculate_scores(valid_logits, valid_labels)
        log.console(f"P@5 ({score_dict['num_matches@5']}/{score_dict['num_preds@5']}): {score_dict['P@5']:.5f}, " +
                    f"R@5 ({score_dict['num_matches@5']}/{score_dict['num_trgs@5']}): {score_dict['R@5']:.5f}, " +
                    f"F1@5: {score_dict['F1@5']:.5f}")
        log.console(f"P@M ({score_dict['num_matches@M']}/{score_dict['num_preds@M']}): {score_dict['P@M']:.5f}, " +
                    f"R@M ({score_dict['num_matches@M']}/{score_dict['num_trgs@M']}): {score_dict['R@M']:.5f}, " +
                    f"F1@M: {score_dict['F1@M']:.5f}")

        if self.args.wandb_on:
            wandb.log({"Valid Loss": valid_loss / total,
                       "P@5": score_dict['P@5'], "R@5": score_dict['R@5'], "F1@5": score_dict['F1@5'],
                       "P@M": score_dict['P@M'], "R@M": score_dict['R@M'], "F1@M": score_dict['F1@M']}, step=self.global_step)

        return valid_loss / total, score_dict


    def calculate_scores(self, logits, labels):
        k = 5
        score_dict = {}

        sorted_logits, sorted_idxes = logits.sort(descending=True)
        preds = torch.zeros_like(sorted_logits).to(sorted_logits)
        preds[labels != -100] = 1.              # assume we predict all values except padding
        num_preds = (preds == 1).cumsum(1)
        num_trgs = (labels == 1).sum(1)
        num_preds_at_k = k
        num_preds_at_m = num_preds
        num_trgs_at_k = torch.minimum(torch.tensor(k, dtype=torch.float).to(num_trgs), num_trgs) if self.args.choi_recall else num_trgs
        num_trgs_at_m = num_trgs.unsqueeze(1).expand((-1, preds.shape[1]))

        sorted_labels = torch.gather(labels, dim=1, index=sorted_idxes)
        num_matches_at_k = sorted_labels[:, :k].sum(1)
        num_matches_at_m = ((preds == 1) * (sorted_labels == 1)).cumsum(1)
        
        # Calculate @k metrics
        precision_at_k = num_matches_at_k / (num_preds_at_k + 1e-20)
        recall_at_k = num_matches_at_k / (num_trgs_at_k + 1e-20)
        f1_at_k = (2 * precision_at_k * recall_at_k / (precision_at_k + recall_at_k + 1e-20))

        # Calculate @M metrics
        precision_at_m = num_matches_at_m / (num_preds_at_m + 1e-20)
        recall_at_m = num_matches_at_m / (num_trgs_at_m + 1e-20)
        f1_at_m = (2 * precision_at_m * recall_at_m / (precision_at_m + recall_at_m + 1e-20))

        f1_at_m[labels == -100] = 0.        # zero out padding values
        best_f1_at_m, best_f1_pos = f1_at_m.max(1)
        precision_at_m = precision_at_m[torch.arange(preds.shape[0]), best_f1_pos]
        recall_at_m = recall_at_m[torch.arange(preds.shape[0]), best_f1_pos]
        num_matches_at_m = num_matches_at_m[torch.arange(preds.shape[0]), best_f1_pos]
        num_preds = num_preds[torch.arange(preds.shape[0]), best_f1_pos]
        thetas = sorted_logits[torch.arange(preds.shape[0]), best_f1_pos]
        topks = best_f1_pos + 1
        f1_at_m = best_f1_at_m

        score_dict[f"P@{k}"] = precision_at_k.mean().item()
        score_dict[f"R@{k}"] = recall_at_k.mean().item()
        score_dict[f"F1@{k}"] = f1_at_k.mean().item()
        score_dict[f"num_matches@{k}"] = num_matches_at_k.sum().long().item()
        score_dict[f"num_preds@{k}"] = k * preds.shape[0]
        score_dict[f"num_trgs@{k}"] = num_trgs_at_k.sum().long().item()
        score_dict[f"P@M"] = precision_at_m.mean().item()
        score_dict[f"R@M"] = recall_at_m.mean().item()
        score_dict[f"F1@M"] = f1_at_m.mean().item()
        score_dict[f"num_matches@M"] = num_matches_at_m.sum().long().item()
        score_dict[f"num_preds@M"] = num_preds.sum().long().item()
        score_dict[f"num_trgs@M"] = num_trgs.sum().long().item()
        score_dict[f"theta"] = thetas.mean().item()
        score_dict[f"topk"] = topks.float().mean().item()

        return score_dict


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="KP Stage 2 Training")
    config.model_args(parser)
    config.data_args(parser)
    config.train_args(parser)
    args = parser.parse_args()
    main(args)