vit_large_patch16_224_evaluate.py

"""
Evaluate google/vit-large-patch16-224 on ImageNet-1k validation set.

We use the HuggingFace Transformers ViT model with the google/vit-large-patch16-224
checkpoint (pretrained on ImageNet-21k, fine-tuned on ImageNet-1k at 224px).

We compute Top-1 / Top-5 accuracy, per-class accuracy, and mAP over the
full 50K ImageNet validation images.

Dataset:  Tsomaros/Imagenet-1k_validation  (validation split, 50K images)
          label indices 0-999 align with the model's output indices.

Usage:
    python vit_large_evaluate.py [--batch_size 64] [--num_workers 4] [--streaming]
"""

import argparse
import os
import time
import numpy as np
import torch
import torch.onnx
from torch.utils.data import Dataset, DataLoader
from datasets import load_dataset
from transformers import ViTForImageClassification, ViTImageProcessor
from sklearn.metrics import average_precision_score, precision_recall_fscore_support


# ---------------------------------------------------------------------------
# Dataset wrapper
# ---------------------------------------------------------------------------

class ImageNetValDataset(Dataset):
    """Wrap the HF ImageNet-1k validation split for evaluation."""

    def __init__(self, hf_dataset, transform=None):
        self.hf_dataset = hf_dataset
        self.transform = transform

    def __len__(self):
        return len(self.hf_dataset)

    def __getitem__(self, idx):
        row = self.hf_dataset[idx]
        img = row["image"]
        if img.mode != "RGB":
            img = img.convert("RGB")
        label = row["label"]
        if self.transform:
            img = self.transform(img)
        return img, label


# ---------------------------------------------------------------------------
# Metrics
# ---------------------------------------------------------------------------

def compute_metrics(logits: np.ndarray, labels: np.ndarray, num_classes: int):
    """
    Compute comprehensive single-label classification metrics.

    Args:
        logits:      (N, C) raw model outputs
        labels:      (N,)   integer ground-truth class indices
        num_classes: total number of classes (1000)

    Returns:
        dict with top1, top5, mAP, precision, recall, f1, per_class_ap
    """
    probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy()  # (N, C)
    preds = probs.argmax(axis=1)

    N = len(labels)

    # --- Top-1 / Top-5 Accuracy ---
    top1_correct = (preds == labels).sum()
    top1 = top1_correct / N

    topk_vals = np.argsort(probs, axis=1)[:, ::-1]  # (N, C) sorted descending
    top5_correct = sum(labels[i] in topk_vals[i, :5] for i in range(N))
    top5 = top5_correct / N

    # --- mAP (per-class AP averaged) ---
    one_hot = np.zeros((N, num_classes), dtype=np.int32)
    one_hot[np.arange(N), labels] = 1

    aps = []
    per_class_ap = np.zeros(num_classes, dtype=np.float64)
    for c in range(num_classes):
        if one_hot[:, c].sum() == 0:
            continue
        try:
            ap = average_precision_score(one_hot[:, c], probs[:, c])
        except ValueError:
            ap = 0.0
        per_class_ap[c] = ap
        aps.append(ap)

    mAP = np.mean(aps) if aps else 0.0

    # --- Precision / Recall / F1 ---
    prec_mac, rec_mac, f1_mac, _ = precision_recall_fscore_support(
        labels, preds, average="macro", zero_division=0
    )
    prec_wt, rec_wt, f1_wt, _ = precision_recall_fscore_support(
        labels, preds, average="weighted", zero_division=0
    )

    return {
        "top1": top1,
        "top5": top5,
        "mAP": mAP,
        "precision_macro": prec_mac,
        "recall_macro": rec_mac,
        "f1_macro": f1_mac,
        "precision_weighted": prec_wt,
        "recall_weighted": rec_wt,
        "f1_weighted": f1_wt,
        "per_class_ap": per_class_ap,
    }


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------

