tasks/trainer.py

import datetime
import json
import logging
import os.path as osp
import time

import pandas as pd
import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import wandb

from dataset import MetaLoader, create_dataset, create_loader, create_sampler
from models.vindlu import VindLU
from tasks.retrieval_utils import evaluation_wrapper
from tasks.shared_utils import get_media_types, setup_model
from utils.basic_utils import (MetricLogger, SmoothedValue,
                               remove_files_if_exist, setup_seed)
from utils.config_utils import setup_main
from utils.distributed import get_rank, get_world_size, is_main_process
from utils.logger import log_dict_to_wandb, setup_wandb

logger = logging.getLogger(__name__)


class PretrainTrainer(object):
    """trainer for pretraining."""

    def __init__(self, config):
        super(PretrainTrainer, self).__init__()
        self.config = config

        self.is_pretrain = config.mode == "pt"
        self.setup()

        self.has_decoder = False
        if config.mode in ["ret", "pt"]:
            self.evaluation_fn = evaluation_wrapper
            self.model_cls = VindLU
        elif config.mode == "vqa":
            raise NotImplementedError("not implemented")
        else:
            raise NotImplementedError("not implemented")

        self.build_dataloaders()
        self.build_model()

    def setup(self):
        """setup for train."""
        config = self.config
        if is_main_process() and config.wandb.enable:
            self.wandb_run = setup_wandb(config)
        else:
            self.wandb_run = None
        setup_seed(config.seed + get_rank())
        self.device = torch.device(config.device)

    @torch.no_grad()
    def evaluate(self, epoch=0):
        """evaluate the model.
        Args:
            model (nn.Module): The model to evaluate.
            loader (DataLoader): dataloader.
            tokenizer (None): tokenizer.
            prefix (str): The str prepended to the keys of return dict.

        Returns: dict. The value is the corresponding evaluation results for the key.
        """
        eval_res = {}
        for test_name, test_loader in self.test_name2loaders.items():
            if test_name not in self.config.test_types:
                logger.info(
                    f"Skip eval {test_name} split. All test_types {self.config.test_types}"
                )
                continue
            with torch.cuda.amp.autocast(enabled=self.config.fp16):
                res = self.evaluation_fn(
                    self.model_without_ddp,
                    test_loader,
                    self.tokenizer,
                    self.device,
                    self.config,
                    test_name,
                )
            eval_res.update(res)

        df = pd.DataFrame(eval_res)
        logger.info(f"Epoch {epoch}")
        logger.info(f"\n{df.transpose().to_string(max_cols=30)}")
        return eval_res

    def build_model(self):
        """TODO: Docstring for build_model.
        Returns: TODO

        """
        (
            self.model,
            self.model_without_ddp,
            self.optimizer,
            self.scheduler,
            self.scaler,
            self.tokenizer,
            self.start_epoch,
            self.global_step,
        ) = setup_model(
            self.config,
            model_cls=self.model_cls,
            has_decoder=self.has_decoder,
            pretrain=self.is_pretrain,
            find_unused_parameters=True,
        )

    def build_dataloaders(self):
        config = self.config
        mode = config.mode
        # train datasets, create a list of data loaders
        logger.info(f"Creating dataset for {mode}")
        train_datasets = create_dataset(f"{mode}_train", config)
        media_types = get_media_types(train_datasets)

        if config.distributed:
            num_tasks = get_world_size()
            global_rank = get_rank()
            samplers = create_sampler(
                train_datasets, [True] * len(media_types), num_tasks, global_rank
            )
        else:
            samplers = [None] * len(media_types)

        train_loaders = create_loader(
            train_datasets,
            samplers,
            batch_size=[config.inputs.batch_size[k] for k in media_types],
            num_workers=[config.num_workers] * len(media_types),
            is_trains=[True] * len(media_types),
            collate_fns=[None] * len(media_types),
        )  # [0]

        # test datasets, a mapping from dataset name to data loader
        test_datasets, test_dataset_names = create_dataset(f"{mode}_eval", config)
        test_loaders = create_loader(
            test_datasets,
            [None] * len(test_datasets),
            batch_size=[config.inputs.batch_size_test[d.media_type] for d in test_datasets],
            num_workers=[config.num_workers] * len(test_datasets),
            is_trains=[False] * len(test_datasets),
            collate_fns=[None] * len(test_datasets),
        )
        test_name2loaders = {k: v for k, v in zip(test_dataset_names, test_loaders)}

        self.train_loaders = train_loaders
        self.test_name2loaders = test_name2loaders
        self.media_types = media_types

        num_steps_per_epoch = sum(len(d) for d in self.train_loaders)
        # update config
        config.scheduler.num_training_steps = num_steps_per_epoch * config.scheduler.epochs
        config.scheduler.num_warmup_steps = (
            num_steps_per_epoch * config.scheduler.warmup_epochs
        )
        self.config = config

    def train(self):
        """train the model."""
        config = self.config

        # set cudnn.benchmark=True only when input size is fixed
        # https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936/3
        cudnn.benchmark = len(self.media_types) == 1

        if is_main_process() and config.wandb.enable:
            wandb.watch(self.model)

        best = 0
        best_epoch = 0

        logger.info("Start training")
        start_time = time.time()
        global_step = self.global_step
        for epoch in range(self.start_epoch, config.scheduler.epochs):
            # train one epoch
            global_step = self.train_one_epoch(epoch, global_step)

            # evaluation.
            eval_res = self.evaluate(epoch)

            if is_main_process():

                # log to wandb
                if config.wandb.enable:
                    for p, v in eval_res.items():
                        log_dict_to_wandb(v, step=global_step, prefix=p)

                if config.stop_key is not None and config.stop_key in eval_res:
                    if config.model.multimodal.enable:
                        cur_r_mean = eval_res[config.stop_key]["r_mean"]
                    else:
                        cur_r_mean = eval_res[config.stop_key.replace("/", "_emb/")]["r_mean"]
                else:  # None
                    cur_r_mean = best + 1  # save the last as the best

                with open(osp.join(config.output_dir, "eval_res_latest.json"), "w") as f:
                    json.dump(eval_res, f)
                # eval_res.to_json(osp.join(config.output_dir, "eval_res_latest.json"))

                save_obj = {
                    "model": self.model_without_ddp.state_dict(),
                    "optimizer": self.optimizer.state_dict(),
                    "scheduler": self.scheduler.state_dict(),
                    "scaler": self.scaler.state_dict(),
                    "config": config,
                    "epoch": epoch,
                    "global_step": global_step,
                }
                torch.save(save_obj, osp.join(config.output_dir, f"ckpt_{epoch:02d}.pth"))

                if cur_r_mean > best:
                    torch.save(save_obj, osp.join(config.output_dir, "ckpt_best.pth"))
                    eval_file = "eval_res_best.json"
                    # eval_res.to_json(osp.join(config.output_dir, eval_file))
                    with open(osp.join(config.output_dir, eval_file), "w") as f:
                        json.dump(eval_res, f)
                    best = cur_r_mean
                    best_epoch = epoch

            dist.barrier()

        total_time = time.time() - start_time
        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
        logger.info(f"Training time {total_time_str}")
        logger.info(f"best epoch {best_epoch} [config.stop_key {config.stop_key}]")
        logger.info(f"Checkpoints and Logs saved at {config.output_dir}")

        if is_main_process() and config.wandb.enable:
            self.wandb_run.finish()

    def train_one_epoch(self, epoch, global_step):
        config = self.config
        self.model.train()

        metric_logger = MetricLogger(delimiter="  ")
        metric_logger.add_meter("lr", SmoothedValue(window=100, fmt="{value:.6f}"))
        metric_logger.add_meter("temperature", SmoothedValue(window=100, fmt="{value:.4f}"))
        loss_names = ["loss_" + k for k, v in config.criterion.loss_weight.items() if v != 0]

        media_types = get_media_types(self.train_loaders)

        for name in loss_names:
            for m in media_types:
                metric_logger.add_meter(
                    f"{m}-{name}", SmoothedValue(window=100, fmt="{value:.4f}")
                )

        header = f"Train Epoch: [{epoch}]"
        log_freq = config.log_freq

        if config.distributed:
            for d in self.train_loaders:
                d.sampler.set_epoch(epoch)
        train_loader = MetaLoader(name2loader=dict(list(zip(media_types, self.train_loaders))))

        model_without_ddp = self.model.module if config.distributed else self.model
        iterator = metric_logger.log_every(train_loader, log_freq, header)
        for i, (media_type, (image, text, idx)) in enumerate(iterator):
            image = image.to(self.device, non_blocking=True)
            idx = idx.to(self.device, non_blocking=True)
            text_input = self.tokenizer(
                text,
                padding="max_length",
                truncation=True,
                max_length=config.inputs.max_txt_l[media_type],
                return_tensors="pt",
            ).to(self.device)

            with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.bfloat16):
                loss_dict = self.model(image, text_input, idx=idx)
                loss = sum(loss_dict.values())

            self.optimizer.zero_grad()
            self.scaler.scale(loss).backward()
            if config.optimizer.max_grad_norm > 0:
                self.scaler.unscale_(self.optimizer)
                torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), config.optimizer.max_grad_norm
                )
            self.scaler.step(self.optimizer)
            self.scaler.update()
            self.scheduler.step()

            # logging
            for name in loss_names:
                value = loss_dict[name]
                value = value if isinstance(value, float) else value.item()
                metric_logger.update(**{f"{media_type}-{name}": value})
            metric_logger.update(lr=self.optimizer.param_groups[0]["lr"])
            metric_logger.update(temperature=model_without_ddp.temp.item())

            if is_main_process() and config.wandb.enable and global_step % log_freq == 0:
                logs = metric_logger.get_global_avg_dict()
                log_dict_to_wandb(logs, step=global_step, prefix="train/")

            global_step += 1

            if config.debug and global_step % 2 == 0:
                logger.info("debug mode, break training loop")
                break

        # gather the stats from all processes
        metric_logger.synchronize_between_processes()
        logger.info(f"Averaged stats: {metric_logger.global_avg()}")
        return global_step


if __name__ == "__main__":
    cfg = setup_main()
    trainer = PretrainTrainer(cfg)
    if cfg.evaluate:
        trainer.evaluate()
    else:
        trainer.train()