tasks/vqa.py

import copy
import datetime
import logging
import os
import time
from os.path import join
from shutil import copyfile

import torch
import torch.backends.cudnn as cudnn
import torch.distributed as dist
import wandb
from omegaconf import OmegaConf

from dataset import (create_dataset, create_loader, create_sampler,
                     vqa_collate_fn)
from dataset.utils import sync_save_result
from models.vindlu_qa import VindLU_QA
from models.vindlu_vit_qa import VindLU_VIT_QA
from models.vindlu_blip_T5 import VindLU_BLIP_T5
from tasks.shared_utils import setup_model
from tasks.vqa_utils import eval_qa_acc
from utils.basic_utils import (MetricLogger, SmoothedValue, save_json,
                               setup_seed)
from utils.config import Config
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__)


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

    metric_logger = MetricLogger(delimiter="  ")
    metric_logger.add_meter("lr", SmoothedValue(window=1, fmt="{value:.6f}"))
    metric_logger.add_meter("loss", SmoothedValue(window=1, fmt="{value:.4f}"))

    header = "Train Epoch: [{}]".format(epoch)
    log_freq = config.log_freq

    if config.distributed:
        train_loader.sampler.set_epoch(epoch)

    iterator = metric_logger.log_every(train_loader, log_freq, header)
    for i, data in enumerate(iterator):
        image, question, answer, weights, n = data
        image = image.to(device, non_blocking=True)
        weights = weights.to(device, non_blocking=True)
        question_input = tokenizer(
            question,
            padding="max_length",
            truncation=True,
            max_length=config.max_q_len,
            return_tensors="pt",
        ).to(device)
        answer_input = tokenizer(
            answer,
            padding="max_length",
            truncation=True,
            max_length=config.max_a_len,
            return_tensors="pt",
        ).to(device)

        with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.bfloat16):
            loss = model(
                image, question_input, answer_input, train=True, k=n, weights=weights,
                raw_question=list(question), raw_answer=list(answer),
            )

        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
        metric_logger.update(loss=loss.item())
        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/")

        global_step += 1

        if config.debug and (i + 1) % 5 == 0:
            break

    metric_logger.synchronize_between_processes()
    logger.info(f"Averaged train stats: {metric_logger.global_avg()}")
    return global_step


@torch.no_grad()
def evaluation(model, data_loader, tokenizer, device, config):
    model.eval()

    metric_logger = MetricLogger(delimiter="  ")
    header = "[evaluation] Generating answers:"
    log_freq = config.log_freq // 2

    result = []

    raw_answer_list = data_loader.dataset.answer_list
    answer_list = [a + " " + config.eos for a in raw_answer_list]
    answer_input = tokenizer(
        answer_list,
        padding="max_length",
        truncation=True,
        max_length=config.max_a_len,
        return_tensors="pt",
    ).to(device)

    logger.info("Start generating results.")

    iterator = metric_logger.log_every(data_loader, log_freq, header)
    for n, data in enumerate(iterator):
        image, question, question_id = data
        image = image.to(device, non_blocking=True)
        question_input = tokenizer(
            question,
            padding="max_length",
            truncation=True,
            max_length=config.max_q_len,
            return_tensors="pt",
        ).to(device)

        topk_ids, topk_probs = model(
            image, question_input, answer_input, train=False, k=config.evaluation.k_test,
            raw_question=list(question), raw_answer=raw_answer_list,
        )

        for ques_id, topk_id, topk_prob in zip(question_id, topk_ids, topk_probs):
            ques_id = int(ques_id.item()) if not isinstance(ques_id, str) else ques_id
            _, pred = topk_prob.max(dim=0)
            result.append({"question_id": ques_id, "answer": raw_answer_list[topk_id[pred]]})

    return result


