main_pretrain.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 os
import pandas as pd
import sys
import time

import torch
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
import wandb

from eval_zeroshot import get_similarity_matrix
from lavila.data import datasets
from lavila.data.video_transforms import Permute
from lavila.models import models
from lavila.utils.meter import AverageMeter, ProgressMeter
from lavila.utils import distributed as dist_utils
from lavila.utils.evaluation_ek100mir import get_mAP, get_nDCG
from lavila.utils.preprocess import generate_tokenizer
from lavila.utils.random import random_seed
from lavila.utils.scheduler import cosine_scheduler


class GroundTruthDataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        self.dataset = dataset

    def __getitem__(self, index):
        return 1, self.dataset[index]

    def __len__(self):
        return len(self.dataset)


class PseudoLabelDataset(torch.utils.data.Dataset):
    def __init__(self, dataset):
        self.dataset = dataset

    def __getitem__(self, index):
        return 0, self.dataset[index]

    def __len__(self):
        return len(self.dataset)


def get_args_parser():
    parser = argparse.ArgumentParser(description='LaVid training and evaluation', add_help=False)
    # Data
    parser.add_argument('--dataset', default='ego4d', type=str, choices=['ego4d'])
    parser.add_argument('--root', default='datasets/Ego4D/video_5min_chunks_288px/',
                        type=str, help='path to dataset root')
    parser.add_argument('--metadata', default='datasets/Ego4D/ego4d_train.pkl',
                        type=str, help='path to metadata file')
    parser.add_argument('--metadata-aux', default=None, nargs='+',
                        type=str, help='path to metadata file (auxiliary data with pseudo narrations)')
    parser.add_argument('--output-dir', default='./', type=str, help='output dir')
    parser.add_argument('--clip-length', default=4, type=int, help='clip length')
    parser.add_argument('--clip-stride', default=16, type=int, help='clip stride')
    parser.add_argument('--sparse-sample', action='store_true', help='switch to sparse sampling')
    parser.add_argument('--narration-selection', default='random',
                        choices=['random', 'concat'],
                        type=str, help='selection strategy if multiple narrations per clip')
    parser.add_argument('--num-hard-neg', default=0, type=int, help='number of hard negatives per video')
    # Model
    parser.add_argument('--model', default='CLIP_OPENAI_TIMESFORMER_BASE', type=str)
    parser.add_argument('--norm-embed', action='store_true', help='norm text and visual embed if set True')
    parser.add_argument('--resume', default='', type=str, help='path to resume from')
    parser.add_argument('--load-visual-pretrained', default=None, type=str,
                        help='path to pretrained model (in1k/in21k/...)')
    parser.add_argument('--project-embed-dim', default=256, type=int, help='embed dim after projection')
    parser.add_argument('--use-cls-token', action='store_true', help='use feature at [CLS] if set True')
    parser.add_argument('--contrastive-use-vissl', action='store_true', help='use contrastive implementation in vissl')
    parser.add_argument('--gated-xattn', action='store_true', help='use gated x-attn in VCLM_GPT2')
    parser.add_argument('--random-init-gpt2', action='store_true', help='random initialize params of text decoder in VCLM_GPT2')
    parser.add_argument('--timesformer-gated-xattn', action='store_true', help='use gated x-attn in TimeSformer')
    parser.add_argument('--timesformer-freeze-space', action='store_true', help='freeze space part in TimeSformer')
    parser.add_argument('--drop-path-rate', default=0., type=float, help='DropPath rate')
    parser.add_argument('--freeze-visual-vclm', action='store_true', help='freeze the visual model in VCLM_GPT2')
    parser.add_argument('--freeze-visual-vclm-temporal', action='store_true', help='freeze the temporal part of visual model in VCLM_GPT2')
    parser.add_argument('--freeze-lm-vclm', action='store_true', help='freeze the lm in VCLM_GPT2')
    parser.add_argument('--find-unused-parameters', action='store_true',
                        help='do this during DDP (useful for models with tied weights)')
    # Training
    parser.add_argument('--epochs', default=5, 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=32, type=int,
                        help='number of samples per-device/per-gpu')
    parser.add_argument('--temperature-init', default=0.07, type=float,
                        help='init. logit temperature for samples')
    parser.add_argument('--freeze-temperature', action='store_true',
                        help='freeze logit temperature')
    parser.add_argument('--pseudo-temperature-init', default=0.07, type=float,
                        help='init. logit temperature for pseudo-narrated samples')
    parser.add_argument('--freeze-pseudo-temperature', action='store_true',
                        help='freeze logit temperature (for pseudo-narrated samples)')
    parser.add_argument('--lr', default=3e-5, 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('--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('--eval-freq', default=99, type=int)
    parser.add_argument('--eval-in-middle-freq', default=-1, type=int)
    parser.add_argument('--save-freq', default=1, 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')
    parser.add_argument('--use-half', action='store_true', help='evaluate using half-precision')
    # System
    parser.add_argument('--print-freq', default=10, type=int, help='print frequency')
    parser.add_argument('-j', '--workers', default=10, type=int, metavar='N',
                        help='number of data loading workers per process')
    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())

    print("=> creating model: {}".format(args.model))
    model = getattr(models, args.model)(
        pretrained=args.load_visual_pretrained,
        pretrained2d=args.load_visual_pretrained is not None,
        text_use_cls_token=args.use_cls_token,
        project_embed_dim=args.project_embed_dim,
        gated_xattn=args.gated_xattn,
        random_init_gpt2=args.random_init_gpt2,
        timesformer_gated_xattn=args.timesformer_gated_xattn,
        timesformer_freeze_space=args.timesformer_freeze_space,
        freeze_lm_vclm=args.freeze_lm_vclm,
        freeze_visual_vclm=args.freeze_visual_vclm,
        freeze_visual_vclm_temporal=args.freeze_visual_vclm_temporal,
        num_frames=args.clip_length,
        drop_path_rate=args.drop_path_rate,
        temperature_init=args.temperature_init,
    )
    if args.freeze_temperature:
        print('Freeze logit temperature')
        model.logit_scale.requires_grad = False
    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
        )

    tokenizer = generate_tokenizer(args.model)

    if args.metadata_aux is None:
        criterion = models.get_loss(args.model, args, tokenizer=tokenizer).cuda(args.gpu)
    else:
        criterion = models.loss.SSLCLIPLoss(
            use_vissl=args.contrastive_use_vissl,
            cache_labels=True,
            rank=args.rank,
            world_size=args.world_size,
            scale_init=args.pseudo_temperature_init,
            freeze_scale=args.freeze_pseudo_temperature,
            ).cuda(args.gpu)

    p_wd, p_non_wd = [], []
    for n, p in model.named_parameters():
        if not p.requires_grad:
            continue  # frozen weights
        if 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)
    for n, p in criterion.named_parameters():
        if not p.requires_grad:
            continue
        p_non_wd.append(p)

    optim_params = [{"params": p_wd, "weight_decay": args.wd},
                    {"params": p_non_wd, "weight_decay": 0}]

    if args.use_zero:
        optimizer = ZeroRedundancyOptimizer(
            optim_params, optimizer_class=torch.optim.AdamW,
            lr=args.lr, betas=args.betas, eps=args.eps, 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
            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 ()
            criterion.load_state_dict(checkpoint['criterion']) if 'criterion' in checkpoint else ()
            best_acc1 = checkpoint['best_acc1']
            print("=> loaded resume checkpoint '{}' (epoch {})"
                  .format(args.resume, epoch))
        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")

    crop_size = 224 if '336PX' not in 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)),
    ]
    if 'OPENAI' in 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)

    # TODO: uncomment when evaluation is done later
    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 args.model else
             transforms_video.NormalizeVideo(mean=[108.3272985, 116.7460125, 104.09373615000001], std=[68.5005327, 66.6321579, 70.32316305]))
        ])

    assert 'train' in args.metadata
    train_dataset = datasets.get_dataset(train_transform, tokenizer, args, is_training=True)
    args.metadata = args.metadata.replace('train', 'val')
    val_dataset = datasets.get_dataset(val_transform, tokenizer, args, is_training=False)
    args.metadata = args.metadata.replace('val', 'train')
    if args.metadata_aux is not None:
        train_dataset = GroundTruthDataset(train_dataset)
        old_metadata = args.metadata
        aux_dataset_list = []
        for aux_i, aux_pkl in enumerate(args.metadata_aux):
            args.metadata = aux_pkl
            aux_dataset = datasets.get_dataset(train_transform, tokenizer, args, is_training=True)
            aux_dataset_list.append(PseudoLabelDataset(aux_dataset))
            print("auxiliary dataset [{}] : source = {}, len(aux_dataset) = {}".format(aux_i, aux_pkl, len(aux_dataset)))
        pseudo_label_dataset = torch.utils.data.ConcatDataset(aux_dataset_list)
        args.metadata = old_metadata
        train_dataset = torch.utils.data.ConcatDataset([train_dataset, pseudo_label_dataset])
        val_dataset = GroundTruthDataset(val_dataset)

    ek100_dataset = datasets.VideoCaptionDatasetCLIP(
        'ek100_mir',
        'datasets/EK100/video_ht256px/',
        'datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test.csv',
        transform=val_transform,
        is_training=False,
        tokenizer=tokenizer,
        clip_length=args.clip_length,
        clip_stride=args.clip_stride,
        sparse_sample=False
    )

    if args.distributed:
        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
        ek100_sampler = torch.utils.data.SequentialSampler(ek100_dataset)
    else:
        train_sampler = None
        val_sampler = None
        ek100_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)))
    ek100_loader = torch.utils.data.DataLoader(
        ek100_dataset, batch_size=args.batch_size * (1 + args.num_hard_neg), shuffle=(ek100_sampler is None),
        num_workers=args.workers, pin_memory=True, sampler=ek100_sampler, drop_last=False
    )
    print('len(ek100_loader) = {}'.format(len(ek100_loader)))

    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='LaVid', 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)

        if hasattr(args, 'eval_in_middle_freq') and args.eval_in_middle_freq > 0:
            train_stats = train(train_loader, model, criterion, optimizer, scaler, epoch, lr_schedule, args,
                                ek100_loader=ek100_loader, eval_in_middle=args.eval_in_middle_freq)
        else:
            train_stats = train(train_loader, model, criterion, optimizer, scaler, epoch, lr_schedule, args)

        if args.model.startswith('CLIP'):
            print('=> 0-shot on EK100')
            similarity_matrix = get_similarity_matrix(ek100_loader, model, use_half=args.use_half)
            similarity_matrix = (similarity_matrix + 1) / 2
            video_id = pd.read_csv("datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test.csv").values[:, 0]
            text_id = pd.read_csv("datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test_sentence.csv").values[:, 0]
            indexes = [video_id.tolist().index(elem) for elem in text_id]
            similarity_matrix = similarity_matrix[:, indexes]
            rel_matrix = pd.read_pickle(
                'datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl'
            )
            vis_map, txt_map, avg_map = get_mAP(similarity_matrix, rel_matrix)
            print('mAP: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_map, txt_map, avg_map))
            vis_ndcg, txt_ndcg, avg_ndcg = get_nDCG(similarity_matrix, rel_matrix)
            print('nDCG: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_ndcg, txt_ndcg, avg_ndcg))
            if avg_map > best_metric:
                is_best = True
                best_metric = avg_map
            else:
                is_best = False
        else:
            is_best = False

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

        if args.distributed and args.use_zero:
            print("=> consolidating state_dict before saving (due to ZeRO)")
            optimizer.consolidate_state_dict()

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

        if (epoch + 1) % args.eval_freq != 0:
            continue

        # TODO: add evaluation
        val_stats = validate(val_loader, model, criterion, args)

        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, ek100_loader=None, eval_in_middle=0):
    batch_time = AverageMeter('Time', ':6.2f')
    data_time = AverageMeter('Data', ':6.2f')
    mem = AverageMeter('Mem (GB)', ':6.1f')
    metric_names = models.get_metric_names(args.model)
    if args.metadata_aux is not None:
        metric_names.extend(['num_gt', 'num_pseudo', 'clip_acc_gt', 'clip_acc_pseudo'])
    iters_per_epoch = len(train_loader) // args.update_freq
    metrics = OrderedDict([(name, AverageMeter(name, ':.2e')) for name in metric_names])
    progress = ProgressMeter(
        iters_per_epoch,
        [batch_time, data_time, mem, *metrics.values()],
        prefix="Epoch: [{}]".format(epoch))

    # switch to train mode
    model.train()

    end = time.time()
    for data_iter, inputs in enumerate(train_loader):
        # evaluate in the middle of training
        if eval_in_middle > 0 and (data_iter > 0 and data_iter % eval_in_middle) and ek100_loader is not None:
            model.eval()
            print('=> 0-shot on EK100 in the middle of training')
            similarity_matrix = get_similarity_matrix(ek100_loader, model, use_half=args.use_half)
            similarity_matrix = (similarity_matrix + 1) / 2
            video_id = pd.read_csv("datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test.csv").values[:, 0]
            text_id = pd.read_csv("datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/EPIC_100_retrieval_test_sentence.csv").values[:, 0]
            indexes = [video_id.tolist().index(elem) for elem in text_id]
            similarity_matrix = similarity_matrix[:, indexes]
            rel_matrix = pd.read_pickle(
                'datasets/EK100/epic-kitchens-100-annotations/retrieval_annotations/relevancy/caption_relevancy_EPIC_100_retrieval_test.pkl'
            )
            vis_map, txt_map, avg_map = get_mAP(similarity_matrix, rel_matrix)
            print('mAP: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_map, txt_map, avg_map))
            vis_ndcg, txt_ndcg, avg_ndcg = get_nDCG(similarity_matrix, rel_matrix)
            print('nDCG: V->T: {:.3f} T->V: {:.3f} AVG: {:.3f}'.format(vis_ndcg, txt_ndcg, avg_ndcg))
            best_metric = avg_map

            print('=> saving checkpoint')
            dist_utils.save_on_master({
                'epoch': epoch + data_iter / len(train_loader),
                'state_dict': model.state_dict(),
                'criterion': criterion.state_dict(),
                'optimizer': optimizer.state_dict(),
                'scaler': scaler.state_dict(),
                'best_acc1': best_metric,
                'args': args,
            }, False, args.output_dir, is_epoch=True)   # save every time (not to conflict the best_metric tracking in the regular validation phrase)
            model.train()

        if args.metadata_aux is not None:
            gt_indicators, inputs = inputs

        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]

        inputs = [tensor.cuda(args.gpu, non_blocking=True) for tensor in inputs]
        _ = inputs.pop()  # loader will a "relevancy" variable which is not needed except ek100_mir

        # compute output
        with amp.autocast(enabled=not args.disable_amp):
            outputs = model(
                *inputs,
                use_checkpoint=args.use_checkpoint,
                norm_embed=args.norm_embed
            )
            if args.metadata_aux is None:
                loss_dict = criterion(outputs)
            else:
                loss_dict = criterion(outputs, gt_indicators)
            loss = loss_dict['loss']
            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 hasattr(dist_utils.get_model(model), 'logit_scale'):
            # clamp logit scale to [0, 100]
            dist_utils.get_model(model).logit_scale.data.clamp_(0, 4.6052)
            logit_scale = dist_utils.get_model(model).logit_scale.exp().item()
        else:
            logit_scale = torch.nan

        for k in loss_dict:
            metrics[k].update(loss_dict[k].item(), args.batch_size)

        # 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({**{k: v.item() for k, v in loss_dict.items()},
                           'scaler': scaler.get_scale(), 'logit': logit_scale})
            progress.display(optim_iter)
    progress.synchronize()
    return {**{k: v.avg for k, v in metrics.items()},
            'lr': optimizer.param_groups[0]['lr'],
            'logit_scale': logit_scale}


