main_finetune_classification.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import argparse
from collections import OrderedDict
import json
import math
import numpy as np
import os
import pandas as pd
import sys
import time

import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.cuda.amp as amp
from torch.distributed.optim import ZeroRedundancyOptimizer
import torch.nn.parallel
import torchvision.transforms as transforms
import torchvision.transforms._transforms_video as transforms_video
from sklearn.metrics import confusion_matrix
import wandb

from lavila.data import datasets
from lavila.data.video_transforms import Permute, SpatialCrop, TemporalCrop
from lavila.models import models
from lavila.models.tokenizer import (MyBertTokenizer, MyDistilBertTokenizer, MyGPT2Tokenizer, SimpleTokenizer)
from lavila.models.utils import inflate_positional_embeds
from lavila.utils import distributed as dist_utils
from lavila.utils.evaluation import accuracy, get_mean_accuracy
from lavila.utils.meter import AverageMeter, ProgressMeter
from lavila.utils.preprocess import generate_label_map
from lavila.utils.random import random_seed
from lavila.utils.scheduler import cosine_scheduler
from lavila.utils.evaluation_ek100cls import get_marginal_indexes, marginalize


def get_args_parser():
    parser = argparse.ArgumentParser(description='lavila finetune and evaluation', add_help=False)
    # Data
    parser.add_argument('--dataset', default='ek100_cls', type=str,
                        choices=['ek100_cls', 'egtea'])
    parser.add_argument('--root',
                        default='datasets/EK100/video_ht256px/',
                        type=str, help='path to dataset root')
    parser.add_argument('--metadata-train',
                        default='datasets/EK100/epic-kitchens-100-annotations/EPIC_100_train.csv',
                        type=str, help='path to metadata file (train set)')
    parser.add_argument('--metadata-val',
                        default='datasets/EK100/epic-kitchens-100-annotations/EPIC_100_validation.csv',
                        type=str, help='path to metadata file (val set)')
    parser.add_argument('--relevancy-path',
                        default='datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl',
                        type=str, help='path to relevancy matrix (val set)')
    parser.add_argument('--output-dir', default='./', type=str, help='output dir')
    parser.add_argument('--num-crops', default=1, type=int, help='number of crops in transforms for val')
    parser.add_argument('--num-clips', default=1, type=int, help='number of clips for val')
    parser.add_argument('--clip-length', default=16, type=int, help='clip length')
    parser.add_argument('--clip-stride', default=2, type=int, help='clip stride')
    parser.add_argument('--sparse-sample', action='store_true', help='switch to sparse sampling')
    # Model
    parser.add_argument('--pretrain-model', default='', type=str, help='path to pretrain model')
    parser.add_argument('--resume', default='', type=str, help='path to resume from')
    parser.add_argument('--find-unused-parameters', action='store_true',
                        help='do this during DDP (useful for models with tied weights)')
    parser.add_argument('--drop-path-rate', default=0.1, type=float, help='drop path ratio')
    parser.add_argument('--dropout-ratio', default=0.5, type=float, help='dropout ratio for the last linear layer')
    parser.add_argument('--num-classes', default=3806, nargs='+', type=int, help='number of classes for the last linear layer')
    parser.add_argument('--use-vn-classifier', action='store_true')
    parser.add_argument('--use-half', action='store_true', help='use half precision at inference')
    # Training
    parser.add_argument('--epochs', default=100, type=int)
    parser.add_argument('--warmup-epochs', default=1, type=int)
    parser.add_argument('--start-epoch', default=0, type=int)
    parser.add_argument('--batch-size', default=16, type=int,
                        help='number of samples per-device/per-gpu')
    parser.add_argument('--use-sgd', action='store_true')
    parser.add_argument('--freeze-temperature', action='store_true', help='freeze temperature if set to True')
    parser.add_argument('--lr', default=3e-3, type=float)
    parser.add_argument('--fix-lr', action='store_true', help='disable cosine lr decay if set True')
    parser.add_argument('--lr-start', default=1e-6, type=float,
                        help='initial warmup lr')
    parser.add_argument('--lr-end', default=1e-5, type=float,
                        help='minimum final lr')
    parser.add_argument('--lr-multiplier-on-backbone', default=0.1, type=float, help='lr multiplier for the backbone')
    parser.add_argument('--clip-grad-type', default='norm', choices=['norm', 'value'])
    parser.add_argument('--clip-grad-value', default=None, type=float, help='')
    parser.add_argument('--update-freq', default=1, type=int,
                        help='optimizer update frequency (i.e. gradient accumulation steps)')
    parser.add_argument('--wd', default=0.01, type=float)
    parser.add_argument('--betas', default=(0.9, 0.999), nargs=2, type=float)
    parser.add_argument('--eps', default=1e-8, type=float)
    parser.add_argument('--label-smoothing', default=0.1, type=float, help='label smoothing')
    parser.add_argument('--eval-freq', default=5, type=int)
    parser.add_argument('--save-freq', default=5, type=int)
    parser.add_argument('--disable-amp', action='store_true',
                        help='disable mixed-precision training (requires more memory and compute)')
    parser.add_argument('--use-zero', action='store_true',
                        help='use ZeroRedundancyOptimizer to save memory')
    parser.add_argument('--use-checkpoint', action='store_true',
                        help='use gradient checkpointing during training for significantly less GPU usage')
    # System
    parser.add_argument('--print-freq', default=100, type=int, help='print frequency')
    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                        help='number of data loading workers per process')
    parser.add_argument('--evaluate', action='store_true', help='eval only')
    parser.add_argument('--world-size', default=1, type=int,
                        help='number of nodes for distributed training')
    parser.add_argument('--rank', default=0, type=int,
                        help='node rank for distributed training')
    parser.add_argument("--local_rank", type=int, default=0)
    parser.add_argument('--dist-url', default='env://', type=str,
                        help='url used to set up distributed training')
    parser.add_argument('--dist-backend', default='nccl', type=str)
    parser.add_argument('--seed', default=0, type=int)
    parser.add_argument('--gpu', default=None, type=int, help='GPU id to use.')
    parser.add_argument('--wandb', action='store_true', help='Enable WandB logging')
    return parser


