train_it_long.py

import datetime
import logging
import time
from os.path import join

import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import numpy as np
import wandb

from dataset import MetaLoader, create_dataset, create_loader, create_sampler
from models.videochat2_it_long import VideoChat2_it_Long
from tasks.shared_utils import get_media_types, setup_model
from utils.basic_utils import (MetricLogger, SmoothedValue, 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
from mv_eval import evaluate

logger = logging.getLogger(__name__)


def train(
    model,
    train_loaders,
    optimizer,
    epoch,
    global_step,
    device,
    scheduler,
    scaler,
    config,
):
    model.train()

    best_avg = 0.0

    metric_logger = MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", SmoothedValue(window=100, fmt="{value:.7f}"))
    loss_names = ["loss"]

    media_types = get_media_types(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
    eval_freq = config.eval_freq

    logger.info(f"Start training epoch {epoch}")
    # log eval_freq
    logger.info(f'eval_freq: {eval_freq}')

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

    iterator = metric_logger.log_every(train_loader, log_freq, header)
    for i, (media_type, (image, text, instruction, _, frames_times)) in enumerate(iterator):

        model.train()

        image = image.to(device, non_blocking=True)

        with torch.cuda.amp.autocast(enabled=config.fp16):
            loss_dict = model(image, text, instruction, frames_times)
            loss = sum(loss_dict.values())

        optimizer.zero_grad()
        scaler.scale(loss).backward()
        if config.optimizer.max_grad_norm > 0:
            scaler.unscale_(optimizer)
            torch.nn.utils.clip_grad_norm_(model.parameters(), config.optimizer.max_grad_norm)
        scaler.step(optimizer)
        scaler.update()
        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=optimizer.param_groups[0]["lr"])

        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/")

        if is_main_process() and (eval_freq > 0 and global_step % eval_freq == 0):
            model.eval()
            with torch.no_grad():
                avgs = []
                for dataset in ['Action Sequence', 'Object Interaction', 'State Change']:
                    avg = evaluate(model.module, 
                            [dataset],
                            base_frame_num=config.base_frame_num,
                            num_frm=-1, 
                            resolution=224,
                            target_video_length=-1,
                            max_num_frm=config.num_frames,
                            output_file=f'{config.output_dir}/eval_{global_step:06d}.json',
                            device=device,
                            )
                    avgs.append(avg)
                logger.info(f"Global step: {global_step} Evaluation: {np.mean(avgs)}, AS, OI, SC: {avgs}")
                # log into metric_logger
                metric_logger.update(eval=avg)
            
            if global_step > 0:

                logger.info(f"Epoch {epoch}")
                param_grad_dic = {
                    k: v.requires_grad for (k, v) in model.module.named_parameters()
                }
                state_dict = model.module.state_dict()
                for k in list(state_dict.keys()):
                    if k in param_grad_dic.keys() and not param_grad_dic[k]:
                        # delete parameters that do not require gradient
                        del state_dict[k]
                save_obj = {
                    "model": state_dict,
                    "optimizer": optimizer.state_dict(),
                    "scheduler": scheduler.state_dict(),
                    "scaler": scaler.state_dict(),
                    "config": config,
                    "epoch": epoch,
                    "global_step": global_step,
                }

                torch.save(save_obj, join(config.output_dir, f"ckpt_step_{global_step:05d}.pth"))
            
        global_step += 1

        if global_step % 100 == 0:
            print("global_step:", global_step)

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

        if config.debug and global_step % (2 * log_freq + 3) == 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


def setup_dataloaders(config, mode="pt"):
    # 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]

    return train_loaders, media_types


def main(config):

    if is_main_process() and config.wandb.enable:
        run = setup_wandb(config)

    logger.info(f"train_file: {config.train_file}")


    # # debug:
    # linear = torch.nn.Linear([16,10]).cuda()
    # x1 = torch.ones([1, 16]).cuda()
    # x2 = torch.ones([1, 16]).cuda()

    # with torch.no_grad():
    #     linear(x1)
    # print(linear(x2).grad_fn)

    # import ipdb; ipdb.set_trace()

    setup_seed(config.seed + get_rank())
    device = torch.device(config.device)

    train_loaders, train_media_types = setup_dataloaders(
        config, mode=config.mode
    )
    num_steps_per_epoch = sum(len(d) for d in train_loaders)
    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
    # 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(train_media_types) == 1

    model_cls = eval(config.model.get('model_cls', 'VideoChat2_pt'))
    (
        model,
        model_without_ddp,
        optimizer,
        scheduler,
        scaler,
        start_epoch,
        global_step,
    ) = setup_model(
        config,
        model_cls=model_cls,
        find_unused_parameters=True,
    )
    if is_main_process() and config.wandb.enable:
        wandb.watch(model)

    logger.info("Start training")
    start_time = time.time()

    print("Starting Epoch:", start_epoch)

    for epoch in range(start_epoch, config.scheduler.epochs):
        if not config.evaluate:
            global_step = train(
                model,
                train_loaders,
                optimizer,
                epoch,
                global_step,
                device,
                scheduler,
                scaler,
                config,
            )

        if is_main_process():

            logger.info(f"Epoch {epoch}")
            param_grad_dic = {
                k: v.requires_grad for (k, v) in model_without_ddp.named_parameters()
            }
            state_dict = model_without_ddp.state_dict()
            for k in list(state_dict.keys()):
                if k in param_grad_dic.keys() and not param_grad_dic[k]:
                    # delete parameters that do not require gradient
                    del state_dict[k]
            save_obj = {
                "model": state_dict,
                "optimizer": optimizer.state_dict(),
                "scheduler": scheduler.state_dict(),
                "scaler": scaler.state_dict(),
                "config": config,
                "epoch": epoch,
                "global_step": global_step,
            }

            torch.save(save_obj, join(config.output_dir, f"ckpt_{epoch:02d}.pth"))

            if config.get("save_latest", False):
                torch.save(save_obj, join(config.output_dir, "ckpt_latest.pth"))

        if config.evaluate:
            break

        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"Checkpoints and Logs saved at {config.output_dir}")

    if is_main_process() and config.wandb.enable:
        run.finish()


if __name__ == "__main__":
    cfg = setup_main()
    main(cfg)