""" Evaluate quantized ONNX models on ImageNet-1k validation set. Uses ONNX Runtime for inference. Loads the cached ImageNet dataset directly from arrow shard files. """ import argparse import os import time import io import numpy as np import onnxruntime as ort import torch from torch.utils.data import Dataset, DataLoader from PIL import Image import pyarrow.ipc as ipc import timm from timm.data import resolve_model_data_config, create_transform from sklearn.metrics import average_precision_score, precision_recall_fscore_support # --------------------------------------------------------------------------- # Dataset that reads directly from arrow shards # --------------------------------------------------------------------------- class ArrowImageNetDataset(Dataset): """Load ImageNet validation data from cached arrow shard files.""" def __init__(self, arrow_dir, transform=None): self.transform = transform self.shards = [] self.offsets = [0] # Load all valid arrow shards shard_files = sorted( f for f in os.listdir(arrow_dir) if f.startswith("imagenet-1k_validation-validation-") and f.endswith(".arrow") ) 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)) print(f" Loaded shard {fname}: {len(table)} rows") except Exception as e: print(f" SKIP shard {fname}: {e}") self.total = self.offsets[-1] print(f" Total images: {self.total}") def __len__(self): return self.total def __getitem__(self, idx): # Binary search for the correct shard 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 = Image.open(io.BytesIO(img_bytes)).convert("RGB") else: img = Image.new("RGB", (224, 224)) label = table.column("label")[local_idx].as_py() if self.transform: img = self.transform(img) return img, label # --------------------------------------------------------------------------- # Metrics (same as model_eval_test.py) # --------------------------------------------------------------------------- def compute_metrics(logits, labels, num_classes): probs = torch.softmax(torch.from_numpy(logits), dim=1).numpy() preds = probs.argmax(axis=1) N = len(labels) top1 = (preds == labels).sum() / N topk_vals = np.argsort(probs, axis=1)[:, ::-1] top5 = sum(labels[i] in topk_vals[i, :5] for i in range(N)) / N one_hot = np.zeros((N, num_classes), dtype=np.int32) one_hot[np.arange(N), labels] = 1 aps = [] 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 aps.append(ap) mAP = np.mean(aps) if aps else 0.0 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, } # --------------------------------------------------------------------------- # ONNX model evaluation # --------------------------------------------------------------------------- def evaluate_onnx(onnx_path, loader, num_classes, print_every=500): """Evaluate an ONNX model using ONNX Runtime with batch=1 (models have static shapes).""" providers = [] if "CUDAExecutionProvider" in ort.get_available_providers(): providers.append(("CUDAExecutionProvider", {"device_id": 0})) providers.append("CPUExecutionProvider") session_options = ort.SessionOptions() session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL session = ort.InferenceSession(onnx_path, sess_options=session_options, providers=providers) input_name = session.get_inputs()[0].name all_logits = [] all_labels = [] total = 0 total_inference_time = 0.0 # strict inference-only timing start = time.time() # process timer (includes data prep, inference, output, metrics) for batch_idx, (imgs, labels) in enumerate(loader): # Run one image at a time (model has static batch=1 in internal Reshape nodes) for i in range(imgs.size(0)): single_img = imgs[i:i+1].numpy() # shape (1, C, H, W) t0 = time.perf_counter() outputs = session.run(None, {input_name: single_img}) t1 = time.perf_counter() total_inference_time += (t1 - t0) all_logits.append(outputs[0]) all_labels.append(np.array([labels[i].item()] if torch.is_tensor(labels[i]) else [labels[i]])) 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 # --------------------------------------------------------------------------- # Main # --------------------------------------------------------------------------- def main(): parser = argparse.ArgumentParser(description="Evaluate quantized ONNX models") parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--num_workers", type=int, default=4) parser.add_argument("--subset", type=int, default=0, help="Evaluate on first N images (0=all)") ALL_MODES = ["fp32", "fp16", "int8", "fp8", "int4"] parser.add_argument( "--mode", type=str, nargs="*", default=ALL_MODES, choices=ALL_MODES, help=f"Quantization mode(s) to evaluate (default: all). Choices: {ALL_MODES}", ) args = parser.parse_args() # Load model config for transforms (pretrained=False to avoid HF download) print("Loading model config for transforms...") model = timm.create_model("hgnetv2_b2.ssld_stage2_ft_in1k", pretrained=False) data_config = resolve_model_data_config(model) transform = create_transform(**data_config, is_training=False) num_classes = 1000 del model # Load dataset from cached arrow shards arrow_dir = os.path.expanduser( "~/.cache/huggingface/datasets/Tsomaros___imagenet-1k_validation/" "default/0.0.0/55405c49dece42420e68ddd5f80174f19b29ebaf/" ) print(f"Loading dataset from arrow shards: {arrow_dir}") dataset = ArrowImageNetDataset(arrow_dir, transform=transform) if args.subset > 0: from torch.utils.data import Subset dataset = Subset(dataset, range(min(args.subset, len(dataset)))) print(f" Using subset: {args.subset} images") loader = DataLoader( dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True, ) # Define models to evaluate models = { "FP32 (baseline)": "hgnetv2_b2_fp32.onnx", "FP16": "fp16/hgnetv2_b2_fp16.onnx", "INT8 entropy": "int8/hgnetv2_b2_int8_entropy.onnx", "INT8 max": "int8/hgnetv2_b2_int8_max.onnx", "FP8 entropy": "fp8/hgnetv2_b2_fp8_entropy.onnx", "FP8 max": "fp8/hgnetv2_b2_fp8_max.onnx", "INT4 awq_clip": "int4/hgnetv2_b2_int4_awq_clip.onnx", "INT4 awq_lite (asym)": "int4/hgnetv2_b2_int4_awq_lite_asym.onnx", "INT4 awq_lite (sym)": "int4/hgnetv2_b2_int4_awq_lite.onnx", "INT4 awq_full": "int4/hgnetv2_b2_int4_awq_full.onnx", "INT4 rtn_dq": "int4/hgnetv2_b2_int4_rtn_dq.onnx", } # Filter by --mode selection mode_prefix = {"fp32": "FP32", "fp16": "FP16", "int8": "INT8", "fp8": "FP8", "int4": "INT4"} selected_prefixes = {mode_prefix[m] for m in args.mode} models = {k: v for k, v in models.items() if any(k.startswith(p) for p in selected_prefixes)} print(f"Evaluating modes: {args.mode}") # Filter to only existing files existing_models = {} for name, path in models.items(): if os.path.exists(path): existing_models[name] = path else: print(f" SKIP: {name} — file not found: {path}") results = {} for name, onnx_path in existing_models.items(): print(f"\n{'='*60}") print(f"Evaluating: {name}") print(f" Model: {onnx_path}") print(f"{'='*60}") try: metrics = evaluate_onnx(onnx_path, loader, num_classes) results[name] = metrics print(f"\n Top-1 Accuracy: {metrics['top1']*100:.3f}%") print(f" Top-5 Accuracy: {metrics['top5']*100:.3f}%") print(f" mAP: {metrics['mAP']:.4f}") print(f" F1 (macro): {metrics['f1_macro']:.4f}") print(f" F1 (weighted): {metrics['f1_weighted']:.4f}") print(f" Time: {metrics['elapsed']:.1f}s") print(f" Avg Process: {metrics['avg_process_ms']:.2f}ms/img") print(f" Avg Inference: {metrics['avg_inference_ms']:.2f}ms/img") except Exception as e: print(f" FAILED: {e}") import traceback traceback.print_exc() results[name] = {"error": str(e)} # Print comparison table print(f"\n\n{'='*100}") print("Evaluation Comparison Table") print(f"{'='*100}") print(f" {'Model':<25s} {'Images':>7s} {'Top-1%':>8s} {'Top-5%':>8s} {'mAP':>8s} {'F1_mac':>8s} {'F1_wt':>8s} {'Proc(ms)':>9s} {'Inf(ms)':>8s} {'Time':>8s}") print(f" {'-'*25} {'-'*7} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*9} {'-'*8} {'-'*8}") for name, m in results.items(): if "error" in m: print(f" {name:<25s} FAILED: {m['error']}") else: print( f" {name:<25s} " f"{m['total_images']:>7d} " f"{m['top1']*100:>8.3f} " f"{m['top5']*100:>8.3f} " f"{m['mAP']:>8.4f} " f"{m['f1_macro']:>8.4f} " f"{m['f1_weighted']:>8.4f} " f"{m['avg_process_ms']:>9.2f} " f"{m['avg_inference_ms']:>8.2f} " f"{m['elapsed']:>7.1f}s" ) print(f"\n Reference (timm model card): Top-1: 82.346% | Top-5: 96.394%") print(f"{'='*90}") # Find best INT8 model and copy as the canonical output int8_results = {k: v for k, v in results.items() if k.startswith("INT8") and "error" not in v} if int8_results: best_int8 = max(int8_results, key=lambda k: int8_results[k]["top1"]) best_path = existing_models[best_int8] print(f"\n Best INT8 model: {best_int8} ({best_path})") print(f" Top-1: {int8_results[best_int8]['top1']*100:.3f}%") print(f" Top-5: {int8_results[best_int8]['top5']*100:.3f}%") if __name__ == "__main__": main()