def validate(val_loader, model, criterion, args):
    batch_time = AverageMeter('Time', ':6.2f')
    data_time = AverageMeter('Data', ':6.2f')
    mem = AverageMeter('Mem (GB)', ':6.1f')
    metric_names = models.get_metric_names(args.model)
    iters_per_epoch = len(val_loader) // args.update_freq
    metrics = OrderedDict([(name, AverageMeter(name, ':.2e')) for name in metric_names])
    progress = ProgressMeter(
        iters_per_epoch,
        [batch_time, data_time, mem, *metrics.values()],
        prefix="Test: "
    )

    # switch to eval mode
    model.eval()

    with torch.no_grad():
        end = time.time()
        for i, inputs in enumerate(val_loader):
            # measure data loading time
            data_time.update(time.time() - end)

            if args.metadata_aux is not None:
                gt_indicators, inputs = inputs

            inputs = [tensor.cuda(args.gpu, non_blocking=True) for tensor in inputs]
            _ = inputs.pop()  # loader will a "relevancy" variable which is not needed except ek100_mir

            # compute output
            outputs = model(
                *inputs,
                use_checkpoint=args.use_checkpoint,
                norm_embed=args.norm_embed
            )
            if args.metadata_aux is None:
                loss_dict = criterion(outputs)
            else:
                loss_dict = criterion(outputs, gt_indicators)

            for k in loss_dict:
                metrics[k].update(loss_dict[k].item(), args.batch_size)

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

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

            if i % args.print_freq == 0:
                if dist_utils.is_main_process() and args.wandb:
                    wandb.log({**{k: v.item() for k, v in loss_dict.items()}})
                progress.display(i)
    progress.synchronize()
    return {**{k: v.avg for k, v in metrics.items()}}


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