def main(args):
    dist_utils.init_distributed_mode(args)

    global best_acc1
    random_seed(args.seed, dist_utils.get_rank())

    if args.pretrain_model:
        ckpt_path = args.pretrain_model
    else:
        raise Exception('no checkpoint found')
    ckpt = torch.load(ckpt_path, map_location='cpu')

    if args.use_vn_classifier:
        assert args.dataset == 'ek100_cls' and len(args.num_classes) == 3

    state_dict = OrderedDict()
    for k, v in ckpt['state_dict'].items():
        state_dict[k.replace('module.', '')] = v

    old_args = ckpt['args']
    print("=> creating model: {}".format(old_args.model))
    model = getattr(models, old_args.model)(
        pretrained=old_args.load_visual_pretrained,
        pretrained2d=old_args.load_visual_pretrained is not None,
        text_use_cls_token=old_args.use_cls_token,
        project_embed_dim=old_args.project_embed_dim,
        timesformer_gated_xattn=False,
        timesformer_freeze_space=False,
        num_frames=args.clip_length,
        drop_path_rate=args.drop_path_rate,
    )
    if 'TIMESFORMER' in old_args.model or 'EGOVLP' in old_args.model:
        # inflate weight
        print('=> inflating PE in models due to different frame numbers')
        state_dict = inflate_positional_embeds(
            model.state_dict(), state_dict,
            num_frames=args.clip_length,
            load_temporal_fix='bilinear',
        )
    model.load_state_dict(state_dict, strict=True)
    print("=> loaded resume checkpoint '{}' (epoch {})".format(ckpt_path, ckpt['epoch']))

    if args.use_vn_classifier:
        model = models.VideoClassifierMultiHead(
            model.visual,
            dropout=args.dropout_ratio,
            num_classes_list=args.num_classes
        )
    else:
        assert len(args.num_classes) == 1
        model = models.VideoClassifier(
            model.visual,
            dropout=args.dropout_ratio,
            num_classes=args.num_classes[0]
        )

    model.cuda(args.gpu)

    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model, device_ids=[args.gpu], bucket_cap_mb=200,
            find_unused_parameters=args.find_unused_parameters
        )

    p_wd, p_non_wd = [], []
    p_head_wd, p_head_non_wd = [], []
    for n, p in model.named_parameters():
        if 'fc_cls' in n:
            if 'bias' in n:
                p_head_non_wd.append(p)
            else:
                p_head_wd.append(p)
        elif not p.requires_grad:
            continue  # frozen weights
        elif p.ndim < 2 or 'bias' in n or 'ln' in n or 'bn' in n:
            p_non_wd.append(p)
        else:
            p_wd.append(p)

    optim_params = [
        {"params": p_wd, "weight_decay": args.wd,  "lr": args.lr * args.lr_multiplier_on_backbone},
        {"params": p_non_wd, "weight_decay": 0, "lr": args.lr * args.lr_multiplier_on_backbone},
        {"params": p_head_wd, "weight_decay": args.wd},
        {"params": p_head_non_wd, "weight_decay": 0}
    ]

    if args.use_zero:
        optimizer = ZeroRedundancyOptimizer(
            optim_params, optimizer_class=torch.optim.SGD if args.use_sgd else torch.optim.AdamW,
            lr=args.lr, betas=args.betas, eps=args.eps, weight_decay=args.wd
        )
    else:
        if args.use_sgd:
            optimizer = torch.optim.SGD(optim_params, lr=args.lr, momentum=args.betas[0], weight_decay=args.wd)
        else:
            optimizer = torch.optim.AdamW(optim_params, lr=args.lr, betas=args.betas,
                                          eps=args.eps, weight_decay=args.wd)
    scaler = amp.GradScaler(enabled=not args.disable_amp)
    # optionally resume from a checkpoint (takes precedence over autoresume)
    latest = os.path.join(args.output_dir, 'checkpoint.pt')
    if os.path.isfile(latest):
        args.resume = ''
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading resume checkpoint '{}'".format(args.resume))
            checkpoint = torch.load(args.resume, map_location='cpu')
            epoch = checkpoint['epoch'] if 'epoch' in checkpoint else 0
            args.start_epoch = epoch
            if not args.distributed:
                state_dict = OrderedDict()
                for k, v in checkpoint['state_dict'].items():
                    state_dict[k.replace('module.', '')] = v
                result = model.load_state_dict(state_dict, strict=False)
            else:
                result = model.load_state_dict(checkpoint['state_dict'], strict=False)
            print(result)
            optimizer.load_state_dict(checkpoint['optimizer']) if 'optimizer' in checkpoint else ()
            scaler.load_state_dict(checkpoint['scaler']) if 'scaler' in checkpoint else ()
            best_acc1 = checkpoint['best_acc1']
            print("=> loaded resume checkpoint '{}' (epoch {}, best_metric = {})"
                  .format(args.resume, epoch, best_acc1))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
    else:
        # auto-resume from latest checkpoint in output directory
        latest = os.path.join(args.output_dir, 'checkpoint.pt')
        if os.path.isfile(latest):
            print("=> loading latest checkpoint '{}'".format(latest))
            latest_checkpoint = torch.load(latest, map_location='cpu')
            args.start_epoch = latest_checkpoint['epoch']
            model.load_state_dict(latest_checkpoint['state_dict'])
            optimizer.load_state_dict(latest_checkpoint['optimizer'])
            scaler.load_state_dict(latest_checkpoint['scaler'])
            best_acc1 = latest_checkpoint['best_acc1']
            print("=> loaded latest checkpoint '{}' (epoch {})"
                  .format(latest, latest_checkpoint['epoch']))

    cudnn.benchmark = True

    # Data loading code
    print("=> creating dataset")
    if old_args.model.endswith('DISTILBERT_BASE'):
        tokenizer = MyDistilBertTokenizer('distilbert-base-uncased')
    elif old_args.model.endswith('BERT_BASE'):
        tokenizer = MyBertTokenizer('bert-base-uncased')
    elif old_args.model.endswith('BERT_LARGE'):
        tokenizer = MyBertTokenizer('bert-large-uncased')
    elif old_args.model.endswith('GPT2'):
        tokenizer = MyGPT2Tokenizer('gpt2')
    elif old_args.model.endswith('GPT2_MEDIUM'):
        tokenizer = MyGPT2Tokenizer('gpt2-medium')
    elif old_args.model.endswith('GPT2_LARGE'):
        tokenizer = MyGPT2Tokenizer('gpt2-large')
    elif old_args.model.endswith('GPT2_XL'):
        tokenizer = MyGPT2Tokenizer('gpt2-xl')
    else:
        print("Using SimpleTokenizer because of model '{}'. "
              "Please check if this is what you want".format(old_args.model))
        tokenizer = SimpleTokenizer()

    criterion = nn.CrossEntropyLoss(label_smoothing=args.label_smoothing).cuda(args.gpu)

    crop_size = 224 if '336PX' not in old_args.model else 336
    transforms_list = [
        Permute([3, 0, 1, 2]),    # T H W C -> C T H W
        transforms.RandomResizedCrop(crop_size, scale=(0.5, 1.0)),
        transforms.RandomHorizontalFlip(p=0.5),
    ]
    if 'OPENAI' in old_args.model:
        transforms_list.append(transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305]))
    else:
        transforms_list.append(transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]))
    train_transform = transforms.Compose(transforms_list)

    val_transform = transforms.Compose([
            Permute([3, 0, 1, 2]),    # T H W C -> C T H W
            transforms.Resize(crop_size),
            transforms.CenterCrop(crop_size),
            (transforms_video.NormalizeVideo(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) if 'OPENAI' not in old_args.model else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305])),
            TemporalCrop(frames_per_clip=args.clip_length, stride=args.clip_length),
            SpatialCrop(crop_size=crop_size, num_crops=args.num_crops),
        ])

    # build dataset
    _, mapping_vn2act = generate_label_map(args.dataset)
    if args.dataset == 'ek100_cls':
        args.mapping_act2v = {i: int(vn.split(':')[0]) for (vn, i) in mapping_vn2act.items()}
        args.mapping_act2n = {i: int(vn.split(':')[1]) for (vn, i) in mapping_vn2act.items()}
        args.actions = pd.DataFrame.from_dict({'verb': args.mapping_act2v.values(), 'noun': args.mapping_act2n.values()})
    num_clips_at_val = args.num_clips
    args.num_clips = 1
    train_dataset = datasets.get_downstream_dataset(
        train_transform, tokenizer, args, subset='train', label_mapping=mapping_vn2act,
    )
    args.num_clips = num_clips_at_val
    val_dataset = datasets.get_downstream_dataset(
        val_transform, tokenizer, args, subset='val', label_mapping=mapping_vn2act,
    )

    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
        val_sampler = torch.utils.data.SequentialSampler(val_dataset)  # disable distributed
    else:
        train_sampler = None
        val_sampler = None

    train_loader = torch.utils.data.DataLoader(
        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=train_sampler, drop_last=True
    )
    print('len(train_loader) = {}'.format(len(train_loader)))
    val_loader = torch.utils.data.DataLoader(
        val_dataset, batch_size=args.batch_size, shuffle=(val_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=val_sampler, drop_last=False
    )
    print('len(val_loader) = {}'.format(len(val_loader)))

    if args.evaluate:
        if args.use_vn_classifier:
            val_stats = validate_multihead(val_loader, model, args)
        else:
            val_stats = validate(val_loader, model, args)
        return

    if args.fix_lr:
        lr_schedule = None
    else:
        lr_schedule = cosine_scheduler(
            args.lr, args.lr_end, args.epochs, len(train_loader) // args.update_freq,
            warmup_epochs=args.warmup_epochs, start_warmup_value=args.lr_start,
        )

    if dist_utils.is_main_process() and args.wandb:
        wandb_id = os.path.split(args.output_dir)[-1]
        wandb.init(project='LaViLa', id=wandb_id, config=args, resume='allow')

    print(args)

    best_metric = 0.
    print("=> beginning training")
    for epoch in range(args.start_epoch, args.epochs):
        if args.distributed:
            train_sampler.set_epoch(epoch)

        train_stats = train(train_loader, model, criterion, optimizer, scaler, epoch, lr_schedule, args)

        is_epoch = ((epoch + 1) % args.save_freq) == 0

        print('=> saving checkpoint')
        dist_utils.save_on_master({
            'epoch': epoch + 1,
            'state_dict': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'scaler': scaler.state_dict(),
            'best_acc1': 0,
            'args': args,
        }, False, args.output_dir, is_epoch=is_epoch)

        if ((epoch + 1) % args.eval_freq) == 0:
            if args.use_vn_classifier:
                val_stats = validate_multihead(val_loader, model, args)
            else:
                val_stats = validate(val_loader, model, args)
            if val_stats['acc1'] > best_metric:
                is_best = True
                best_metric = val_stats['acc1']
            else:
                is_best = False

            print('=> saving checkpoint')
            dist_utils.save_on_master({
                'epoch': epoch + 1,
                'state_dict': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scaler': scaler.state_dict(),
                'best_acc1': best_metric,
                'args': args,
            }, is_best, args.output_dir, is_epoch=is_epoch)

            log_stats = {**{f'train_{k}': v for k, v in train_stats.items()},
                         **{f'test_{k}': v for k, v in val_stats.items()},
                         'epoch': epoch}

            if dist_utils.is_main_process():
                if args.wandb:
                    wandb.log(log_stats)
                with open(os.path.join(args.output_dir, 'log.txt'), 'a') as f:
                    f.write(json.dumps(log_stats) + '\n')


def train(train_loader, model, criterion, optimizer, scaler, epoch, lr_schedule, args):
    batch_time = AverageMeter('Time', ':6.2f')
    data_time = AverageMeter('Data', ':6.2f')
    mem = AverageMeter('Mem (GB)', ':6.1f')
    iters_per_epoch = len(train_loader) // args.update_freq
    losses = AverageMeter('Loss', ':.4e')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    top1_noun = AverageMeter('Noun Acc@1', ':6.2f')
    top1_verb = AverageMeter('Verb Acc@1', ':6.2f')
    progress = ProgressMeter(
        iters_per_epoch,
        [batch_time, data_time, mem, losses, top1, top5, top1_noun, top1_verb],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    for data_iter, (images, target) in enumerate(train_loader):
        optim_iter = data_iter // args.update_freq

        # measure data loading time
        data_time.update(time.time() - end)

        # update weight decay and learning rate according to their schedule
        it = iters_per_epoch * epoch + optim_iter  # global training iteration
        for k, param_group in enumerate(optimizer.param_groups):
            if lr_schedule is not None:
                param_group['lr'] = lr_schedule[it] * args.lr_multiplier_on_backbone
            else:
                param_group['lr'] = lr_schedule[it]

        images = images.cuda(args.gpu, non_blocking=True)
        target = target.cuda(args.gpu, non_blocking=True)

        # compute output
        with amp.autocast(enabled=not args.disable_amp):
            output = model(images, use_checkpoint=args.use_checkpoint)
            if isinstance(output, list):
                assert len(output) == 3
                target_to_verb = torch.tensor([args.mapping_act2v[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
                loss = criterion(output[0], target_to_verb)
                target_to_noun = torch.tensor([args.mapping_act2n[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
                loss += criterion(output[1], target_to_noun)
                loss += criterion(output[2], target)
            else:
                loss = criterion(output, target)
            loss /= args.update_freq

        if not math.isfinite(loss.item()):
            print("Loss is {}, stopping training".format(loss.item()))
            sys.exit(1)

        scaler.scale(loss).backward()

        if (data_iter + 1) % args.update_freq != 0:
            continue

        if args.clip_grad_value is not None:
            scaler.unscale_(optimizer)
            if args.clip_grad_type == 'norm':
                torch.nn.utils.clip_grad_norm_(
                    model.parameters(), args.clip_grad_value, norm_type=2.
                )
            elif args.clip_grad_type == 'value':
                torch.nn.utils.clip_grad_value_(model.parameters(), args.clip_grad_value)
            else:
                assert False, f"Unknown clip mode ({args.clip_grad_type})."
        # compute gradient and do SGD step
        scaler.step(optimizer)
        scaler.update()
        model.zero_grad(set_to_none=True)

        if isinstance(output, list):
            target_to_verb = torch.tensor([args.mapping_act2v[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
            acc1_verb, _ = accuracy(output[0], target_to_verb, topk=(1, 5))
            top1_verb.update(acc1_verb.item(), images.size(0))
            target_to_noun = torch.tensor([args.mapping_act2n[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
            acc1_noun, _ = accuracy(output[1], target_to_noun, topk=(1, 5))
            top1_noun.update(acc1_noun.item(), images.size(0))
            acc1, acc5 = accuracy(output[2], target, topk=(1, 5))
            losses.update(loss.item(), images.size(0))
            top1.update(acc1.item(), images.size(0))
            top5.update(acc5.item(), images.size(0))
        else:
            output = torch.softmax(output, dim=1)
            acc1, acc5 = accuracy(output, target, topk=(1, 5))
            losses.update(loss.item(), images.size(0))
            top1.update(acc1.item(), images.size(0))
            top5.update(acc5.item(), images.size(0))
            if args.dataset == 'ek100_cls':
                vi = get_marginal_indexes(args.actions, 'verb')
                ni = get_marginal_indexes(args.actions, 'noun')
                verb_scores = torch.tensor(marginalize(output.detach().cpu().numpy(), vi)).cuda(args.gpu, non_blocking=True)
                noun_scores = torch.tensor(marginalize(output.detach().cpu().numpy(), ni)).cuda(args.gpu, non_blocking=True)
                target_to_verb = torch.tensor([args.mapping_act2v[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
                target_to_noun = torch.tensor([args.mapping_act2n[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
                acc1_verb, _ = accuracy(verb_scores, target_to_verb, topk=(1, 5))
                acc1_noun, _ = accuracy(noun_scores, target_to_noun, topk=(1, 5))
                top1_verb.update(acc1_verb.item(), images.size(0))
                top1_noun.update(acc1_noun.item(), images.size(0))
            else:
                top1_verb.update(0., images.size(0))
                top1_noun.update(0., images.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        mem.update(torch.cuda.max_memory_allocated() // 1e9)

        if optim_iter % args.print_freq == 0:
            if dist_utils.is_main_process() and args.wandb:
                wandb.log({
                    'acc1': top1.avg, 'acc5': top5.avg, 'loss': losses.avg,
                    'acc1_verb': top1_verb.avg, 'acc1_noun': top1_noun.avg,
                })
            progress.display(optim_iter)
    progress.synchronize()
    return {
        'acc1': top1.avg, 'acc5': top5.avg, 'loss': losses.avg,
        'acc1_verb': top1_verb.avg, 'acc1_noun': top1_noun.avg,
        'lr': optimizer.param_groups[0]['lr'],
    }


def validate(val_loader, model, args):
    batch_time = AverageMeter('Time', ':6.2f')
    data_time = AverageMeter('Data', ':6.2f')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, top1, top5],
        prefix='Test: '
    )

    # switch to eval mode
    model.eval()
    if args.use_half:
        model.half()

    all_outputs = []
    all_targets = []
    with torch.no_grad():
        end = time.time()
        for i, (images, target) in enumerate(val_loader):
            # measure data loading time
            data_time.update(time.time() - end)
            if isinstance(images, list):
                logit_allcrops = []
                for crop in images:
                    crop = crop.cuda(args.gpu, non_blocking=True)
                    if args.use_half:
                        crop = crop.half()
                    logit = model(crop, use_checkpoint=args.use_checkpoint)
                    logit_allcrops.append(logit)
                logit_allcrops = torch.stack(logit_allcrops, 0)
                logit = logit_allcrops.mean(0)
                logit = torch.softmax(logit, dim=1)
                target = target.cuda(args.gpu, non_blocking=True)

                acc1, acc5 = accuracy(logit, target, topk=(1, 5))
                top1.update(acc1.item(), target.size(0))
                top5.update(acc5.item(), target.size(0))
            else:
                images = images.cuda(args.gpu, non_blocking=True)
                target = target.cuda(args.gpu, non_blocking=True)
                if args.use_half:
                    images = images.half()

                logit = model(images, use_checkpoint=args.use_checkpoint)
                logit = torch.softmax(logit, dim=1)

                acc1, acc5 = accuracy(logit, target, topk=(1, 5))
                top1.update(acc1.item(), images.size(0))
                top5.update(acc5.item(), images.size(0))

            all_outputs.append(logit)
            all_targets.append(target)
            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                progress.display(i)
    progress.synchronize()
    if args.dataset == 'ek100_cls':
        print('EK100 * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))
    else:
        print('EGTEA * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))
    all_outputs = torch.cat(all_outputs).cpu().numpy()
    all_targets = torch.cat(all_targets).cpu().numpy()
    cm = confusion_matrix(all_targets, all_outputs.argmax(axis=1))
    mean_acc, acc = get_mean_accuracy(cm)
    print('Mean Acc. = {:.3f}, Top-1 Acc. = {:.3f}'.format(mean_acc, acc))

    if args.dataset == 'ek100_cls':
        vi = get_marginal_indexes(args.actions, 'verb')
        ni = get_marginal_indexes(args.actions, 'noun')
        verb_scores = marginalize(all_outputs, vi)
        noun_scores = marginalize(all_outputs, ni)
        target_to_verb = np.array([args.mapping_act2v[a] for a in all_targets.tolist()])
        target_to_noun = np.array([args.mapping_act2n[a] for a in all_targets.tolist()])
        cm = confusion_matrix(target_to_verb, verb_scores.argmax(axis=1))
        _, acc = get_mean_accuracy(cm)
        print('Verb Acc@1: {:.3f}'.format(acc))
        cm = confusion_matrix(target_to_noun, noun_scores.argmax(axis=1))
        _, acc = get_mean_accuracy(cm)
        print('Noun Acc@1: {:.3f}'.format(acc))
    return {'acc1': top1.avg, 'acc5': top5.avg, 'mean_acc': mean_acc}


def validate_multihead(val_loader, model, args):
    batch_time = AverageMeter('Time', ':6.2f')
    data_time = AverageMeter('Data', ':6.2f')
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    top1_verb = AverageMeter('Verb Acc@1', ':6.2f')
    top1_noun = AverageMeter('Noun Acc@1', ':6.2f')
    progress = ProgressMeter(
        len(val_loader),
        [batch_time, top1, top5, top1_verb, top1_noun],
        prefix='Test: '
    )

    # switch to eval mode
    model.eval()
    if args.use_half:
        model.half()

    all_verb_outputs = []
    all_noun_outputs = []
    all_action_outputs = []
    all_verb_targets = []
    all_noun_targets = []
    all_action_targets = []
    with torch.no_grad():
        end = time.time()
        for i, (images, target) in enumerate(val_loader):
            # measure data loading time
            data_time.update(time.time() - end)
            if isinstance(images, torch.Tensor):
                images = [images, ]
            logit_verb_allcrops = []
            logit_noun_allcrops = []
            logit_action_allcrops = []
            for crop in images:
                crop = crop.cuda(args.gpu, non_blocking=True)
                if args.use_half:
                    crop = crop.half()
                logit = model(crop, use_checkpoint=args.use_checkpoint)
                logit_verb_allcrops.append(logit[0])
                logit_noun_allcrops.append(logit[1])
                logit_action_allcrops.append(logit[2])
            logit_verb_allcrops = torch.stack(logit_verb_allcrops, 0)
            logit_noun_allcrops = torch.stack(logit_noun_allcrops, 0)
            logit_action_allcrops = torch.stack(logit_action_allcrops, 0)
            logit_verb = logit_verb_allcrops.mean(0)
            logit_noun = logit_noun_allcrops.mean(0)
            logit_action = logit_action_allcrops.mean(0)
            logit_noun = torch.softmax(logit_noun, dim=1)
            logit_verb = torch.softmax(logit_verb, dim=1)
            logit_action = torch.softmax(logit_action, dim=1)
            target = target.cuda(args.gpu, non_blocking=True)
            target_to_verb = torch.tensor([args.mapping_act2v[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)
            target_to_noun = torch.tensor([args.mapping_act2n[a] for a in target.tolist()]).cuda(args.gpu, non_blocking=True)

            acc1, acc5 = accuracy(logit_action, target, topk=(1, 5))
            acc1_verb, _ = accuracy(logit_verb, target_to_verb, topk=(1, 5))
            acc1_noun, _ = accuracy(logit_noun, target_to_noun, topk=(1, 5))
            top1.update(acc1.item(), target.size(0))
            top5.update(acc5.item(), target.size(0))
            top1_verb.update(acc1_verb.item(), target_to_verb.size(0))
            top1_noun.update(acc1_noun.item(), target_to_noun.size(0))

            all_verb_outputs.append(logit_verb)
            all_noun_outputs.append(logit_noun)
            all_action_outputs.append(logit_action)
            all_verb_targets.append(target_to_verb)
            all_noun_targets.append(target_to_noun)
            all_action_targets.append(target)
            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if i % args.print_freq == 0:
                progress.display(i)
    progress.synchronize()
    print('EK100 * Verb Acc@1 {top1.avg:.3f}'.format(top1=top1_verb))
    print('EK100 * Noun Acc@1 {top1.avg:.3f}'.format(top1=top1_noun))
    print('EK100 * Action Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'.format(top1=top1, top5=top5))
    return {'acc1': top1.avg, 'acc5': top5.avg, 'acc1_verb': top1_verb.avg, 'acc1_noun': top1_noun.avg}


if __name__ == '__main__':
    parser = argparse.ArgumentParser('lavila finetune and evaluation', parents=[get_args_parser()])
    args = parser.parse_args()
    os.makedirs(args.output_dir, exist_ok=True)
    main(args)