diff --git a/examples/torch_onnx/torch_quant_to_onnx.py b/examples/torch_onnx/torch_quant_to_onnx.py index bd935146e0e..a9a441e1a6d 100644 --- a/examples/torch_onnx/torch_quant_to_onnx.py +++ b/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: diff --git a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py b/tests/examples/torch_onnx/test_torch_quant_to_onnx.py index 28de24e9429..9fa4aa958d1 100644 --- a/tests/examples/torch_onnx/test_torch_quant_to_onnx.py +++ b/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"