Spaces:
Sleeping
Sleeping
| import sys | |
| import os | |
| import argparse | |
| import pathlib | |
| from tqdm import tqdm | |
| import json | |
| import torch | |
| import torch.nn as nn | |
| # https://github.com/PyTorchLightning/pytorch-lightning/issues/11663 | |
| import sentencepiece; import pytorch_lightning as pl; import clip | |
| from transfer_experiments.train import LinearClassifier | |
| from load_aokvqa import load_aokvqa | |
| from evaluation.remap_predictions import map_to_choices | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument('--split', type=str, choices=['train', 'val', 'test'], required=True) | |
| parser.add_argument('--aokvqa-dir', type=pathlib.Path, required=True, dest='aokvqa_dir') | |
| parser.add_argument('--features', type=pathlib.Path, required=True) | |
| parser.add_argument('--out', type=argparse.FileType('w'), dest='output_file') | |
| # | |
| parser_weights = parser.add_mutually_exclusive_group(required=True) | |
| parser_weights.add_argument('--ckpt', type=pathlib.Path, dest='checkpoint_path') | |
| parser_weights.add_argument('--zero-shot', action='store_true', dest='clip_zero_shot') | |
| parser.add_argument('--inputs', nargs='+', type=str, choices=['question', 'image'], required=('--zero-shot' in sys.argv)) | |
| # | |
| parser.add_argument('--vocab', type=argparse.FileType('r')) | |
| parser.add_argument('--vocab-features', type=pathlib.Path, dest='vocab_features') | |
| parser.add_argument('--mc', action='store_true', dest='multiple_choice') | |
| parser.add_argument('--clip-model-type', type=str, | |
| choices=['RN50', 'RN50x4', 'RN50x16', 'RN50x64', 'RN101', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px'], | |
| dest='clip_model_type', required=('--zero-shot' in sys.argv and '--mc' in sys.argv)) | |
| # | |
| args = parser.parse_args() | |
| ## Load dataset | |
| aokvqa_set = load_aokvqa(args.aokvqa_dir, args.split) | |
| ## Load models | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| if args.checkpoint_path is not None: | |
| classifier = LinearClassifier.load_from_checkpoint(args.checkpoint_path) | |
| classifier.to(device) | |
| hp = classifier.hparams | |
| elif args.clip_zero_shot: | |
| classifier = nn.Identity().to(device) | |
| hp = pl.utilities.AttributeDict(backbone='clip', clip_model_type=args.clip_model_type, objective='zero-shot', inputs=args.inputs) | |
| # Load input features | |
| embeddings = torch.load(args.features) | |
| if hp.backbone == 'clip': | |
| for q in embeddings.keys(): | |
| embeddings[q]['question'] = embeddings[q]['question'] / embeddings[q]['question'].norm(dim=-1, keepdim=True) | |
| embeddings[q]['image'] = embeddings[q]['image'] / embeddings[q]['image'].norm(dim=-1, keepdim=True) | |
| # Load vocab, vocab features, clip | |
| if (hp.objective == 'classifier') or \ | |
| (hp.objective in ['contrastive', 'zero-shot'] and args.multiple_choice is False): | |
| vocab = args.vocab.read().splitlines() | |
| if hp.objective in ['contrastive', 'zero-shot']: | |
| if args.multiple_choice is False: | |
| vocab_features = torch.load(args.vocab_features).cpu() | |
| vocab_features /= vocab_features.norm(dim=-1, keepdim=True) | |
| else: | |
| clip_model = clip.load(hp.clip_model_type, device=device)[0] | |
| logit_scale = clip_model.logit_scale.exp().cpu() | |
| ## Prediction loop | |
| predictions = {} | |
| with torch.no_grad(): | |
| for o in tqdm(aokvqa_set): | |
| q = o['question_id'] | |
| # Load input embedding (from question / image) | |
| if hp.objective == 'zero-shot' and ('question' in hp.inputs and 'image' in hp.inputs): | |
| e = embeddings[q]['question'] + embeddings[q]['image'] | |
| elif 'question' in hp.inputs and 'image' in hp.inputs: | |
| e = torch.cat((embeddings[q]['question'], embeddings[q]['image'])) | |
| elif 'question' in hp.inputs: | |
| e = embeddings[q]['question'] | |
| elif 'image' in hp.inputs: | |
| e = embeddings[q]['image'] | |
| # Pass inputs through model | |
| e = e.unsqueeze(0).to(device) | |
| x = classifier(e)[0].cpu() | |
| # Predict | |
| if hp.objective in ['contrastive', 'zero-shot']: | |
| if args.multiple_choice: | |
| vocab = o['choices'] | |
| # Encode choices | |
| vocab_features = clip.tokenize(vocab).to(device) | |
| vocab_features = torch.stack([ | |
| clip_model.encode_text(v.unsqueeze(0)) for v in vocab_features | |
| ], dim=1)[0] | |
| vocab_features /= vocab_features.norm(dim=-1, keepdim=True) | |
| vocab_features = vocab_features.float().cpu() | |
| x = logit_scale * x @ vocab_features.t() | |
| x = x.softmax(dim=-1) | |
| predictions[q] = vocab[x.argmax().item()] | |
| ## Save and evaluate predictions | |
| # Map prediction to nearest neighbor choice (by word embeddings) | |
| if args.multiple_choice and hp.objective == 'classifier': | |
| predictions = map_to_choices(aokvqa_set, predictions) | |
| json.dump(predictions, args.output_file) | |