def setup_dataloaders(config, dataset_name="vqa"):
    logger.info(f"Creating {dataset_name} QA datasets")
    train_dataset = create_dataset("qa_train", config)
    test_datasets, test_dataset_names = create_dataset("qa_eval", config)
    all_datasets = [train_dataset] + test_datasets

    if config.distributed:
        num_tasks = get_world_size()
        global_rank = get_rank()
        samplers = create_sampler(
            all_datasets,
            shuffles=[True] + [False] * len(test_datasets),
            num_tasks=num_tasks,
            global_rank=global_rank,
        )
    else:
        samplers = [None] * len(all_datasets)

    train_bsz = config.inputs.batch_size[train_dataset.media_type]
    test_bsz_list = [config.inputs.batch_size_test[d.media_type] for d in test_datasets]
    loaders = create_loader(
        all_datasets,
        samplers,
        batch_size=[train_bsz] + test_bsz_list,
        num_workers=[
            config.num_workers,
        ]
        * len(all_datasets),
        is_trains=[True] + [False] * len(test_datasets),
        collate_fns=[vqa_collate_fn] + [None] * len(test_datasets),
    )
    train_loader, test_loaders = loaders[0], loaders[1:]
    test_name2loaders = {k: v for k, v in zip(test_dataset_names, test_loaders)}
    return train_loader, test_name2loaders


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

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

    setup_seed(config.seed + get_rank())
    device = torch.device(config.device)
    cudnn.benchmark = True

    train_loader, test_name2loaders = setup_dataloaders(config)
    config.scheduler.num_training_steps = len(train_loader) * config.scheduler.epochs
    config.scheduler.num_warmup_steps = len(train_loader) * config.scheduler.warmup_epochs

    model_cls = eval(config.model.get('model_cls', 'VindLU_QA'))
    (
        model,
        model_without_ddp,
        optimizer,
        scheduler,
        scaler,
        tokenizer,
        start_epoch,
        global_step,
    ) = setup_model(
        config,
        model_cls=model_cls,
        has_decoder=True,
        pretrain=False,
        # find_unused_parameters=True,
        find_unused_parameters=False,
    )
    if is_main_process() and config.wandb.enable:
        wandb.watch(model)

    best = 0
    best_epoch = 0
    has_gt = config.dataset_name != "vqa" or config.test_types[0] == "minival"

    logger.info("Start " + "evaluation" if config.evaluate else "training")
    start_time = time.time()
    for epoch in range(start_epoch, config.scheduler.epochs):
        if not config.evaluate:
            global_step = train(
                model,
                train_loader,
                optimizer,
                tokenizer,
                epoch,
                global_step,
                device,
                scheduler,
                scaler,
                config,
            )

        if config.get('no_test', False) and not config.evaluate:
            if is_main_process():
                save_obj = {
                    "model": model_without_ddp.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, "ckpt_best.pth"))
                best_epoch = epoch
                
        else:
            with torch.cuda.amp.autocast(enabled=config.fp16, dtype=torch.bfloat16):
                eval_res = {}
                pred_name2file = {}
                for test_name, test_loader in test_name2loaders.items():
                    if test_name not in config.test_types:
                        logger.info(
                            f"Skip eval {test_name} split. All test_types {config.test_types}"
                        )
                        continue
                    logger.info(f"Evaluating {test_name} split...")
                    qa_result = evaluation(model, test_loader, tokenizer, device, config)

                    # sync/gather results and eval
                    pred_file, gathered_result = sync_save_result(
                        qa_result, config.result_dir, f"{test_name}_latest"
                    )
                    pred_name2file[test_name] = pred_file
                    if is_main_process() and has_gt:
                        eval_res[test_name] = eval_qa_acc(
                            test_loader.dataset.anno_list,
                            gathered_result,
                            is_vqa=config.dataset_name == "vqa",
                        )

            if is_main_process():
                if len(eval_res) > 0:
                    logger.info(f"eval_res {eval_res}")

                    if config.wandb.enable:
                        for name, acc_dict in eval_res.items():
                            log_dict_to_wandb(acc_dict, step=global_step, prefix=f"{name}/")

                if not config.evaluate and has_gt and eval_res[config.stop_key]["overall"] > best:
                    save_obj = {
                        "model": model_without_ddp.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "scheduler": scheduler.state_dict(),
                        "scaler": scaler.state_dict(),
                        "config": config,
                        "epoch": epoch,
                        "global_step": global_step,
                    }
                    for name, pred_file in pred_name2file.items():
                        copyfile(pred_file, pred_file.replace("latest", "best"))
                    save_json(
                        eval_res, join(config.output_dir, "eval_res_best.json"), save_pretty=True
                    )
                    torch.save(save_obj, join(config.output_dir, "ckpt_best.pth"))
                    best = eval_res[config.stop_key]["overall"]
                    best_epoch = epoch
                if config.evaluate:
                    save_json(
                        eval_res, join(config.output_dir, "eval_res_best.json"), save_pretty=True
                    )

                    for name, pred_file in pred_name2file.items():
                        copyfile(pred_file, pred_file.replace("latest", "eval"))
                        if has_gt:
                            save_path = join(config.output_dir, f"{name}_acc_eval.json")
                            save_json(eval_res[name], save_path, save_pretty=True)

        if config.evaluate or config.debug:
            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"best epoch {best_epoch}")
    logger.info(f"Checkpoints and Logs saved at {config.output_dir}")

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


def eval_after_training(train_config):
    # general config for all
    train_config.wandb.enable = False
    train_config.evaluate = True
    train_config.pretrained_path = join(train_config.output_dir, "ckpt_best.pth")
    
    if train_config.get('num_frames_test_final', False):
        train_config.num_frames_test = train_config.num_frames_test_final
        train_config.batch_size = train_config.batch_size_final
        train_config.inputs.video_input.num_frames_test = train_config.num_frames_test_final
        train_config.model.vision_encoder.num_frames = train_config.num_frames_test_final

    eval_config = copy.deepcopy(train_config)
    eval_config.test_types = (
        ["val", "test"] if eval_config.dataset_name != "vqa" else ["minival"]
    )
    eval_config.output_dir = join(eval_config.output_dir, f"eval_after_training")
    eval_config.result_dir = eval_config.output_dir
    if is_main_process():
        os.makedirs(eval_config.output_dir, exist_ok=False)
        Config.dump(eval_config, os.path.join(eval_config.output_dir, "config.json"))
    logger.info(f"===========> START eval_after_training [{eval_config.test_types}]")
    main(eval_config)


if __name__ == "__main__":
    cfg = setup_main()
    cfg.result_dir = cfg.output_dir
    main(cfg)
    if not cfg.evaluate:
        eval_after_training(cfg)