bioclip / examples /zero_shot.py

Samuel Stevens

Add zero-shot example.

7b4abf1 6 months ago

No virus

10.3 kB

	"""
	Do zero-shot image classification.

	Writes the output to a plaintext and JSON format in the logs directory.
	"""
	import argparse
	import ast
	import contextlib
	import json
	import logging
	import os
	import random
	import sys

	import numpy as np
	import open_clip
	import torch
	import torch.nn.functional as F
	from torchvision import datasets
	from tqdm import tqdm

	log_format = "[%(asctime)s] [%(levelname)s] [%(name)s] %(message)s"
	logging.basicConfig(level=logging.INFO, format=log_format)
	logger = logging.getLogger("main")

	openai_templates = [
	lambda c: f"a bad photo of a {c}.",
	lambda c: f"a photo of many {c}.",
	lambda c: f"a sculpture of a {c}.",
	lambda c: f"a photo of the hard to see {c}.",
	lambda c: f"a low resolution photo of the {c}.",
	lambda c: f"a rendering of a {c}.",
	lambda c: f"graffiti of a {c}.",
	lambda c: f"a bad photo of the {c}.",
	lambda c: f"a cropped photo of the {c}.",
	lambda c: f"a tattoo of a {c}.",
	lambda c: f"the embroidered {c}.",
	lambda c: f"a photo of a hard to see {c}.",
	lambda c: f"a bright photo of a {c}.",
	lambda c: f"a photo of a clean {c}.",
	lambda c: f"a photo of a dirty {c}.",
	lambda c: f"a dark photo of the {c}.",
	lambda c: f"a drawing of a {c}.",
	lambda c: f"a photo of my {c}.",
	lambda c: f"the plastic {c}.",
	lambda c: f"a photo of the cool {c}.",
	lambda c: f"a close-up photo of a {c}.",
	lambda c: f"a black and white photo of the {c}.",
	lambda c: f"a painting of the {c}.",
	lambda c: f"a painting of a {c}.",
	lambda c: f"a pixelated photo of the {c}.",
	lambda c: f"a sculpture of the {c}.",
	lambda c: f"a bright photo of the {c}.",
	lambda c: f"a cropped photo of a {c}.",
	lambda c: f"a plastic {c}.",
	lambda c: f"a photo of the dirty {c}.",
	lambda c: f"a jpeg corrupted photo of a {c}.",
	lambda c: f"a blurry photo of the {c}.",
	lambda c: f"a photo of the {c}.",
	lambda c: f"a good photo of the {c}.",
	lambda c: f"a rendering of the {c}.",
	lambda c: f"a {c} in a video game.",
	lambda c: f"a photo of one {c}.",
	lambda c: f"a doodle of a {c}.",
	lambda c: f"a close-up photo of the {c}.",
	lambda c: f"a photo of a {c}.",
	lambda c: f"the origami {c}.",
	lambda c: f"the {c} in a video game.",
	lambda c: f"a sketch of a {c}.",
	lambda c: f"a doodle of the {c}.",
	lambda c: f"a origami {c}.",
	lambda c: f"a low resolution photo of a {c}.",
	lambda c: f"the toy {c}.",
	lambda c: f"a rendition of the {c}.",
	lambda c: f"a photo of the clean {c}.",
	lambda c: f"a photo of a large {c}.",
	lambda c: f"a rendition of a {c}.",
	lambda c: f"a photo of a nice {c}.",
	lambda c: f"a photo of a weird {c}.",
	lambda c: f"a blurry photo of a {c}.",
	lambda c: f"a cartoon {c}.",
	lambda c: f"art of a {c}.",
	lambda c: f"a sketch of the {c}.",
	lambda c: f"a embroidered {c}.",
	lambda c: f"a pixelated photo of a {c}.",
	lambda c: f"itap of the {c}.",
	lambda c: f"a jpeg corrupted photo of the {c}.",
	lambda c: f"a good photo of a {c}.",
	lambda c: f"a plushie {c}.",
	lambda c: f"a photo of the nice {c}.",
	lambda c: f"a photo of the small {c}.",
	lambda c: f"a photo of the weird {c}.",
	lambda c: f"the cartoon {c}.",
	lambda c: f"art of the {c}.",
	lambda c: f"a drawing of the {c}.",
	lambda c: f"a photo of the large {c}.",
	lambda c: f"a black and white photo of a {c}.",
	lambda c: f"the plushie {c}.",
	lambda c: f"a dark photo of a {c}.",
	lambda c: f"itap of a {c}.",
	lambda c: f"graffiti of the {c}.",
	lambda c: f"a toy {c}.",
	lambda c: f"itap of my {c}.",
	lambda c: f"a photo of a cool {c}.",
	lambda c: f"a photo of a small {c}.",
	lambda c: f"a tattoo of the {c}.",
	]


	def parse_args(args):
	class ParseKwargs(argparse.Action):
	def __call__(self, parser, namespace, values, option_string=None):
	kw = {}
	for value in values:
	key, value = value.split("=")
	try:
	kw[key] = ast.literal_eval(value)
	except (ValueError, SyntaxError):
	# fallback to string (avoid need to escape on command line)
	kw[key] = str(value)
	setattr(namespace, self.dest, kw)

	parser = argparse.ArgumentParser()
	parser.add_argument(
	"--datasets",
	type=str,
	default=None,
	nargs="+",
	help="Path to dirs(s) with validation data. In the format NAME=PATH.",
	action=ParseKwargs,
	)
	parser.add_argument(
	"--logs", type=str, default="./logs", help="Where to write logs"
	)
	parser.add_argument(
	"--exp", type=str, default="bioclip-zero-shot", help="Experiment name."
	)
	parser.add_argument(
	"--workers", type=int, default=8, help="Number of dataloader workers per GPU."
	)
	parser.add_argument(
	"--batch-size", type=int, default=64, help="Batch size per GPU."
	)
	parser.add_argument(
	"--precision",
	choices=["amp", "amp_bf16", "amp_bfloat16", "bf16", "fp32"],
	default="amp",
	help="Floating point precision.",
	)
	parser.add_argument("--seed", type=int, default=0, help="Default random seed.")
	args = parser.parse_args(args)
	os.makedirs(os.path.join(args.logs, args.exp), exist_ok=True)

	return args


	def make_txt_features(model, classnames, templates, args):
	tokenizer = open_clip.get_tokenizer("hf-hub:imageomics/bioclip")
	with torch.no_grad():
	txt_features = []
	for classname in tqdm(classnames):
	classname = " ".join(word for word in classname.split("_") if word)
	texts = [template(classname) for template in templates] # format with class
	texts = tokenizer(texts).to(args.device) # tokenize
	class_embeddings = model.encode_text(texts)
	class_embedding = F.normalize(class_embeddings, dim=-1).mean(dim=0)
	class_embedding /= class_embedding.norm()
	txt_features.append(class_embedding)
	txt_features = torch.stack(txt_features, dim=1).to(args.device)
	return txt_features


	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 [correct[:k].reshape(-1).float().sum(0, keepdim=True).item() for k in topk]


	def get_autocast(precision):
	if precision == "amp":
	return torch.cuda.amp.autocast
	elif precision == "amp_bfloat16" or precision == "amp_bf16":
	# amp_bfloat16 is more stable than amp float16 for clip training
	return lambda: torch.cuda.amp.autocast(dtype=torch.bfloat16)
	else:
	return contextlib.suppress


	def run(model, txt_features, dataloader, args):
	autocast = get_autocast(args.precision)
	cast_dtype = open_clip.get_cast_dtype(args.precision)

	top1, top5, n = 0.0, 0.0, 0.0

	with torch.no_grad():
	for images, targets in tqdm(dataloader, unit_scale=args.batch_size):
	images = images.to(args.device)
	if cast_dtype is not None:
	images = images.to(dtype=cast_dtype)
	targets = targets.to(args.device)

	with autocast():
	image_features = model.encode_image(images)
	image_features = F.normalize(image_features, dim=-1)
	logits = model.logit_scale.exp() * image_features @ txt_features

	# Measure accuracy
	acc1, acc5 = accuracy(logits, targets, topk=(1, 5))
	top1 += acc1
	top5 += acc5
	n += images.size(0)

	top1 = top1 / n
	top5 = top5 / n
	return top1, top5


	def evaluate(model, data, args):
	results = {}

	logger.info("Starting zero-shot classification.")

	for split in data:
	logger.info("Building zero-shot %s classifier.", split)

	classnames = data[split].dataset.classes
	classnames = [name.replace("_", " ") for name in classnames]

	txt_features = make_txt_features(model, classnames, openai_templates, args)

	logger.info("Got text features.")
	top1, top5 = run(model, txt_features, data[split], args)

	logger.info("%s-top1: %.3f", split, top1 * 100)
	logger.info("%s-top5: %.3f", split, top5 * 100)

	results[f"{split}-top1"] = top1 * 100
	results[f"{split}-top5"] = top5 * 100

	logger.info("Finished zero-shot %s.", split)

	logger.info("Finished zero-shot classification.")

	return results


	if __name__ == "__main__":
	args = parse_args(sys.argv[1:])

	if torch.cuda.is_available():
	# This enables tf32 on Ampere GPUs which is only 8% slower than
	# float16 and almost as accurate as float32
	# This was a default in pytorch until 1.12
	torch.backends.cuda.matmul.allow_tf32 = True
	torch.backends.cudnn.benchmark = True
	torch.backends.cudnn.deterministic = False

	# Init torch device
	if torch.cuda.is_available():
	device = "cuda:0"
	torch.cuda.set_device(device)
	else:
	device = "cpu"
	args.device = device

	# Random seeding
	torch.manual_seed(args.seed)
	np.random.seed(args.seed)
	random.seed(args.seed)

	# Load model.
	model, preprocess_train, preprocess_val = open_clip.create_model_and_transforms(
	"hf-hub:imageomics/bioclip"
	)

	# Write datasets
	params_file = os.path.join(args.logs, args.exp, "params.json")
	with open(params_file, "w") as fd:
	params = {name: getattr(args, name) for name in vars(args)}
	json.dump(params, fd, sort_keys=True, indent=4)

	# Initialize datasets.
	data = {}
	for split, path in args.datasets.items():
	data[split] = torch.utils.data.DataLoader(
	datasets.ImageFolder(path, transform=preprocess_val),
	batch_size=args.batch_size,
	num_workers=args.workers,
	sampler=None,
	shuffle=False,
	)

	model.eval()
	results = evaluate(model, data, args)

	results_file = os.path.join(args.logs, args.exp, "results.json")
	with open(results_file, "w") as fd:
	json.dump(results, fd, indent=4, sort_keys=True)