-
Notifications
You must be signed in to change notification settings - Fork 135
/
multimodal_eval.py
108 lines (91 loc) · 3.55 KB
/
multimodal_eval.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import logging
import torch
from flava.data import default_text_transform, VL_MAX_LENGTH_DEFAULT
from flava.data.imagenet_zeroshot_data import (
imagenet_classnames,
openai_imagenet_template,
)
from pytorch_lightning import Callback, LightningDataModule
from pytorch_lightning.utilities import rank_zero_only
from tqdm import tqdm
logger = logging.getLogger(__name__)
def _zero_shot_classifier(model, device, text_transform, *args, **kwargs):
zeroshot_weights = []
for classname in tqdm(imagenet_classnames):
texts = text_transform(
[template(classname) for template in openai_imagenet_template]
)["input_ids"]
texts = texts.to(device)
class_embeddings = model.encode_text(texts)
class_embeddings /= class_embeddings.norm(dim=-1, keepdim=True)
class_embedding = class_embeddings.mean(dim=0)
class_embedding /= class_embedding.norm()
zeroshot_weights.append(class_embedding)
zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(device)
return zeroshot_weights
def _accuracy(output, target, topk=(1,)):
pred = output.topk(max(topk), 1, True, True)[1].t()
correct = pred.eq(target.view(1, -1).expand_as(pred))
return [
float(correct[:k].reshape(-1).float().sum(0, keepdim=True).cpu().numpy())
for k in topk
]
@rank_zero_only
def run_imagenet_zero_shot(model, dataloader, device, text_transform, *args, **kwargs):
logger.info("Starting ImageNet Zero-Shot Eval")
logger.info("Building classifier")
classifier = _zero_shot_classifier(model, device, text_transform)
logger.info("Classifier built")
top1, top5, n = 0.0, 0.0, 0.0
for sample in tqdm(dataloader):
images = sample["image"]
target = sample["label"]
images = images.to(device)
target = target.to(device)
# predict
# if hasattr(model, "module"):
# image_features = model.module.encode_image({"image": images})
# else:
image_features = model.encode_image(images)
image_features /= image_features.norm(dim=-1, keepdim=True)
logits = 100.0 * image_features @ classifier
# measure accuracy
acc1, acc5 = _accuracy(logits, target, topk=(1, 5))
top1 += acc1
top5 += acc5
n += images.size(0)
top1 = top1 / n
top5 = top5 / n
results = {}
results["imagenet-zeroshot-val-top1"] = top1
results["imagenet-zeroshot-val-top5"] = top5
return results
class MultimodalEvalCallback(Callback):
def __init__(self, imagenet_datamodule: LightningDataModule, *args, **kwargs):
super().__init__()
self.imagenet_val_dataloader = imagenet_datamodule.val_dataloader()
self.text_transform = default_text_transform(
max_text_length=VL_MAX_LENGTH_DEFAULT
)
@torch.no_grad()
def on_validation_start(self, trainer, pl_module, **kwargs) -> None:
metrics = run_imagenet_zero_shot(
pl_module.model,
self.imagenet_val_dataloader,
pl_module.device,
self.text_transform,
)
if metrics is not None:
for key in metrics:
self.log(
f"val/{key}",
metrics[key],
prog_bar=True,
logger=True,
rank_zero_only=True,
)