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