@torch.no_grad()
def evaluate(model, loader, device, num_classes, print_every=500):
    """Run full evaluation and return metrics dict."""
    model.eval()
    all_logits = []
    all_labels = []
    total = 0
    total_inference_time = 0.0  # strict inference-only timing

    start = time.time()  # process timer (includes data transfer, inference, output, metrics)
    for batch_idx, (imgs, labels) in enumerate(loader):
        imgs = imgs.to(device)
        t0 = time.perf_counter()
        outputs = model(pixel_values=imgs)
        # ViTForImageClassification returns logits in outputs.logits
        logits = outputs.logits
        t1 = time.perf_counter()
        total_inference_time += (t1 - t0)
        all_logits.append(logits.cpu().numpy())
        all_labels.append(np.array(labels))
        total += imgs.size(0)

        if print_every and (batch_idx + 1) % print_every == 0:
            elapsed = time.time() - start
            speed = total / elapsed
            print(f"  [{total:>6d} images]  {speed:.1f} img/s")

    all_logits = np.concatenate(all_logits, axis=0)
    all_labels = np.concatenate(all_labels, axis=0)

    metrics = compute_metrics(all_logits, all_labels, num_classes)
    metrics["total_images"] = total
    elapsed = time.time() - start
    metrics["elapsed"] = elapsed
    metrics["avg_process_ms"] = elapsed / total * 1000 if total > 0 else 0.0
    metrics["avg_inference_ms"] = total_inference_time / total * 1000 if total > 0 else 0.0
    return metrics


# ---------------------------------------------------------------------------
# Calibration data
# ---------------------------------------------------------------------------

