diff --git a/tests/_test_utils/onnx/lib_test_models.py b/tests/_test_utils/onnx/lib_test_models.py index 57a0621c5da..12dba85b20f 100644 --- a/tests/_test_utils/onnx/lib_test_models.py +++ b/tests/_test_utils/onnx/lib_test_models.py @@ -260,12 +260,12 @@ def find_init(onnx_model: onnx.ModelProto, init_name: str) -> np.ndarray: # conv_4 # | # -def build_conv_concat_model(): +def build_conv_concat_model(channels=128, height=240, width=64): # Define your model inputs and outputs input_names = ["input_0"] output_names = ["conv4_conv/Conv2D:0"] - input_shapes = [(1, 128, 240, 64)] - output_shapes = [(1, 256, 240, 64)] + input_shapes = [(1, channels, height, width)] + output_shapes = [(1, 2 * channels, height, width)] inputs = [ helper.make_tensor_value_info(input_name, onnx.TensorProto.FLOAT, input_shape) @@ -340,26 +340,26 @@ def build_conv_concat_model(): helper.make_tensor( name="ReadVariableOp:0", data_type=onnx.TensorProto.FLOAT, - dims=(128, 128, 3, 3), - vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3), + dims=(channels, channels, 3, 3), + vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3), ), helper.make_tensor( name="ReadVariableOp:1", data_type=onnx.TensorProto.FLOAT, - dims=(128, 128, 3, 3), - vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3), + dims=(channels, channels, 3, 3), + vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3), ), helper.make_tensor( name="ReadVariableOp:2", data_type=onnx.TensorProto.FLOAT, - dims=(128, 128, 3, 3), - vals=np.random.uniform(low=0.5, high=1.0, size=128 * 128 * 3 * 3), + dims=(channels, channels, 3, 3), + vals=np.random.uniform(low=0.5, high=1.0, size=channels * channels * 3 * 3), ), helper.make_tensor( name="ReadVariableOp:3", data_type=onnx.TensorProto.FLOAT, - dims=(256, 384, 1, 1), - vals=np.random.uniform(low=0.5, high=1.0, size=256 * 384 * 1 * 1), + dims=(2 * channels, 3 * channels, 1, 1), + vals=np.random.uniform(low=0.5, high=1.0, size=2 * channels * 3 * channels), ), ] diff --git a/tests/_test_utils/onnx/quantization/autotune/models.py b/tests/_test_utils/onnx/quantization/autotune/models.py index 84a8b4ab82e..89721661b4a 100644 --- a/tests/_test_utils/onnx/quantization/autotune/models.py +++ b/tests/_test_utils/onnx/quantization/autotune/models.py @@ -25,13 +25,13 @@ from onnx import helper -def _create_simple_conv_onnx_model(): +def _create_simple_conv_onnx_model(batch_size=64, spatial_size=224): """Build ONNX model: Input -> Conv -> Relu -> Output (minimal for autotuner tests).""" input_tensor = helper.make_tensor_value_info( - "input", onnx.TensorProto.FLOAT, [64, 32, 224, 224] + "input", onnx.TensorProto.FLOAT, [batch_size, 32, spatial_size, spatial_size] ) output_tensor = helper.make_tensor_value_info( - "output", onnx.TensorProto.FLOAT, [64, 64, 224, 224] + "output", onnx.TensorProto.FLOAT, [batch_size, 64, spatial_size, spatial_size] ) conv_node = helper.make_node( "Conv", diff --git a/tests/gpu/onnx/quantization/autotune/test_benchmark.py b/tests/gpu/onnx/quantization/autotune/test_benchmark.py index 012bf495ab6..47311e25219 100644 --- a/tests/gpu/onnx/quantization/autotune/test_benchmark.py +++ b/tests/gpu/onnx/quantization/autotune/test_benchmark.py @@ -46,8 +46,8 @@ @pytest.fixture def simple_conv_model_bytes(): - """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils).""" - model = _create_simple_conv_onnx_model() + """ONNX model bytes: Input -> Conv -> Relu -> Output (from _test_utils), small input.""" + model = _create_simple_conv_onnx_model(batch_size=1, spatial_size=32) return model.SerializeToString() @@ -431,7 +431,12 @@ def test_run_with_path(self, simple_conv_model_path): with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f: cache_path = f.name try: - benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2) + benchmark = TrtExecBenchmark( + timing_cache_file=cache_path, + warmup_runs=1, + timing_runs=2, + trtexec_args=["--duration=0"], + ) latency_ms = benchmark.run(simple_conv_model_path) assert isinstance(latency_ms, float) assert latency_ms > 0 @@ -445,7 +450,12 @@ def test_run_with_bytes(self, simple_conv_model_bytes): with tempfile.NamedTemporaryFile(suffix=".cache", delete=False) as f: cache_path = f.name try: - benchmark = TrtExecBenchmark(timing_cache_file=cache_path, warmup_runs=1, timing_runs=2) + benchmark = TrtExecBenchmark( + timing_cache_file=cache_path, + warmup_runs=1, + timing_runs=2, + trtexec_args=["--duration=0"], + ) latency_ms = benchmark.run(simple_conv_model_bytes) assert isinstance(latency_ms, float) assert latency_ms > 0 diff --git a/tests/gpu/onnx/quantization/autotune/test_workflow.py b/tests/gpu/onnx/quantization/autotune/test_workflow.py index eebc42f6c33..b323402f1eb 100644 --- a/tests/gpu/onnx/quantization/autotune/test_workflow.py +++ b/tests/gpu/onnx/quantization/autotune/test_workflow.py @@ -59,7 +59,10 @@ def test_export_quantized_model(use_trtexec, simple_conv_model): output_path = output_dir + ".quant.onnx" try: - init_benchmark_instance(use_trtexec=use_trtexec, timing_runs=100) + # trtexec measures for at least 3 s by default; --duration=0 stops after timing_runs iterations + init_benchmark_instance( + use_trtexec=use_trtexec, timing_runs=100, trtexec_args=["--duration=0"] + ) autotuner = region_pattern_autotuning_workflow(baseline_model_path, Path(output_dir)) # Export model with Q/DQ insertion diff --git a/tests/gpu/onnx/quantization/test_concat_elim.py b/tests/gpu/onnx/quantization/test_concat_elim.py index 68179051909..de9feafe7e7 100644 --- a/tests/gpu/onnx/quantization/test_concat_elim.py +++ b/tests/gpu/onnx/quantization/test_concat_elim.py @@ -23,10 +23,19 @@ from modelopt.onnx.quantization.quantize import quantize +# Keep every Conv above the 16-channel limit for INT8/FP8 quantization (graph_selection.py) so the +# same nodes are quantized as with the default 128-channel model, but calibrate a small input. +SMALL_MODEL_SHAPE = {"channels": 32, "height": 16, "width": 16} + def _check_concat_qdq_status(onnx_path, quantize_mode): - # Quantize the input model - quantize(onnx_path, quantize_mode=quantize_mode, passes="concat_elimination") + # Quantize the input model. Only the graph rewrite is checked, so calibrate on CPU. + quantize( + onnx_path, + quantize_mode=quantize_mode, + passes="concat_elimination", + calibration_eps=["cpu"], + ) # Output model should be produced in the same tmpdir output_onnx_path = onnx_path.replace(".onnx", ".quant.onnx") @@ -61,7 +70,7 @@ def _check_concat_qdq_status(onnx_path, quantize_mode): def test_concat_elim_int8(tmp_path): # Copy the test model to the tmp_path - model = build_conv_concat_model() + model = build_conv_concat_model(**SMALL_MODEL_SHAPE) this_function_name = sys._getframe().f_code.co_name onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx") onnx.save(model, onnx_path) @@ -70,7 +79,7 @@ def test_concat_elim_int8(tmp_path): def test_concat_elim_fp8(tmp_path): # Copy the test model to the tmp_path - model = build_conv_concat_model() + model = build_conv_concat_model(**SMALL_MODEL_SHAPE) this_function_name = sys._getframe().f_code.co_name onnx_path = os.path.join(tmp_path, f"{this_function_name}.onnx") onnx.save(model, onnx_path) diff --git a/tests/gpu/onnx/quantization/test_dla_qdq_rules.py b/tests/unit/onnx/quantization/test_dla_qdq_rules.py similarity index 100% rename from tests/gpu/onnx/quantization/test_dla_qdq_rules.py rename to tests/unit/onnx/quantization/test_dla_qdq_rules.py diff --git a/tests/gpu/onnx/quantization/test_dla_scalar_constants.py b/tests/unit/onnx/quantization/test_dla_scalar_constants.py similarity index 100% rename from tests/gpu/onnx/quantization/test_dla_scalar_constants.py rename to tests/unit/onnx/quantization/test_dla_scalar_constants.py diff --git a/tests/gpu/onnx/quantization/test_quantization_initializer_names.py b/tests/unit/onnx/quantization/test_quantization_initializer_names.py similarity index 100% rename from tests/gpu/onnx/quantization/test_quantization_initializer_names.py rename to tests/unit/onnx/quantization/test_quantization_initializer_names.py