examples/zero_shot.py

"""
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)