# 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)