feat: finished with regnet_x_800mf PTQ scripts.

regnet_x_800mf_eval_quantized.py

@@ -0,0 +1,306 @@
+"""
+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
+from sklearn.metrics import average_precision_score, precision_recall_fscore_support
+import torchvision
+from torchvision.models import RegNet_X_800MF_Weights
+
+
+# ---------------------------------------------------------------------------
+# 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
+
+    start = time.time()
+    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)
+            outputs = session.run(None, {input_name: single_img})
+            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
+    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)")
+    args = parser.parse_args()
+
+    # Load model config for transforms
+    print("Loading model config for transforms...")
+    weights = RegNet_X_800MF_Weights.IMAGENET1K_V2
+    # model = torchvision.models.regnet_x_800mf(weights=weights)
+    # model.eval()
+
+    transform = weights.transforms()
+    input_size = (3, 224, 224)  # (C, H, W)
+    
+    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)": "regnet_x_800mf_fp32.onnx",
+        "FP16": "fp16/regnet_x_800mf_fp16.onnx",
+        "INT8 entropy": "int8/regnet_x_800mf_int8_entropy.onnx",
+        "INT8 max": "int8/regnet_x_800mf_int8_max.onnx",
+        "FP8 entropy": "fp8/regnet_x_800mf_fp8_entropy.onnx",
+        "FP8 max": "fp8/regnet_x_800mf_fp8_max.onnx",
+        "INT4 awq_clip": "int4/regnet_x_800mf_int4_awq_clip.onnx",
+        "INT4 awq_lite (asym)": "int4/regnet_x_800mf_int4_awq_lite_asym.onnx",
+        "INT4 awq_lite (sym)": "int4/regnet_x_800mf_int4_awq_lite.onnx",
+        "INT4 awq_full": "int4/regnet_x_800mf_int4_awq_full.onnx",
+        "INT4 rtn_dq": "int4/regnet_x_800mf_int4_rtn_dq.onnx",
+    }
+
+    # 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")
+        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{'='*90}")
+    print("Evaluation Comparison Table")
+    print(f"{'='*90}")
+    print(f"  {'Model':<25s} {'Top-1%':>8s} {'Top-5%':>8s} {'mAP':>8s} {'F1_mac':>8s} {'F1_wt':>8s} {'Time':>8s}")
+    print(f"  {'-'*25} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*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['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['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()

regnet_x_800mf_evaluate.py

@@ -0,0 +1,454 @@
+"""
+Evaluate mlx-vision/regnet_x_800mf-mlxim on ImageNet-1k validation set.
+
+The mlx-vision model is a port of torchvision RegNet_X_800MF (IMAGENET1K_V2)
+weights to Apple MLX. Since this evaluation runs on a standard PyTorch
+environment, we load the equivalent torchvision model with identical weights.
+
+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 torchvision model output indices.
+
+Usage:
+    python regnet_x_800mf_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
+import torchvision
+from torchvision.models import RegNet_X_800MF_Weights
+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
+
+    start = time.time()
+    for batch_idx, (imgs, labels) in enumerate(loader):
+        imgs = imgs.to(device)
+        outputs = model(imgs)
+        all_logits.append(outputs.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
+    return metrics
+
+
+# ---------------------------------------------------------------------------
+# Calibration data
+# ---------------------------------------------------------------------------
+
+def save_calibration_data(k=20, save_path="regnet_x_800mf_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:
+        k:         number of calibration images to sample
+        save_path: output .npy file path
+    """
+    # Load model & transforms (same as main)
+    print("Loading torchvision RegNet_X_800MF (IMAGENET1K_V2) ...")
+    weights = RegNet_X_800MF_Weights.IMAGENET1K_V2
+    model = torchvision.models.regnet_x_800mf(weights=weights)
+    model.eval()
+
+    transform = weights.transforms()
+    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")
+        tensor = transform(img)  # (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 mlx-vision/regnet_x_800mf-mlxim (via torchvision) 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 & transforms
+    # ------------------------------------------------------------------
+    print("Loading torchvision RegNet_X_800MF (IMAGENET1K_V2) ...")
+    print("  (equivalent to mlx-vision/regnet_x_800mf-mlxim weights)")
+    weights = RegNet_X_800MF_Weights.IMAGENET1K_V2
+    model = torchvision.models.regnet_x_800mf(weights=weights)
+    model = model.to(device)
+    model.eval()
+
+    transform = weights.transforms()
+    num_classes = 1000
+    input_size = (1, 3, 224, 224)  # (B, C, H, W)
+    print(f"  Classes: {num_classes}, Input size: {input_size}")
+
+    if args.save_model:
+        inputs = torch.randn(input_size).to(device)
+        torch.save(model.state_dict(), 'regnet_x_800mf_fp32.pth')
+
+        torch.onnx.export(model, inputs, "regnet_x_800mf_fp32.onnx")
+
+        save_calibration_data()
+
+    # ------------------------------------------------------------------
+    # 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("Evaluation Results — regnet_x_800mf (mlx-vision/regnet_x_800mf-mlxim)")
+    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}")
+
+    # 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 torchvision weights meta
+        class_names = weights.meta.get("categories", None)
+
+        def _class_label(idx):
+            if class_names:
+                return class_names[idx][:40]
+            return f"class_{idx}"
+
+        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 torchvision model card
+    print(f"\n{'='*60}")
+    print("Reference (torchvision RegNet_X_800MF IMAGENET1K_V2):")
+    print("  Top-1: 77.054%  |  Top-5: 93.530%  (224px val crop)")
+    print(f"{'='*60}")
+
+
+if __name__ == "__main__":
+    main()

regnet_x_800mf_quantize_all.py

@@ -0,0 +1,213 @@
+"""
+Quantize regnet_x_800mf_fp32.onnx using Nvidia Model Optimizer.
+
+Quantization modes attempted:
+  - FP16:  convert_to_f16
+  - INT8:  entropy / max calibration
+  - FP8:   entropy / max calibration
+  - INT4:  awq_clip / awq_lite  (symmetric + asymmetric) / awq_full / rtn_dq
+
+Calibration data: regnet_x_800mf_calibration.npy (20 images, shape 20x3x224x224)
+"""
+
+import os
+import sys
+import time
+import numpy as np
+import onnx
+import onnxruntime as ort
+
+# ---------------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------------
+ONNX_PATH = "regnet_x_800mf_fp32.onnx"
+CALIB_PATH = "regnet_x_800mf_calibration.npy"
+MODEL_NAME = "regnet_x_800mf"
+
+# Load calibration data
+calib_data = np.load(CALIB_PATH)
+print(f"Calibration data: shape={calib_data.shape}, dtype={calib_data.dtype}")
+
+# ---------------------------------------------------------------------------
+# FP16 Quantization
+# ---------------------------------------------------------------------------
+def run_fp16():
+    print("\n" + "=" * 70)
+    print("FP16 Quantization")
+    print("=" * 70)
+
+    from modelopt.onnx.quantization.int8 import convert_to_f16
+
+    onnx_model = onnx.load(ONNX_PATH, load_external_data=True)
+    fp16_model = convert_to_f16(onnx_model, keep_io_types=True)
+
+    output_path = "fp16/regnet_x_800mf_fp16.onnx"
+    if not os.path.exists("fp16"):
+        os.mkdir("fp16")
+        
+    onnx.save(fp16_model, output_path)
+    size_mb = os.path.getsize(output_path) / 1e6
+    print(f"  Saved: {output_path} ({size_mb:.1f} MB)")
+    return output_path
+
+
+# ---------------------------------------------------------------------------
+# INT8 Quantization
+# ---------------------------------------------------------------------------
+def run_int8(method, label=None):
+    """Run INT8 quantization with a given calibration method."""
+    if label is None:
+        label = f"int8_{method}{zp_tag}"
+
+    print("\n" + "=" * 70)
+    print(f"INT8 Quantization — method={method}")
+    print("=" * 70)
+
+    from modelopt.onnx.quantization import quantize
+
+    output_path = f"int8/regnet_x_800mf_{label}.onnx"
+    if not os.path.exists("int8"):
+        os.mkdir("int8")
+    t0 = time.time()
+    quantize(
+        onnx_path=ONNX_PATH,
+        quantize_mode="int8",
+        calibration_data=calib_data,
+        calibration_method=method,
+        output_path=output_path,
+        use_external_data_format=True,
+        calibration_eps=["cuda:0", "cpu"],
+    )
+    elapsed = time.time() - t0
+
+    # Calculate total size (onnx + .data)
+    size_mb = os.path.getsize(output_path) / 1e6
+    data_path = output_path + ".data"
+    if os.path.exists(data_path):
+        size_mb += os.path.getsize(data_path) / 1e6
+    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")
+    return output_path
+
+
+# ---------------------------------------------------------------------------
+# FP8 Quantization
+# ---------------------------------------------------------------------------
+def run_fp8(method):
+    print("\n" + "=" * 70)
+    print(f"FP8 Quantization — method={method}")
+    print("=" * 70)
+
+    from modelopt.onnx.quantization import quantize
+
+    output_path = f"fp8/regnet_x_800mf_fp8_{method}.onnx"
+    if not os.path.exists("fp8"):
+        os.mkdir("fp8")
+    t0 = time.time()
+    quantize(
+        onnx_path=ONNX_PATH,
+        quantize_mode="fp8",
+        calibration_data=calib_data,
+        calibration_method=method,
+        output_path=output_path,
+        use_external_data_format=True,
+        calibration_eps=["cuda:0", "cpu"],
+    )
+    elapsed = time.time() - t0
+
+    size_mb = os.path.getsize(output_path) / 1e6
+    data_path = output_path + ".data"
+    if os.path.exists(data_path):
+        size_mb += os.path.getsize(data_path) / 1e6
+    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")
+    return output_path
+
+
+# ---------------------------------------------------------------------------
+# INT4 Quantization
+# ---------------------------------------------------------------------------
+def run_int4(method, use_zero_point=False):
+    print("\n" + "=" * 70)
+    print(f"INT4 Quantization — method={method}, zero_point={use_zero_point}")
+    print("=" * 70)
+
+    from modelopt.onnx.quantization import quantize
+
+    zp_tag = "_asym" if use_zero_point else ""
+    output_path = f"int4/regnet_x_800mf_int4_{method}{zp_tag}.onnx"
+    if not os.path.exists("int4"):
+        os.mkdir("int4")
+    t0 = time.time()
+    quantize(
+        onnx_path=ONNX_PATH,
+        quantize_mode="int4",
+        calibration_data=calib_data,
+        calibration_method=method,
+        use_zero_point=use_zero_point,
+        output_path=output_path,
+        use_external_data_format=True,
+        calibration_eps=["cuda:0", "cpu"],
+    )
+    elapsed = time.time() - t0
+
+    size_mb = os.path.getsize(output_path) / 1e6
+    data_path = output_path + ".data"
+    if os.path.exists(data_path):
+        size_mb += os.path.getsize(data_path) / 1e6
+    print(f"  Saved: {output_path} ({size_mb:.1f} MB) in {elapsed:.1f}s")
+    return output_path
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+if __name__ == "__main__":
+    results = {}
+
+    # --- FP16 ---
+    try:
+        results["fp16"] = run_fp16()
+    except Exception as e:
+        print(f"  FP16 FAILED: {e}")
+
+    # --- INT8 variants ---
+    for method in ["entropy", "max"]:
+        label = f"int8_{method}"
+        try:
+            results[label] = run_int8(method, label=label)
+        except Exception as e:
+            print(f"  {label} FAILED: {e}")
+
+    # --- FP8 variants ---
+    for method in ["entropy", "max"]:
+        label = f"fp8_{method}"
+        try:
+            results[label] = run_fp8(method)
+        except Exception as e:
+            print(f"  {label} FAILED: {e}")
+
+    # --- INT4 variants ---
+    for method in ["awq_clip", "awq_full", "rtn_dq"]:
+        label = f"int4_{method}"
+        try:
+            results[label] = run_int4(method, use_zero_point=False)
+        except Exception as e:
+            print(f"  {label} FAILED: {e}")
+    # special treat for awq_lite
+    method = "awq_lite"
+    for zp in [False, True]:
+        label = f"int4_{method}" + ("_asym" if zp else "")
+        try:
+            results[label] = run_int4(method, use_zero_point=zp)
+        except Exception as e:
+            print(f"  {label} FAILED: {e}")
+            
+    # --- Summary ---
+    print("\n" + "=" * 70)
+    print("Quantization Summary")
+    print("=" * 70)
+    for label, path in results.items():
+        size_mb = os.path.getsize(path) / 1e6
+        data_path = path + ".data"
+        if os.path.exists(data_path):
+            size_mb += os.path.getsize(data_path) / 1e6
+        print(f"  {label:30s} → {path:50s} ({size_mb:.1f} MB)")

remote_eval.py

@@ -0,0 +1,563 @@
+#!/usr/bin/env python3
+"""
+Remote TensorRT evaluation script for HGNetV2-B2 quantized ONNX models.
+
+Steps:
+1. Upload ONNX models + calibration data to remote
+2. Build TensorRT engines with various precision flags
+3. Run inference on ImageNet validation set via ONNX Runtime (on remote)
+4. Collect and compare accuracy metrics
+5. Log everything to /data/rilin/hgnetv2/evaluate.log
+"""
+
+import paramiko
+import os
+import time
+import sys
+import json
+import scp
+import io
+import numpy as np
+import torch
+import timm
+from timm.data import resolve_model_data_config, create_transform
+from PIL import Image
+import pyarrow.ipc as ipc
+from sklearn.metrics import average_precision_score, precision_recall_fscore_support
+
+# ===================== CONFIG =====================
+REMOTE_HOST = "192.168.8.108"
+REMOTE_USER = "nvidia"
+REMOTE_PASS = "nvidia"
+REMOTE_BASE = "/data/rilin/hgnetv2"
+TRTEXEC = "/data/rilin/tlr/aicompiler"
+LOCAL_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"
+}
+# ===================== END CONFIG =====================
+
+
+def ssh_exec(ssh, cmd, timeout=600, log_output=True):
+    """Execute command on remote and return stdout, stderr."""
+    stdin, stdout, stderr = ssh.exec_command(cmd, timeout=timeout)
+    out = stdout.read().decode()
+    err = stderr.read().decode()
+    if log_output and out.strip():
+        print(f"  [OUT] {out.strip()[:500]}")
+    if log_output and err.strip():
+        print(f"  [ERR] {err.strip()[:500]}")
+    return out.strip(), err.strip()
+
+
+def main():
+    # Connect
+    print(f"Connecting to {REMOTE_HOST}...")
+    ssh = paramiko.SSHClient()
+    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
+    ssh.connect(REMOTE_HOST, username=REMOTE_USER, password=REMOTE_PASS, timeout=30)
+    print("Connected!")
+
+    # Create remote dirs
+    ssh_exec(ssh, f"mkdir -p {REMOTE_BASE}/models {REMOTE_BASE}/results", log_output=False)
+
+    # =====================================================================
+    # PHASE 1: Upload models and calibration data
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("PHASE 1: Uploading models to remote")
+    print("=" * 70)
+
+    scp_client = scp.SCPClient(ssh.get_transport())
+
+    # Upload calibration data
+    calib_local = "hgnetv2_b2_calibration.npy"
+    calib_remote = f"{REMOTE_BASE}/models/hgnetv2_b2_calibration.npy"
+    print(f"  Uploading {calib_local}...")
+    scp_client.put(calib_local, calib_remote)
+
+    # Upload ONNX models + external data
+    uploaded = []
+    for name, local_path in LOCAL_MODELS.items():
+        if not os.path.exists(local_path):
+            print(f"  SKIP {name}: {local_path} not found")
+            continue
+
+        remote_path = f"{REMOTE_BASE}/models/{os.path.basename(local_path)}"
+        print(f"  Uploading {name}: {local_path} -> {remote_path}")
+        scp_client.put(local_path, remote_path)
+        uploaded.append((name, remote_path))
+
+        # Also upload external data files (.data or _data)
+        for ext in [".data", "_data"]:
+            data_local = local_path + ext if ext == ".data" else local_path.replace(".onnx", ".onnx" + ext)
+            # Check various naming patterns
+            for candidate in [local_path + ".data", local_path.replace(".onnx", ".onnx.data"),
+                              local_path + "_data"]:
+                if os.path.exists(candidate):
+                    data_remote = f"{REMOTE_BASE}/models/{os.path.basename(candidate)}"
+                    print(f"    + external data: {candidate}")
+                    scp_client.put(candidate, data_remote)
+                    break
+
+    scp_client.close()
+    print(f"  Uploaded {len(uploaded)} models.")
+
+    # Verify uploads
+    out, _ = ssh_exec(ssh, f"ls -lh {REMOTE_BASE}/models/", log_output=False)
+    print(f"  Remote files:\n{out}")
+
+    # =====================================================================
+    # PHASE 2: Build TensorRT engines
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("PHASE 2: Building TensorRT engines")
+    print("=" * 70)
+
+    # FP32 baseline: just build engine
+    # FP16: --fp16
+    # INT8 QDQ models: --int8 (they already have QDQ nodes)
+    # FP8 QDQ models: --fp8
+    # INT4 QDQ models: --int4
+
+    engine_configs = []
+
+    # FP32 baseline
+    engine_configs.append({
+        "name": "fp32",
+        "onnx": f"{REMOTE_BASE}/models/hgnetv2_b2_fp32.onnx",
+        "engine": f"{REMOTE_BASE}/results/hgnetv2_b2_fp32.plan",
+        "flags": "",
+    })
+
+    # FP16
+    engine_configs.append({
+        "name": "fp16",
+        "onnx": f"{REMOTE_BASE}/models/hgnetv2_b2_fp16.onnx",
+        "engine": f"{REMOTE_BASE}/results/hgnetv2_b2_fp16.plan",
+        "flags": "--fp16",
+    })
+
+    # INT8 QDQ models (already quantized, just need --int8 flag for TRT to recognize QDQ)
+    for name in ["int8_entropy", "int8_entropy_asym", "int8_max", "int8_max_asym"]:
+        onnx_file = f"{REMOTE_BASE}/models/hgnetv2_b2_{name}.onnx"
+        engine_configs.append({
+            "name": name,
+            "onnx": onnx_file,
+            "engine": f"{REMOTE_BASE}/results/hgnetv2_b2_{name}.plan",
+            "flags": "--int8 --fp16",
+        })
+
+    # FP8 QDQ models
+    for name in ["fp8_entropy", "fp8_max"]:
+        onnx_file = f"{REMOTE_BASE}/models/hgnetv2_b2_{name}.onnx"
+        engine_configs.append({
+            "name": name,
+            "onnx": onnx_file,
+            "engine": f"{REMOTE_BASE}/results/hgnetv2_b2_{name}.plan",
+            "flags": "--fp8 --fp16",
+        })
+
+    # INT4 QDQ models
+    for name in ["int4_rtn_dq", "int4_awq_clip", "int4_awq_lite", "int4_awq_full"]:
+        onnx_file = f"{REMOTE_BASE}/models/hgnetv2_b2_{name}.onnx"
+        engine_configs.append({
+            "name": name,
+            "onnx": onnx_file,
+            "engine": f"{REMOTE_BASE}/results/hgnetv2_b2_{name}.plan",
+            "flags": "--int4 --fp16",
+        })
+
+    build_results = {}
+    for cfg in engine_configs:
+        name = cfg["name"]
+        onnx = cfg["onnx"]
+        engine = cfg["engine"]
+        flags = cfg["flags"]
+
+        print(f"\n--- Building: {name} (flags: {flags or 'none'}) ---")
+
+        # Check if ONNX file exists on remote
+        check_out, _ = ssh_exec(ssh, f"test -f {onnx} && echo EXISTS || echo MISSING", log_output=False)
+        if "MISSING" in check_out:
+            print(f"  SKIP: ONNX file not found on remote: {onnx}")
+            build_results[name] = {"status": "skipped", "reason": "onnx missing"}
+            continue
+
+        cmd = f"{TRTEXEC} --onnx={onnx} --saveEngine={engine} {flags} --tacticSources=+CUBLAS,+CUBLASLT 2>&1"
+        print(f"  CMD: {cmd}")
+        t0 = time.time()
+        out, err = ssh_exec(ssh, cmd, timeout=600)
+        elapsed = time.time() - t0
+
+        # Check if engine was built
+        check_out, _ = ssh_exec(ssh, f"test -f {engine} && echo EXISTS || echo MISSING", log_output=False)
+        if "EXISTS" in check_out:
+            size_out, _ = ssh_exec(ssh, f"ls -lh {engine}", log_output=False)
+            print(f"  SUCCESS: {size_out} ({elapsed:.1f}s)")
+            build_results[name] = {"status": "success", "engine": engine, "time": elapsed}
+        else:
+            print(f"  FAILED ({elapsed:.1f}s)")
+            build_results[name] = {"status": "failed", "time": elapsed, "output": out[-500:] if out else ""}
+
+    # =====================================================================
+    # PHASE 3: Install onnxruntime on remote and evaluate
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("PHASE 3: Installing dependencies and evaluating on ImageNet")
+    print("=" * 70)
+
+    # Check if pip packages are available
+    pip_check, _ = ssh_exec(ssh, "pip3 list 2>/dev/null | grep -iE 'onnx|numpy|pillow|sklearn'", log_output=False)
+    if not pip_check.strip():
+        print("  Installing Python packages on remote...")
+        ssh_exec(ssh, "pip3 install onnxruntime numpy pillow scikit-learn 2>&1 | tail -5", timeout=300)
+
+    # Prepare calibration data as numpy on remote for evaluation
+    # We need to generate evaluation data on the remote side
+    # Since the remote may not have ImageNet, we'll upload a pre-built evaluation batch
+
+    print("\n  Preparing evaluation data locally and uploading...")
+
+    # Build evaluation dataset from cached arrow shards
+    model = timm.create_model("hgnetv2_b2.ssld_stage2_ft_in1k", pretrained=True)
+    data_config = resolve_model_data_config(model)
+    transform = create_transform(**data_config, is_training=False)
+    del model
+
+    # Load from arrow shards
+    arrow_dir = os.path.expanduser(
+        "~/.cache/huggingface/datasets/Tsomaros___imagenet-1k_validation/"
+        "default/0.0.0/55405c49dece42420e68ddd5f80174f19b29ebaf/"
+    )
+
+    # Build a manageable subset for remote eval - use 2000 images for speed
+    # (full 41K would take very long over SSH)
+    N_EVAL = 2000
+    print(f"  Loading {N_EVAL} evaluation images from arrow shards...")
+
+    all_images = []
+    all_labels = []
+
+    shard_files = sorted(
+        f for f in os.listdir(arrow_dir)
+        if f.startswith("imagenet-1k_validation-validation-") and f.endswith(".arrow")
+    )
+
+    count = 0
+    for fname in shard_files:
+        if count >= N_EVAL:
+            break
+        path = os.path.join(arrow_dir, fname)
+        try:
+            with open(path, "rb") as f:
+                reader = ipc.RecordBatchStreamReader(f)
+                table = reader.read_all()
+        except Exception:
+            continue
+
+        for i in range(len(table)):
+            if count >= N_EVAL:
+                break
+            img_bytes = table.column("image")[i].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:
+                continue
+            label = table.column("label")[i].as_py()
+            tensor = transform(img)
+            all_images.append(tensor.numpy())
+            all_labels.append(label)
+            count += 1
+            if count % 500 == 0:
+                print(f"    [{count}/{N_EVAL}] images loaded")
+
+    images_np = np.stack(all_images)  # (N, C, H, W)
+    labels_np = np.array(all_labels)  # (N,)
+    print(f"  Evaluation data: images={images_np.shape}, labels={labels_np.shape}")
+
+    # Save and upload
+    eval_images_path = "/tmp/hgnetv2_eval_images.npy"
+    eval_labels_path = "/tmp/hgnetv2_eval_labels.npy"
+    np.save(eval_images_path, images_np)
+    np.save(eval_labels_path, labels_np)
+
+    scp_client = scp.SCPClient(ssh.get_transport())
+    scp_client.put(eval_images_path, f"{REMOTE_BASE}/models/eval_images.npy")
+    scp_client.put(eval_labels_path, f"{REMOTE_BASE}/models/eval_labels.npy")
+    scp_client.close()
+    print("  Uploaded evaluation data.")
+
+    # =====================================================================
+    # PHASE 4: Run evaluation on remote via ONNX Runtime + TensorRT EP
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("PHASE 4: Running evaluation on remote platform")
+    print("=" * 70)
+
+    # Write evaluation script and upload
+    eval_script = r'''
+import sys
+import os
+import time
+import numpy as np
+
+# Try ONNX Runtime with TensorRT EP first, fall back to CPU
+try:
+    import onnxruntime as ort
+    print(f"ONNX Runtime version: {ort.__version__}")
+    print(f"Available providers: {ort.get_available_providers()}")
+except ImportError:
+    print("ERROR: onnxruntime not installed")
+    sys.exit(1)
+
+from sklearn.metrics import average_precision_score, precision_recall_fscore_support
+import torch
+
+BASE = "/data/rilin/hgnetv2"
+
+# Load eval data
+print("Loading evaluation data...")
+images = np.load(f"{BASE}/models/eval_images.npy")
+labels = np.load(f"{BASE}/models/eval_labels.npy")
+print(f"  Images: {images.shape}, Labels: {labels.shape}")
+N = len(labels)
+num_classes = 1000
+
+# Compute metrics
+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: 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": float(top1), "top5": float(top5), "mAP": float(mAP),
+            "f1_macro": float(f1_mac), "f1_weighted": float(f1_wt)}
+
+# Models to evaluate
+models = {
+    "fp32": f"{BASE}/models/hgnetv2_b2_fp32.onnx",
+    "fp16": f"{BASE}/models/hgnetv2_b2_fp16.onnx",
+    "int8_entropy": f"{BASE}/models/hgnetv2_b2_int8_entropy.onnx",
+    "int8_entropy_asym": f"{BASE}/models/hgnetv2_b2_int8_entropy_asym.onnx",
+    "int8_max": f"{BASE}/models/hgnetv2_b2_int8_max.onnx",
+    "int8_max_asym": f"{BASE}/models/hgnetv2_b2_int8_max_asym.onnx",
+    "fp8_entropy": f"{BASE}/models/hgnetv2_b2_fp8_entropy.onnx",
+    "fp8_max": f"{BASE}/models/hgnetv2_b2_fp8_max.onnx",
+    "int4_rtn_dq": f"{BASE}/models/hgnetv2_b2_int4_rtn_dq.onnx",
+    "int4_awq_clip": f"{BASE}/models/hgnetv2_b2_int4_awq_clip.onnx",
+    "int4_awq_lite": f"{BASE}/models/hgnetv2_b2_int4_awq_lite.onnx",
+    "int4_awq_full": f"{BASE}/models/hgnetv2_b2_int4_awq_full.onnx",
+}
+
+# Evaluate each model
+results = {}
+for name, onnx_path in models.items():
+    if not os.path.exists(onnx_path):
+        print(f"\nSKIP {name}: file not found")
+        results[name] = {"error": "file not found"}
+        continue
+
+    print(f"\n{'='*50}")
+    print(f"Evaluating: {name}")
+    print(f"  ONNX: {onnx_path}")
+
+    # Try TRT EP first, then CPU
+    sess = None
+    for providers in [
+        [("TensorrtExecutionProvider", {"trt_max_workspace_size": 2147483648}), "CPUExecutionProvider"],
+        ["CPUExecutionProvider"],
+    ]:
+        try:
+            opts = ort.SessionOptions()
+            opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_DISABLE_ALL
+            sess = ort.InferenceSession(onnx_path, sess_options=opts, providers=providers)
+            print(f"  Session created with providers: {sess.get_providers()}")
+            break
+        except Exception as e:
+            print(f"  Provider {providers} failed: {e}")
+            sess = None
+            continue
+
+    if sess is None:
+        print(f"  FAILED: could not create session")
+        results[name] = {"error": "session creation failed"}
+        continue
+
+    input_name = sess.get_inputs()[0].name
+    all_logits = []
+    batch_size = 1  # batch=1 for safety with static shapes
+
+    t0 = time.time()
+    for i in range(N):
+        single = images[i:i+1]
+        out = sess.run(None, {input_name: single})
+        all_logits.append(out[0])
+        if (i+1) % 500 == 0:
+            elapsed = time.time() - t0
+            speed = (i+1) / elapsed
+            print(f"  [{i+1}/{N}] {speed:.1f} img/s")
+
+    all_logits = np.concatenate(all_logits, axis=0)
+    elapsed = time.time() - t0
+
+    metrics = compute_metrics(all_logits, labels, num_classes)
+    metrics["elapsed"] = elapsed
+    metrics["speed"] = N / elapsed
+    results[name] = metrics
+
+    print(f"  Top-1: {metrics['top1']*100:.3f}%")
+    print(f"  Top-5: {metrics['top5']*100:.3f}%")
+    print(f"  mAP:   {metrics['mAP']:.4f}")
+    print(f"  Time:  {elapsed:.1f}s ({metrics['speed']:.1f} img/s)")
+
+# Print comparison table
+print(f"\n\n{'='*80}")
+print("Evaluation Comparison Table (TensorRT Platform)")
+print(f"{'='*80}")
+print(f"  {'Model':<25s} {'Top-1%':>8s} {'Top-5%':>8s} {'mAP':>8s} {'F1_mac':>8s} {'Speed':>10s}")
+print(f"  {'-'*25} {'-'*8} {'-'*8} {'-'*8} {'-'*8} {'-'*10}")
+
+for name, m in results.items():
+    if "error" in m:
+        print(f"  {name:<25s} FAILED: {m['error']}")
+    else:
+        print(f"  {name:<25s} {m['top1']*100:>8.3f} {m['top5']*100:>8.3f} {m['mAP']:>8.4f} {m['f1_macro']:>8.4f} {m['speed']:>9.1f}/s")
+
+print(f"\n  Reference (timm): Top-1: 82.346% | Top-5: 96.394%")
+print(f"{'='*80}")
+
+# Save results as JSON
+import json
+with open(f"{BASE}/results/eval_results.json", "w") as f:
+    json.dump(results, f, indent=2)
+print(f"\nResults saved to {BASE}/results/eval_results.json")
+'''
+
+    # Upload eval script
+    script_remote = f"{REMOTE_BASE}/eval_trt.py"
+    scp_client = scp.SCPClient(ssh.get_transport())
+    scp_client.put(io.BytesIO(eval_script.encode()), script_remote)
+    scp_client.close()
+
+    # Run evaluation on remote
+    print("\n  Running remote evaluation script...")
+    print("  (This will take a while - evaluating 12 models x 2000 images)")
+
+    # Run in background and poll
+    cmd = f"cd {REMOTE_BASE} && python3 {script_remote} 2>&1 | tee {REMOTE_BASE}/evaluate.log"
+    t0 = time.time()
+
+    stdin, stdout, stderr = ssh.exec_command(cmd, timeout=3600)
+
+    # Stream output
+    full_output = ""
+    while True:
+        line = stdout.readline()
+        if not line:
+            break
+        line = line.strip()
+        if line:
+            print(f"  {line}")
+            full_output += line + "\n"
+
+    exit_code = stdout.channel.recv_exit_status()
+    elapsed = time.time() - t0
+    print(f"\n  Remote evaluation completed in {elapsed:.1f}s (exit code: {exit_code})")
+
+    # =====================================================================
+    # PHASE 5: Also test with trtexec for performance benchmarking
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("PHASE 5: TensorRT engine build + trtexec benchmarking")
+    print("=" * 70)
+
+    # For each successfully built engine, run trtexec inference benchmark
+    for cfg in engine_configs:
+        name = cfg["name"]
+        engine = cfg["engine"]
+
+        check_out, _ = ssh_exec(ssh, f"test -f {engine} && echo EXISTS || echo MISSING", log_output=False)
+        if "EXISTS" not in check_out:
+            print(f"  SKIP {name}: engine not built")
+            continue
+
+        print(f"\n--- Benchmarking: {name} ---")
+        cmd = f"{TRTEXEC} --loadEngine={engine} --iterations=100 --warmUp=10 2>&1 | tail -30"
+        out, err = ssh_exec(ssh, cmd, timeout=120)
+
+    # =====================================================================
+    # SUMMARY
+    # =====================================================================
+    print("\n" + "=" * 70)
+    print("FINAL SUMMARY")
+    print("=" * 70)
+
+    # Fetch results JSON from remote
+    try:
+        scp_client = scp.SCPClient(ssh.get_transport())
+        local_results_path = "/tmp/hgnetv2_remote_eval_results.json"
+        scp_client.get(f"{REMOTE_BASE}/results/eval_results.json", local_results_path)
+        scp_client.close()
+
+        with open(local_results_path) as f:
+            remote_results = json.load(f)
+
+        print("\nRemote Evaluation Results (ORT + TensorRT EP):")
+        print(f"  {'Model':<25s} {'Top-1%':>8s} {'Top-5%':>8s} {'mAP':>8s}")
+        print(f"  {'-'*25} {'-'*8} {'-'*8} {'-'*8}")
+        for name, m in remote_results.items():
+            if "error" in m:
+                print(f"  {name:<25s} FAILED: {m['error']}")
+            else:
+                print(f"  {name:<25s} {m['top1']*100:>8.3f} {m['top5']*100:>8.3f} {m['mAP']:>8.4f}")
+    except Exception as e:
+        print(f"  Could not fetch results: {e}")
+
+    # Fetch the remote log
+    try:
+        scp_client = scp.SCPClient(ssh.get_transport())
+        scp_client.get(f"{REMOTE_BASE}/evaluate.log", "/tmp/hgnetv2_evaluate.log")
+        scp_client.close()
+        print(f"\nRemote log saved to /tmp/hgnetv2_evaluate.log")
+    except Exception as e:
+        print(f"  Could not fetch log: {e}")
+
+    # Print build results
+    print("\nTensorRT Engine Build Results:")
+    print(f"  {'Model':<25s} {'Status':>10s} {'Time':>8s}")
+    print(f"  {'-'*25} {'-'*10} {'-'*8}")
+    for name, r in build_results.items():
+        status = r["status"]
+        t = f"{r.get('time', 0):.1f}s" if "time" in r else "-"
+        print(f"  {name:<25s} {status:>10s} {t:>8s}")
+
+    ssh.close()
+    print("\nDone!")
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,181 @@
+accelerate==1.13.0
+aiohappyeyeballs==2.6.2
+aiohttp==3.14.0
+aiosignal==1.4.0
+anaconda-anon-usage @ file:///home/task_177490767889375/croot/anaconda-anon-usage_1774907925443/work
+anaconda-auth @ file:///home/task_177681132732760/croot/anaconda-cloud-auth-split_1776812130933/work
+anaconda-cli-base @ file:///home/task_177546922169020/croot/anaconda-cli-base_1775469250358/work
+annotated-doc==0.0.4
+annotated-types @ file:///home/task_176174487847081/conda-bld/annotated-types_1761744928471/work
+antlr4-python3-runtime==4.9.3
+anyio @ file:///home/task_177313440431767/croot/anyio_1773134428538/work
+archspec @ file:///home/task_175812491784513/conda-bld/archspec_1758124989039/work
+attrs==26.1.0
+audioop-lts==0.2.2
+bcrypt==5.0.0
+boltons @ file:///croot/boltons_1751383775419/work
+brotli==1.2.0
+brotlicffi @ file:///home/task_176496126124857/conda-bld/brotlicffi_1764961326114/work
+certifi @ file:///home/task_177697098542695/croot/certifi_1776971111806/work/certifi
+cffi==2.0.0
+charset-normalizer @ file:///home/task_176174481524083/conda-bld/charset-normalizer_1761744828711/work
+click==8.4.1
+colorama==0.4.6
+conda @ file:///home/task_177585116263960/croot/conda_1775852100438/work/conda-src
+conda-anaconda-telemetry @ file:///croot/conda-anaconda-telemetry_1755883788794/work
+conda-anaconda-tos @ file:///croot/conda-anaconda-tos_1755123313635/work
+conda-content-trust @ file:///home/task_177526675282743/croot/conda-content-trust_1775266795331/work
+conda-libmamba-solver @ file:///home/task_177680598199370/croot/conda-libmamba-solver_1776806057595/work/src
+conda-package-handling @ file:///home/task_176236637997838/conda-bld/conda-package-handling_1762366405159/work
+conda_package_streaming @ file:///home/task_176236152622180/conda-bld/conda-package-streaming_1762361633858/work
+cppimport==26.4.17
+cryptography @ file:///home/task_176009247093364/conda-bld/cryptography-split_1760092601830/work
+cuda-bindings==13.3.1
+cuda-pathfinder==1.5.5
+cuda-toolkit==13.0.2
+cupy-cuda12x==14.1.1
+datasets==4.8.5
+dill==0.4.1
+distro @ file:///work/perseverance-python-buildout/croot/distro_1728396110052/work
+evaluate==0.4.6
+fastapi==0.136.3
+filelock==3.29.1
+flatbuffers==25.12.19
+frozendict @ file:///home/task_176175049677026/conda-bld/frozendict_1761750665451/work
+frozenlist==1.8.0
+fsspec==2026.2.0
+gradio==6.16.0
+gradio_client==2.5.0
+groovy==0.1.2
+h11 @ file:///home/task_176193111687894/conda-bld/h11_1761931259172/work
+hf-gradio==0.4.1
+hf-xet==1.5.0
+httpcore @ file:///croot/httpcore_1748526048470/work
+httpx @ file:///home/task_176044729083783/conda-bld/httpx_1760447344657/work
+huggingface_hub==1.17.0
+idna @ file:///home/task_176191175488988/conda-bld/idna_1761911981595/work
+invoke==3.0.3
+jaraco.classes @ file:///croot/jaraco.classes_1755516327832/work
+jaraco.context @ file:///home/task_176947759528748/croot/jaraco.context_1769477842447/work
+jaraco.functools @ file:///home/task_176838874890348/croot/jaraco.functools_1768389366336/work
+jeepney @ file:///tmp/build/80754af9/jeepney_1627537048313/work
+Jinja2==3.1.6
+joblib==1.5.3
+jsonpatch @ file:///work/perseverance-python-buildout/croot/jsonpatch_1728399595941/work
+jsonpointer @ file:///home/task_177443285493046/croot/jsonpointer_1774432880631/work
+keyring @ file:///home/task_176363713576885/conda-bld/keyring_1763637184857/work
+libmambapy @ file:///home/task_177669999816871/croot/mamba-split_1776700366233/work/libmambapy
+lief==0.17.6
+Mako==1.3.12
+markdown-it-py @ file:///home/task_176700162971740/croot/markdown-it-py_1767001650253/work
+MarkupSafe==3.0.3
+mdurl @ file:///home/task_175855216352428/conda-bld/mdurl_1758552176751/work
+menuinst @ file:///home/task_176538230483207/conda-bld/menuinst_1765382366331/work
+ml_dtypes==0.5.4
+mlx==0.31.2
+more-itertools @ file:///home/task_177581903763312/croot/more-itertools_1775819395840/work
+mpmath==1.3.0
+msgpack @ file:///croot/msgpack-python_1750958623279/work
+multidict==6.7.1
+multiprocess==0.70.19
+narwhals==2.22.0
+networkx==3.6.1
+ninja==1.13.0
+numpy==2.4.6
+nvidia-cublas==13.1.1.3
+nvidia-cuda-cupti==13.0.85
+nvidia-cuda-nvrtc==13.0.88
+nvidia-cuda-runtime==13.0.96
+nvidia-cudnn-cu13==9.20.0.48
+nvidia-cufft==12.0.0.61
+nvidia-cufile==1.15.1.6
+nvidia-curand==10.4.0.35
+nvidia-cusolver==12.0.4.66
+nvidia-cusparse==12.6.3.3
+nvidia-cusparselt-cu13==0.8.1
+nvidia-ml-py==13.610.43
+nvidia-modelopt==0.44.0
+nvidia-nccl-cu13==2.29.7
+nvidia-nvjitlink==13.0.88
+nvidia-nvshmem-cu13==3.4.5
+nvidia-nvtx==13.0.85
+omegaconf==2.3.0
+onnx==1.21.0
+onnx-ir==0.2.1
+onnx_graphsurgeon==0.6.1
+onnxconverter-common==1.16.0
+onnxruntime-gpu==1.24.4
+onnxscript==0.7.0
+onnxslim==0.1.94
+opencv-python==4.13.0.92
+orjson==3.11.9
+packaging==26.0
+pandas==3.0.3
+paramiko==5.0.0
+pillow==12.2.0
+pkce @ file:///work/perseverance-python-buildout/croot/pkce_1728388093137/work
+platformdirs @ file:///home/task_177365214898855/croot/platformdirs_1773652963129/work
+pluggy @ file:///home/task_177697287633370/croot/pluggy_1776972917169/work
+polygraphy==0.50.3
+propcache==0.5.2
+protobuf==7.35.0
+psutil==7.2.2
+PuLP==3.3.2
+pyarrow==24.0.0
+pybind11==3.0.4
+pycosat @ file:///croot/pycosat_1736868416091/work
+pycparser @ file:///home/task_177495969264897/croot/pycparser_1774959744365/work
+pydantic @ file:///home/task_177676175720061/croot/pydantic_1776761799833/work
+pydantic-settings @ file:///home/task_176416496069032/conda-bld/pydantic-settings_1764165165402/work
+pydantic_core @ file:///home/task_177669498737047/croot/pydantic-core_1776695032910/work
+pydub==0.25.1
+Pygments @ file:///home/task_177512698095713/croot/pygments_1775127017410/work
+PyJWT @ file:///home/task_177377321285166/croot/pyjwt_1773773790894/work
+PyNaCl==1.6.2
+PySocks @ file:///home/task_176175235884818/conda-bld/pysocks_1761753009495/work
+python-dateutil==2.9.0.post0
+python-dotenv @ file:///home/task_176855885656798/croot/python-dotenv_1768559688221/work
+python-multipart==0.0.30
+pytz==2026.2
+PyYAML==6.0.3
+readchar @ file:///home/task_176061336965958/conda-bld/readchar_1760613437376/work
+regex==2026.5.9
+requests @ file:///home/task_177506603358355/croot/requests_1775066075144/work
+rich @ file:///home/task_176037557612320/conda-bld/rich_1760375632000/work
+ruamel.yaml @ file:///home/task_176253595903435/conda-bld/ruamel.yaml_1762535972567/work
+ruamel.yaml.clib @ file:///home/task_176253026438192/conda-bld/ruamel.yaml.clib_1762530278005/work
+safehttpx==0.1.7
+safetensors==0.7.0
+scikit-learn==1.9.0
+scipy==1.17.1
+scp==0.15.0
+SecretStorage @ file:///home/task_177512694364126/croot/secretstorage_1775127010526/work
+semantic-version==2.10.0
+semver @ file:///home/task_176190313959014/conda-bld/semver_1761903196627/work
+setuptools==81.0.0
+shellingham @ file:///home/task_176191203629823/conda-bld/shellingham_1761912204430/work
+six==1.17.0
+socksio==1.0.0
+starlette==1.2.1
+sympy==1.14.0
+threadpoolctl==3.6.0
+timm==1.0.27
+tokenizers==0.22.2
+tomli @ file:///home/task_176829646049141/croot/tomli_1768296482268/work
+tomlkit @ file:///home/task_176289666312502/conda-bld/tomlkit_1762896676945/work
+torch==2.12.0
+torchvision==0.27.0
+tqdm @ file:///home/task_177139827237424/croot/tqdm_1771398701478/work
+transformers==5.10.2
+triton==3.7.0
+truststore @ file:///home/task_176252074050774/conda-bld/truststore_1762520903993/work
+typer==0.20.0
+typer-slim==0.20.0
+typing-inspection @ file:///home/task_176061348700415/conda-bld/typing-inspection_1760614166875/work
+typing_extensions @ file:///croot/typing_extensions_1756280817316/work
+urllib3 @ file:///home/task_176781010304586/croot/urllib3_1767810267320/work
+uvicorn==0.49.0
+wheel==0.46.3
+xxhash==3.7.0
+yarl==1.24.2
+zstandard @ file:///home/task_177696025748869/croot/zstandard_1776960539790/work

run.sh

@@ -0,0 +1,11 @@
+python3 hgnetv2_b2_evaluate.py --save_model
+python3 hgnetv2_b2_quantize_all.py
+python3 hgnetv2_b2_eval_quantized.py
+
+python remote_eval.py #ssh nvidia@192.168.8.108 (nvidia) 导出tensorrt，运行，把logits回传并进行分析。 仍有问题在改进
+
+python3 regnet_x_800mf_evaluate.py --save_model
+python3 regnet_x_800mf_quantize_all.py
+python3 regnet_x_800mf_eval_quantized.py
+
+python remote_eval.py #ssh nvidia@192.168.8.108 (nvidia) 导出tensorrt，运行，把logits回传并进行分析。 仍有问题在改进