File size: 3,429 Bytes
bdab1da |
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 |
# NOTE: This script is currently not supported for CLAP.
import logging
from contextlib import suppress
import torch
import torch.nn.functional as F
from tqdm import tqdm
from open_clip import tokenize
from .imagenet_zeroshot_data import imagenet_classnames, openai_imagenet_template
def zero_shot_classifier(model, classnames, templates, args):
with torch.no_grad():
zeroshot_weights = []
for classname in tqdm(classnames):
texts = [template(classname) for template in templates] # format with class
texts = tokenize(texts).to(args.device) # tokenize
if args.distributed and not args.horovod:
class_embeddings = model.module.encode_text(texts)
else:
class_embeddings = model.encode_text(texts)
class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
class_embedding /= class_embedding.norm()
zeroshot_weights.append(class_embedding)
zeroshot_weights = torch.stack(zeroshot_weights, dim=1).to(args.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
]
def run(model, classifier, dataloader, args):
autocast = torch.cuda.amp.autocast if args.precision == "amp" else suppress
with torch.no_grad():
top1, top5, n = 0.0, 0.0, 0.0
for images, target in tqdm(dataloader, unit_scale=args.batch_size):
images = images.to(args.device)
target = target.to(args.device)
with autocast():
# predict
if args.distributed and not args.horovod:
image_features = model.module.encode_image(images)
else:
image_features = model.encode_image(images)
image_features = F.normalize(image_features, dim=-1)
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
return top1, top5
def zero_shot_eval(model, data, epoch, args):
if "imagenet-val" not in data and "imagenet-v2" not in data:
return {}
if args.zeroshot_frequency == 0:
return {}
if (epoch % args.zeroshot_frequency) != 0 and epoch != args.epochs:
return {}
logging.info("Starting zero-shot imagenet.")
logging.info("Building zero-shot classifier")
classifier = zero_shot_classifier(
model, imagenet_classnames, openai_imagenet_template, args
)
logging.info("Using classifier")
results = {}
if "imagenet-val" in data:
top1, top5 = run(model, classifier, data["imagenet-val"].dataloader, args)
results["imagenet-zeroshot-val-top1"] = top1
results["imagenet-zeroshot-val-top5"] = top5
if "imagenet-v2" in data:
top1, top5 = run(model, classifier, data["imagenet-v2"].dataloader, args)
results["imagenetv2-zeroshot-val-top1"] = top1
results["imagenetv2-zeroshot-val-top5"] = top5
logging.info("Finished zero-shot imagenet.")
return results
|