# 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 os import os.path as osp import pickle import time import torch import torchvision.transforms as transforms import torchvision.transforms._transforms_video as transforms_video from lavila.data import datasets from lavila.data.video_transforms import Permute from lavila.models import models from lavila.utils.preprocess import generate_tokenizer from lavila.utils import distributed as dist_utils from eval_narrator import decode_one class IndexedDataset(torch.utils.data.Dataset): def __init__(self, dataset): self.dataset = dataset def __getitem__(self, index): return index, self.dataset[index] def __len__(self): return len(self.dataset) def get_args_parser(): parser = argparse.ArgumentParser(description='lavila infer narrator', add_help=False) 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('--output-dir', default='./', type=str, help='output dir') parser.add_argument('--batch-size', default=64, type=int) parser.add_argument('--use-half', action='store_true') 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('--resume', default='', type=str, help='path to latest checkpoint') parser.add_argument('--caption-sample', default='multinomial_sample', choices=['multinomial_sample', 'beam_sample', 'group_beam_search']) parser.add_argument('--caption-top-k', default=None, type=int) parser.add_argument('--caption-top-p', default=0.95, type=float) parser.add_argument('--caption-num-beams', default=1, type=int) parser.add_argument('--caption-num-beam-groups', default=1, type=int) parser.add_argument('--caption-temperature', default=0.7, type=float) parser.add_argument('--caption-length-penalty', default=1.0, type=float) parser.add_argument('--caption-num-return-sequences', default=10, type=int) parser.add_argument('--caption-max-len', default=77, type=int) parser.add_argument('--caption-early-stop', action='store_true', help='early stopping to save computation') # 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('--gpu', default=None, type=int, help='GPU id to use.') return parser def main(args): dist_utils.init_distributed_mode(args) print(args) if args.resume: ckpt_path = args.resume elif osp.isfile(osp.join(args.output_dir, 'checkpoint_best.pt')): ckpt_path = osp.join(args.output_dir, 'checkpoint_best.pt') else: raise Exception('no checkpoint found') ckpt = torch.load(ckpt_path, map_location='cpu') state_dict = OrderedDict() for k, v in ckpt['state_dict'].items(): state_dict[k.replace('module.', '')] = v # create model old_args = ckpt['args'] print('=> creating model: {}'.format(old_args.model)) model = getattr(models, old_args.model)( text_use_cls_token=old_args.use_cls_token, gated_xattn=old_args.gated_xattn, timesformer_gated_xattn=old_args.timesformer_gated_xattn, num_frames=old_args.clip_length, drop_path_rate=0, ) model.cuda() model.load_state_dict(state_dict, strict=True) print("=> loaded resume checkpoint '{}' (epoch {})".format(args.resume, ckpt['epoch'])) torch.backends.cudnn.benchmark = True # Data loading print("=> creating dataset") tokenizer = generate_tokenizer(old_args.model) crop_size = 224 if '336PX' not in old_args.model else 336 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])), ]) val_dataset = datasets.VideoCaptionDatasetCLIP( args.dataset, args.root, args.metadata, transform=val_transform, is_training=False, tokenizer=tokenizer, clip_length=args.clip_length, clip_stride=args.clip_stride, sparse_sample=False, subsample_stride=1, ) val_dataset = IndexedDataset(val_dataset) print(len(val_dataset)) if args.distributed: val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset, shuffle=False) else: val_sampler = None val_loader = torch.utils.data.DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True, sampler=val_sampler, drop_last=False ) print('len(val_loader) = {}'.format(len(val_loader))) model.eval() if args.use_half: model.half() id_offset = 0 all_captions_cache = [] end = time.time() with torch.no_grad(): for data_iter, (indices, inputs) in enumerate(val_loader): indices = indices.tolist() if data_iter % args.print_freq == 0: print("finished {}/{} in {}".format(data_iter, len(val_loader), time.time() - end)) end = time.time() if len(inputs) == 2 or len(inputs) == 3: images = inputs[0].cuda(non_blocking=True) if args.use_half: images = images.half() image_features = dist_utils.get_model(model).encode_image(images) if not isinstance(image_features, (list, tuple)): image_tokens = image_features else: image_tokens = image_features[1] if args.caption_sample == 'multinomial_sample': generated_text_ids, ppls = dist_utils.get_model(model).generate( image_tokens, tokenizer, target=None, max_text_length=args.caption_max_len, top_k=args.caption_top_k, top_p=args.caption_top_p, num_return_sequences=args.caption_num_return_sequences, temperature=args.caption_temperature, early_stopping=args.caption_early_stop, ) elif args.caption_sample == 'beam_sample': generated_text_ids, ppls = dist_utils.get_model(model).beam_sample( image_tokens, tokenizer, target=None, max_text_length=args.caption_max_len, top_k=args.caption_top_k, top_p=args.caption_top_p, temperature=args.caption_temperature, length_penalty=args.caption_length_penalty, num_beams=args.caption_num_beams, num_return_sequences=args.caption_num_return_sequences, ) elif args.caption_sample == 'group_beam_search': assert args.caption_num_beam_groups > 1 and args.caption_num_beams % args.caption_num_beam_groups == 0 generated_text_ids, ppls = dist_utils.get_model(model).group_beam_search( image_tokens, tokenizer, target=None, max_text_length=args.caption_max_len, top_k=args.caption_top_k, top_p=args.caption_top_p, temperature=args.caption_temperature, length_penalty=args.caption_length_penalty, num_beams=args.caption_num_beams, num_beam_groups=args.caption_num_beam_groups, num_return_sequences=args.caption_num_return_sequences, ) for j in range(generated_text_ids.shape[0] // args.caption_num_return_sequences): generated_text_str_list = [] ppls_list = [] for k in range(args.caption_num_return_sequences): jj = j * args.caption_num_return_sequences + k generated_text_str = decode_one(generated_text_ids[jj], tokenizer) generated_text_str_list.append(generated_text_str) ppls_list.append(ppls[jj].item()) video_uid, t_start, t_end, _ = val_loader.dataset.dataset.samples[indices[j]] if args.caption_num_return_sequences == 1: all_captions_cache.append((video_uid, t_start, t_end, generated_text_str, ppls[jj].item())) else: all_captions_cache.append((video_uid, t_start, t_end, generated_text_str_list, ppls_list)) id_offset += generated_text_ids.shape[0] pickle.dump(all_captions_cache, open(osp.join(args.output_dir, 'cache.{}.pkl'.format(args.rank)), 'wb')) torch.distributed.barrier() disorded_list = [] total_num = 0 if args.rank == 0: for i in range(args.world_size): print('=> reading {}'.format(osp.join(args.output_dir, f'cache.{i}.pkl'))) sublist = pickle.load(open(osp.join(args.output_dir, f'cache.{i}.pkl'), 'rb')) disorded_list.append(sublist) total_num += len(sublist) ordered_list = [] for i in range(total_num): ordered_list.append(disorded_list[i % args.world_size][i // args.world_size]) print(f"{len(val_dataset)}/{len(ordered_list)}") ordered_list = ordered_list[:len(val_dataset)] pickle.dump(ordered_list, open(osp.join(args.output_dir, 'total.pkl'), 'wb')) for i in range(args.world_size): print('=> deleting {}'.format(osp.join(args.output_dir, f'cache.{i}.pkl'))) os.remove(osp.join(args.output_dir, f'cache.{i}.pkl')) if __name__ == '__main__': parser = argparse.ArgumentParser('lavila infer narrator', parents=[get_args_parser()]) args = parser.parse_args() main(args)