85 lines
2.8 KiB
Python
85 lines
2.8 KiB
Python
#! pip install ftfy regex tqdm
|
|
#! pip install git+https://github.com/openai/CLIP.git
|
|
|
|
import numpy as np
|
|
import torch
|
|
import clip
|
|
from tqdm.notebook import tqdm
|
|
from pkg_resources import packaging
|
|
|
|
print("Torch version:", torch.__version__)
|
|
|
|
|
|
clip.available_models()
|
|
|
|
model, preprocess = clip.load("ViT-B/32")
|
|
|
|
input_resolution = model.visual.input_resolution
|
|
context_length = model.context_length
|
|
vocab_size = model.vocab_size
|
|
|
|
print("Model parameters:", f"{np.sum([int(np.prod(p.shape)) for p in model.parameters()]):,}")
|
|
print("Input resolution:", input_resolution)
|
|
print("Context length:", context_length)
|
|
print("Vocab size:", vocab_size)
|
|
|
|
import json
|
|
imagenet_data = json.loads(open("imagenet_data.json","r").read())
|
|
imagenet_classes = imagenet_data['imagenet_classes']
|
|
imagenet_templates = imagenet_data['imagenet_templates']
|
|
|
|
print(f"{len(imagenet_classes)} classes, {len(imagenet_templates)} templates")
|
|
|
|
# execute:
|
|
# ! pip install git+https://github.com/modestyachts/ImageNetV2_pytorch
|
|
|
|
from imagenetv2_pytorch import ImageNetV2Dataset
|
|
|
|
images = ImageNetV2Dataset(transform=preprocess)
|
|
loader = torch.utils.data.DataLoader(images, batch_size=32, num_workers=2)
|
|
|
|
def zeroshot_classifier(classnames, templates):
|
|
with torch.no_grad():
|
|
zeroshot_weights = []
|
|
for classname in tqdm(classnames):
|
|
texts = [template.format(classname) for template in templates] #format with class
|
|
texts = clip.tokenize(texts).cuda() #tokenize
|
|
class_embeddings = model.encode_text(texts) #embed with text encoder
|
|
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).cuda()
|
|
return zeroshot_weights
|
|
|
|
|
|
zeroshot_weights = zeroshot_classifier(imagenet_classes, imagenet_templates)
|
|
|
|
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]
|
|
|
|
with torch.no_grad():
|
|
top1, top5, n = 0., 0., 0.
|
|
for i, (images, target) in enumerate(tqdm(loader)):
|
|
images = images.cuda()
|
|
target = target.cuda()
|
|
|
|
# predict
|
|
image_features = model.encode_image(images)
|
|
image_features /= image_features.norm(dim=-1, keepdim=True)
|
|
logits = 100. * image_features @ zeroshot_weights
|
|
|
|
# measure accuracy
|
|
acc1, acc5 = accuracy(logits, target, topk=(1, 5))
|
|
top1 += acc1
|
|
top5 += acc5
|
|
n += images.size(0)
|
|
|
|
top1 = (top1 / n) * 100
|
|
top5 = (top5 / n) * 100
|
|
|
|
print(f"Top-1 accuracy: {top1:.2f}")
|
|
print(f"Top-5 accuracy: {top5:.2f}")
|