From c7ae9b977af4676be9505c5bfefa66d6014cc431 Mon Sep 17 00:00:00 2001 From: Adrian Lizarraga Date: Mon, 9 Sep 2024 12:05:41 -0700 Subject: [PATCH] [Quantization] Apply workaround for crash when using histogram-based calibrators (#21972) ### Description - Applies a workaround that prevents the histogram-based calibrators (percentile, entropy, distribution) from crashing. The workaround involves copying inference outputs that come directly from model inputs. A description of the bug is here: https://github.com/microsoft/onnxruntime/issues/21922. **This PR does not fix the root bug, but instead provides a workaround to _unblock_ users using histogram-based calibration.** - Adds a unit test that runs all histogram-based calibrators to help catch future regressions. We didn't have unit tests that ran these calibration methods. ### Motivation and Context Trying to quantize a model with the percentile, entropy, or distribution calibration methods raises an exception: ```shell File "/.../site-packages/onnxruntime/quantization/quantize.py", line 691, in quantize quantize_static( File "/.../site-packages/onnxruntime/quantization/quantize.py", line 525, in quantize_static calibrator.collect_data(calibration_data_reader) File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 571, in collect_data self.collector.collect(clean_merged_dict) File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 746, in collect return self.collect_value(name_to_arr) File "/.../site-packages/onnxruntime/quantization/calibrate.py", line 836, in collect_value hist, hist_edges = np.histogram(data_arr, self.num_bins, range=(-threshold, threshold)) File "<__array_function__ internals>", line 180, in histogram File ".../site-packages/numpy/lib/histograms.py", line 793, in histogram bin_edges, uniform_bins = _get_bin_edges(a, bins, range, weights) File "/.../site-packages/numpy/lib/histograms.py", line 426, in _get_bin_edges first_edge, last_edge = _get_outer_edges(a, range) File "/.../site-packages/numpy/lib/histograms.py", line 315, in _get_outer_edges raise ValueError( ValueError: supplied range of [nan, nan] is not finite ``` The calibrators create an augmented model with all tensors (including model inputs) set as model outputs. The data for outputs that are also model inputs is corrupted as described in https://github.com/microsoft/onnxruntime/issues/21922. The corrupted data sometimes contains `NaN` values that cause numpy's histogram utilities to raise an exception. --- .../python/tools/quantization/calibrate.py | 17 +++++- .../python/quantization/test_calibration.py | 53 +++++++++++++++---- 2 files changed, 59 insertions(+), 11 deletions(-) diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 703accbcc1..74d064204e 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -565,16 +565,29 @@ class HistogramCalibrater(CalibraterBase): """ Entropy Calibrator collects operators' tensors as well as generates tensor histogram for each operator. """ + input_names_set = {node_arg.name for node_arg in self.infer_session.get_inputs()} + output_names = [node_arg.name for node_arg in self.infer_session.get_outputs()] + while True: inputs = data_reader.get_next() if not inputs: break - self.intermediate_outputs.append(self.infer_session.run(None, inputs)) + outputs = self.infer_session.run(None, inputs) + + # Copy np.ndarray only for graph outputs that are also graph inputs to workaround bug: + # https://github.com/microsoft/onnxruntime/issues/21922 + fixed_outputs = [] + for output_index, output in enumerate(outputs): + if output_names[output_index] in input_names_set: + fixed_outputs.append(copy.copy(output)) + else: + fixed_outputs.append(output) + + self.intermediate_outputs.append(fixed_outputs) if len(self.intermediate_outputs) == 0: raise ValueError("No data is collected.") - output_names = [self.infer_session.get_outputs()[i].name for i in range(len(self.intermediate_outputs[0]))] output_dicts_list = [ dict(zip(output_names, intermediate_output)) for intermediate_output in self.intermediate_outputs ] diff --git a/onnxruntime/test/python/quantization/test_calibration.py b/onnxruntime/test/python/quantization/test_calibration.py index b36ba45cdb..e7f5a959a3 100644 --- a/onnxruntime/test/python/quantization/test_calibration.py +++ b/onnxruntime/test/python/quantization/test_calibration.py @@ -14,7 +14,7 @@ import onnx from onnx import TensorProto, helper, numpy_helper import onnxruntime -from onnxruntime.quantization.calibrate import CalibrationDataReader, create_calibrator +from onnxruntime.quantization.calibrate import CalibrationDataReader, CalibrationMethod, create_calibrator def generate_input_initializer(tensor_shape, tensor_dtype, input_name): @@ -275,7 +275,7 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase): for output in added_outputs: self.assertTrue(output in augmented_model_outputs) - def construct_test_compute_data_model(self, test_model_path, opset_version=13): + def construct_test_compute_data_model(self, test_model_path, opset_version=13, augmented=True): # (input) # | # Relu @@ -290,12 +290,19 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase): # | # (X6) input = helper.make_tensor_value_info("input", TensorProto.FLOAT, [1, 3, 1, 3]) - x1_output = helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3]) - x2_output = helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3]) - x3_output = helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3]) - x4_output = helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3]) - x5_output = helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3]) - x6_output = helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3]) + graph_outputs = None + if augmented: + graph_outputs = [ + helper.make_tensor_value_info("X1", TensorProto.FLOAT, [1, 3, 1, 3]), + helper.make_tensor_value_info("X2", TensorProto.FLOAT, [1, 3, 1, 3]), + helper.make_tensor_value_info("X3", TensorProto.FLOAT, [1, 3, 1, 3]), + helper.make_tensor_value_info("X4", TensorProto.FLOAT, [1, 3, 1, 3]), + helper.make_tensor_value_info("X5", TensorProto.FLOAT, [1, 3, 1, 3]), + helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3]), + ] + else: + graph_outputs = [helper.make_tensor_value_info("X6", TensorProto.FLOAT, [1, 3, 1, 3])] + w1 = generate_input_initializer([3, 3, 1, 1], np.float32, "W1") b1 = generate_input_initializer([3], np.float32, "B1") w3 = generate_input_initializer([3, 3, 1, 1], np.float32, "W3") @@ -312,7 +319,7 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase): [relu_node_1, conv_node_1, relu_node_2, conv_node_2, conv_node_3, add_node], "test_graph_4", [input], - [x1_output, x2_output, x3_output, x4_output, x5_output, x6_output], + graph_outputs, ) graph.initializer.add().CopyFrom(w1) graph.initializer.add().CopyFrom(b1) @@ -357,6 +364,34 @@ class TestCalibrateMinMaxCalibrator(unittest.TestCase): for output_name in output_min_max_dict: self.assertEqual(output_min_max_dict[output_name], tensors_range[output_name].range_value) + def test_histogram_calibrators_run(self): + """ + Runs all histogram-based calibrators (Percentile, Entropy, Distribution) and checks that they run + and generate the expected number of tensor ranges. Does not check correctness of range values. + """ + # Create test model. + test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_4.onnx") + self.construct_test_compute_data_model(test_model_path.as_posix(), augmented=False) + + # Count the number of tensors in the model. + model = onnx.load_model(test_model_path) + model = onnx.shape_inference.infer_shapes(model) + num_tensors = len(model.graph.value_info) + len(model.graph.input) + len(model.graph.output) + + # Run all histogram calibration methods. + data_reader = TestDataReader() + calibration_methods = [CalibrationMethod.Percentile, CalibrationMethod.Entropy, CalibrationMethod.Distribution] + for calibration_method in calibration_methods: + with self.subTest(calibration_method=calibration_method): + data_reader.rewind() + augmented_model_path = Path(self._tmp_model_dir.name).joinpath(f"augmented_{calibration_method}.onnx") + calibrator = create_calibrator( + test_model_path, calibrate_method=calibration_method, augmented_model_path=augmented_model_path + ) + calibrator.collect_data(data_reader) + tensors_range = calibrator.compute_data() + self.assertEqual(len(tensors_range.items()), num_tensors) # A range for every tensor in the graph. + def test_augment_graph_with_zero_value_dimension(self): """TEST_CONFIG_5""" # Conv