def save_calibration_data(processor, k=20, save_path="vit_large_patch16_224_calibration.npy"):
    """
    Randomly sample k images from Tsomaros/Imagenet-1k_validation validation split as
    calibration data for model quantization, and save as a .npy file.

    Args:
        processor: ViTImageProcessor for preprocessing
        k:         number of calibration images to sample
        save_path: output .npy file path
    """
    input_size = (3, 224, 224)  # (C, H, W)
    print(f"  Input size: {input_size}")

    # Load dataset
    print("Loading Tsomaros/Imagenet-1k_validation validation ...")
    ds = load_dataset("Tsomaros/Imagenet-1k_validation", split="validation")
    total = len(ds)
    print(f"  Total images: {total}")

    # Random sample k indices
    rng = np.random.default_rng(42)
    indices = rng.choice(total, size=min(k, total), replace=False)
    indices.sort()  # sorted for sequential HF dataset access

    # Collect calibration images
    calibration = np.empty((len(indices), *input_size), dtype=np.float32)
    for i, idx in enumerate(indices):
        row = ds[int(idx)]
        img = row["image"]
        if img.mode != "RGB":
            img = img.convert("RGB")
        inputs = processor(images=img, return_tensors="pt")
        tensor = inputs["pixel_values"].squeeze(0)  # (C, H, W) float32
        calibration[i] = tensor.numpy()
        if (i + 1) % 20 == 0 or (i + 1) == len(indices):
            print(f"  [{i + 1:>4d}/{len(indices)}] images collected")

    np.save(save_path, calibration)
    print(f"Saved calibration data: {calibration.shape} -> {save_path}")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def main():
    parser = argparse.ArgumentParser(
        description="Evaluate google/vit-large-patch16-224 on ImageNet-1k val"
    )
    parser.add_argument("--batch_size", type=int, default=64, help="Batch size")
    parser.add_argument("--num_workers", type=int, default=4, help="DataLoader workers")
    parser.add_argument("--streaming", action="store_true", help="Stream dataset instead of downloading")
    parser.add_argument("--subset", type=int, default=0, help="Evaluate on first N images only (0 = all 50K)")
    parser.add_argument("--save_model", action="store_true", help="Save pth, onnx, and calibration data")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device: {device}")

    # ------------------------------------------------------------------
    # Load model & processor
    # ------------------------------------------------------------------
    model_name = "google/vit-large-patch16-224"
    print(f"Loading {model_name} ...")

    processor = ViTImageProcessor.from_pretrained(model_name)
    model = ViTForImageClassification.from_pretrained(model_name)
    model = model.to(device)
    model.eval()

    num_classes = model.config.num_labels  # 1000
    input_size = (1, 3, 224, 224)  # (B, C, H, W)
    print(f"  Classes: {num_classes}, Input size: {input_size}")

    # Build a transform callable from the HF processor for use in DataLoader
    def transform(img):
        inputs = processor(images=img, return_tensors="pt")
        return inputs["pixel_values"].squeeze(0)  # (C, H, W)

    if args.save_model:
        inputs = torch.randn(input_size).to(device)
        torch.save(model.state_dict(), 'vit_large_patch16_224_fp32.pth')

        torch.onnx.export(
            model,
            (inputs,),
            "vit_large_patch16_224_fp32.onnx",
            input_names=["pixel_values"],
            output_names=["logits"],
            dynamic_axes={
                "pixel_values": {0: "batch_size"},
                "logits": {0: "batch_size"},
            },
        )

        save_calibration_data(processor, 10)

    # ------------------------------------------------------------------
    # Load dataset
    # ------------------------------------------------------------------
    print("Loading Tsomaros/Imagenet-1k_validation validation split ...")
    ds = None
    try:
        ds = load_dataset(
            "Tsomaros/Imagenet-1k_validation",
            split="validation",
            streaming=args.streaming,
        )
    except Exception as e:
        print(f"  load_dataset failed: {type(e).__name__}: {str(e)[:150]}")
        print("  Falling back to loading from cached arrow shard files ...")
        ds = None

    if ds is None:
        # Fallback: load directly from arrow shard files
        import io as _io
        import pyarrow.ipc as _ipc
        from PIL import Image as _PILImage

        # Try to locate the cached arrow directory for Tsomaros/Imagenet-1k_validation
        _cache_base = os.path.expanduser("~/.cache/huggingface/datasets/")
        _arrow_dir = None
        # Search for the ILSVRC imagenet-1k dataset cache
        for root, dirs, files in os.walk(_cache_base):
            if "imagenet" in root.lower() and any(f.endswith(".arrow") for f in files):
                _arrow_dir = root
                break

        if _arrow_dir is None or not os.path.isdir(_arrow_dir):
            print(f"  ERROR: Arrow cache directory not found under {_cache_base}")
            print("  Please run with network access first to download the dataset.")
            return

        _shard_files = sorted(
            f for f in os.listdir(_arrow_dir)
            if f.endswith(".arrow")
        )
        print(f"  Found {len(_shard_files)} arrow shard files in {_arrow_dir}")

        class _ArrowShardDataset(Dataset):
            """Load ImageNet validation data from cached arrow shard files."""

            def __init__(self, arrow_dir, shard_files, transform=None):
                self.transform = transform
                self.shards = []
                self.offsets = [0]
                for fname in shard_files:
                    path = os.path.join(arrow_dir, fname)
                    try:
                        with open(path, "rb") as f:
                            reader = _ipc.RecordBatchStreamReader(f)
                            table = reader.read_all()
                        self.shards.append(table)
                        self.offsets.append(self.offsets[-1] + len(table))
                    except Exception as ex:
                        print(f"    SKIP shard {fname}: {ex}")
                self.total = self.offsets[-1]
                print(f"    Total images loaded from shards: {self.total}")

            def __len__(self):
                return self.total

            def __getitem__(self, idx):
                lo, hi = 0, len(self.shards) - 1
                while lo < hi:
                    mid = (lo + hi) // 2
                    if self.offsets[mid + 1] <= idx:
                        lo = mid + 1
                    else:
                        hi = mid
                shard_idx = lo
                local_idx = idx - self.offsets[shard_idx]
                table = self.shards[shard_idx]
                img_bytes = table.column("image")[local_idx].as_py()
                if isinstance(img_bytes, dict):
                    img_bytes = img_bytes.get("bytes", img_bytes.get("path", b""))
                if isinstance(img_bytes, bytes):
                    img = _PILImage.open(_io.BytesIO(img_bytes)).convert("RGB")
                else:
                    img = _PILImage.new("RGB", (224, 224))
                label = table.column("label")[local_idx].as_py()
                if self.transform:
                    img = self.transform(img)
                return img, label

        dataset = _ArrowShardDataset(_arrow_dir, _shard_files, transform=transform)
        total_images = len(dataset)
        if args.subset > 0:
            from torch.utils.data import Subset
            dataset = Subset(dataset, range(min(args.subset, total_images)))
            total_images = min(args.subset, total_images)
            print(f"  Using subset: {total_images} images")
    elif args.streaming:
        if args.subset > 0:
            from itertools import islice
            ds = list(islice(ds, args.subset))
            print(f"  Using subset: {args.subset} images (streamed)")

        # Materialize for DataLoader compatibility
        if not isinstance(ds, list):
            ds = list(ds)
        from torch.utils.data import Dataset as _DS

        class _StreamDS(_DS):
            def __init__(self, items, tfm):
                self.items = items
                self.tfm = tfm

            def __len__(self):
                return len(self.items)

            def __getitem__(self, idx):
                row = self.items[idx]
                img = row["image"]
                if img.mode != "RGB":
                    img = img.convert("RGB")
                return self.tfm(img), row["label"]

        dataset = _StreamDS(ds, transform)
        total_images = len(ds)
    else:
        if args.subset > 0:
            ds = ds.select(range(args.subset))
            print(f"  Using subset: {args.subset} images")
        dataset = ImageNetValDataset(ds, transform=transform)
        total_images = len(dataset)

    print(f"  Total images: {total_images}")

    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True,
    )

    # ------------------------------------------------------------------
    # Evaluate
    # ------------------------------------------------------------------
    print(f"\n{'='*60}")
    print("Running ImageNet-1k validation evaluation ...")
    print(f"{'='*60}\n")

    metrics = evaluate(model, loader, device, num_classes=num_classes)

    # ------------------------------------------------------------------
    # Print results
    # ------------------------------------------------------------------
    print(f"\n{'='*60}")
    print(f"Evaluation Results — google/vit-large-patch16-224")
    print(f"{'='*60}")
    print(f"\n  Dataset:         ImageNet-1k validation ({metrics['total_images']} images)")
    print(f"  Top-1 Accuracy:  {metrics['top1']*100:.3f}%")
    print(f"  Top-5 Accuracy:  {metrics['top5']*100:.3f}%")
    print(f"  mAP:             {metrics['mAP']:.4f}")
    print(f"\n  Precision (macro):    {metrics['precision_macro']:.4f}")
    print(f"  Recall (macro):       {metrics['recall_macro']:.4f}")
    print(f"  F1 (macro):           {metrics['f1_macro']:.4f}")
    print(f"  Precision (weighted): {metrics['precision_weighted']:.4f}")
    print(f"  Recall (weighted):    {metrics['recall_weighted']:.4f}")
    print(f"  F1 (weighted):        {metrics['f1_weighted']:.4f}")
    print(f"\n  Total Time:           {metrics['elapsed']:.1f}s")
    print(f"  Avg Process Time:     {metrics['avg_process_ms']:.2f}ms/img")
    print(f"  Avg Inference Time:   {metrics['avg_inference_ms']:.2f}ms/img")

    # Per-class AP summary
    per_class_ap = metrics["per_class_ap"]
    valid_mask = per_class_ap > 0
    valid_aps = per_class_ap[valid_mask]
    if len(valid_aps) > 0:
        print(f"\n  Per-class AP (over {len(valid_aps)} classes present in val):")
        print(f"    Mean:   {valid_aps.mean():.4f}")
        print(f"    Median: {np.median(valid_aps):.4f}")
        print(f"    Min:    {valid_aps.min():.4f}")
        print(f"    Max:    {valid_aps.max():.4f}")
        print(f"    Std:    {valid_aps.std():.4f}")

        # Top-10 / Bottom-10 classes by AP
        present_indices = np.where(valid_mask)[0]
        sorted_by_ap = present_indices[np.argsort(per_class_ap[present_indices])]

        # Load class names from model config id2label
        id2label = model.config.id2label

        def _class_label(idx):
            name = id2label.get(idx, f"class_{idx}")
            return name[:40]

        print(f"\n  Top-10 classes by AP:")
        for idx in sorted_by_ap[-10:][::-1]:
            print(f"    [{idx:>3d}] {_class_label(idx):42s} AP = {per_class_ap[idx]:.4f}")
        print(f"\n  Bottom-10 classes by AP:")
        for idx in sorted_by_ap[:10]:
            print(f"    [{idx:>3d}] {_class_label(idx):42s} AP = {per_class_ap[idx]:.4f}")

    # Reference accuracy from the model card
    print(f"\n{'='*60}")
    print("Reference (google/vit-large-patch16-224):")
    print("  Top-1: 85.14%  |  Top-5: 97.64%  (224px val crop)")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()