| |
|
| | 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]
|
| | texts = tokenize(texts).to(args.device)
|
| | 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():
|
| |
|
| | 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
|
| |
|
| |
|
| | 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
|
| |
|
