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