hgnetv2_b2_eval_quantized.py

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