[9/11] Trim the torch_onnx example lane

examples/torch_onnx/torch_quant_to_onnx.py

@@ -27,7 +27,6 @@
 
 import timm
 import torch
-import torch.multiprocessing as mp
 import torch.nn.functional as F
 from datasets import load_dataset
 from download_example_onnx import export_to_onnx
@@ -56,9 +55,6 @@
 """
 
 
-mp.set_start_method("spawn", force=True)  # Needed for data loader with multiple workers
-
-
 _FP8_CONV_OVERRIDE: list = [
     {
         "parent_class": "nn.Conv2d",
@@ -311,16 +307,17 @@ def load_calibration_data(model, data_size, batch_size, device, with_labels=Fals
     calib_tensor = [transforms(img) for img in images]
     calib_tensor = [t.to(device) for t in calib_tensor]
 
+    # The samples already sit on ``device``; worker processes would only add spawn and import cost.
     if with_labels:
         labels = dataset["train"][:data_size]["label"]
         labels = torch.tensor(labels, device=device)
         calib_dataset = [{"image": img, "label": lbl} for img, lbl in zip(calib_tensor, labels)]
         return torch.utils.data.DataLoader(
-            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_dataset, batch_size=batch_size, shuffle=True, num_workers=0
         )
     else:
         return torch.utils.data.DataLoader(
-            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=4
+            calib_tensor, batch_size=batch_size, shuffle=True, num_workers=0
         )
 
 
@@ -761,6 +758,10 @@ def build_trt_engine(onnx_path, optimization_level=4):
         f"--onnx={onnx_path}",
         "--stronglyTyped",
         f"--builderOptimizationLevel={optimization_level}",
+        # One inference iteration shows the engine runs; the default 3 s timed benchmark is not shown.
+        "--warmUp=0",
+        "--duration=0",
+        "--iterations=1",
     ]
     print(f"\nBuilding TensorRT engine: {' '.join(cmd)}")
     try:

tests/examples/torch_onnx/test_torch_quant_to_onnx.py

@@ -15,6 +15,7 @@
 
 
 import importlib
+import os
 from collections import defaultdict
 from pathlib import Path
 
@@ -38,6 +39,10 @@
     "resnet50": ("resnet50", None),
 }
 
+# PRs run swinv2 for fp8 and nvfp4 only (swin_tiny covers every format); scheduled runs keep all.
+_PR_CI = os.environ.get("GITHUB_REF", "").startswith("refs/heads/pull-request/")
+_SKIPPED_ON_PR = {("swinv2_tiny", "int8"), ("swinv2_tiny", "mxfp8"), ("swinv2_tiny", "auto")}
+
 
 def _assert_residual_inputs_are_quantized(onnx_save_path):
     model = onnx.load(onnx_save_path)
@@ -74,6 +79,8 @@ def _assert_residual_inputs_are_quantized(onnx_save_path):
 def test_torch_onnx(tmp_path, model_key, qformat):
     if model_key == "resnet50" and qformat not in _RESNET_RECIPE_QFORMATS:
         pytest.skip("Only FP8 and INT8 quantization are supported for ResNet")
+    if _PR_CI and (model_key, qformat) in _SKIPPED_ON_PR:
+        pytest.skip("Covered by swin_tiny on PRs; runs in the scheduled CI")
 
     timm_model_name, model_kwargs = _MODELS[model_key]
     onnx_save_path = tmp_path / f"{model_key}.{qformat}.onnx"