diff --git a/onnxruntime/python/tools/quantization/calibrate.py b/onnxruntime/python/tools/quantization/calibrate.py index 894fbaa97a..38d7c22fbf 100644 --- a/onnxruntime/python/tools/quantization/calibrate.py +++ b/onnxruntime/python/tools/quantization/calibrate.py @@ -7,25 +7,17 @@ # -------------------------------------------------------------------------- import abc import itertools +import uuid from enum import Enum from pathlib import Path import numpy as np import onnx -from onnx import ModelProto, TensorProto, helper -from onnx import onnx_pb as onnx_proto +from onnx import ModelProto, TensorProto, helper, numpy_helper import onnxruntime -from .quant_utils import ( - QuantType, - apply_plot, - clone_model_with_shape_infer, - load_model, - model_has_infer_metadata, - smooth_distribution, -) -from .registry import QLinearOpsRegistry +from .quant_utils import apply_plot, clone_model_with_shape_infer, load_model, smooth_distribution class CalibrationMethod(Enum): @@ -98,9 +90,7 @@ class CalibraterBase: sess_options = onnxruntime.SessionOptions() sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL self.infer_session = onnxruntime.InferenceSession( - self.augmented_model_path, - sess_options=sess_options, - providers=self.execution_providers, + self.augmented_model_path, sess_options=sess_options, providers=self.execution_providers, ) def select_tensors_to_calibrate(self, model): @@ -180,11 +170,7 @@ class MinMaxCalibrater(CalibraterBase): :param averaging_constant: constant smoothing factor to use when computing the moving average. """ super(MinMaxCalibrater, self).__init__( - model, - op_types_to_calibrate, - augmented_model_path, - symmetric, - use_external_data_format, + model, op_types_to_calibrate, augmented_model_path, symmetric, use_external_data_format, ) self.intermediate_outputs = [] self.calibrate_tensors_range = None @@ -203,56 +189,39 @@ class MinMaxCalibrater(CalibraterBase): """ model = clone_model_with_shape_infer(self.model) - added_nodes = [] - added_outputs = [] - tensors, value_infos = self.select_tensors_to_calibrate(model) - - for tensor in tensors: + tensors, _ = self.select_tensors_to_calibrate(model) + reshape_shape_name = str(uuid.uuid4()) + reshape_shape = numpy_helper.from_array(np.array([1], dtype=np.int64), reshape_shape_name) + model.graph.initializer.append(reshape_shape) + def add_reduce_min_max(tensor_name, reduce_op_name): # When doing ReduceMax/ReduceMin, ORT can't reduce on dim with value of 0 if 'keepdims' is false. # To make the code simple, we always let keepdims to be 1. keepdims = 1 - # dim could be: - # [dim_param: "batch_size", dim_value: 256, dim_value: 36, dim_value: 64], - # [dim_value: 0], - # ... - # Please see the definition of TensorShapeProto https://github.com/onnx/onnx/blob/master/onnx/onnx.proto#L651 - dim = value_infos[tensor].type.tensor_type.shape.dim - shape = (1,) if len(dim) == 1 else tuple(1 for i in range(len(dim))) - - # Adding ReduceMin nodes - reduce_min_name = tensor + "_ReduceMin" - reduce_min_node = onnx.helper.make_node( - "ReduceMin", - [tensor], - [tensor + "_ReduceMin"], - reduce_min_name, - keepdims=keepdims, + # Adding ReduceMin/ReduceMax nodes: ReduceMin/ReduceMax -> Reshape-> (output) + reduce_output = tensor_name + "_" + reduce_op_name + intermediate_output = reduce_output + "_Reshape" + reduce_node = onnx.helper.make_node( + reduce_op_name, [tensor_name], [intermediate_output], keepdims=keepdims, name=reduce_output ) - added_nodes.append(reduce_min_node) - added_outputs.append(helper.make_tensor_value_info(reduce_min_node.output[0], TensorProto.FLOAT, shape)) - - # Adding ReduceMax nodes - reduce_max_name = tensor + "_ReduceMax" - reduce_max_node = onnx.helper.make_node( - "ReduceMax", - [tensor], - [tensor + "_ReduceMax"], - reduce_max_name, - keepdims=keepdims, + reshape_node = onnx.helper.make_node( + "Reshape", + inputs=[intermediate_output, reshape_shape_name], + outputs=[reduce_output], + name=intermediate_output, ) - added_nodes.append(reduce_max_node) - added_outputs.append(helper.make_tensor_value_info(reduce_max_node.output[0], TensorProto.FLOAT, shape)) + model.graph.node.extend([reduce_node, reshape_node]) + model.graph.output.append(helper.make_tensor_value_info(reduce_output, TensorProto.FLOAT, [1])) + + for tensor in tensors: + add_reduce_min_max(tensor, "ReduceMin") + add_reduce_min_max(tensor, "ReduceMax") - model.graph.node.extend(added_nodes) - model.graph.output.extend(added_outputs) onnx.save( - model, - self.augmented_model_path, - save_as_external_data=self.use_external_data_format, + model, self.augmented_model_path, save_as_external_data=self.use_external_data_format, ) self.augment_model = model @@ -399,9 +368,7 @@ class HistogramCalibrater(CalibraterBase): model.graph.node.extend(added_nodes) model.graph.output.extend(added_outputs) onnx.save( - model, - self.augmented_model_path, - save_as_external_data=self.use_external_data_format, + model, self.augmented_model_path, save_as_external_data=self.use_external_data_format, ) self.augment_model = model diff --git a/onnxruntime/test/python/quantization/test_calibration.py b/onnxruntime/test/python/quantization/test_calibration.py index 9a7bd657cc..b62935555e 100644 --- a/onnxruntime/test/python/quantization/test_calibration.py +++ b/onnxruntime/test/python/quantization/test_calibration.py @@ -57,21 +57,16 @@ class TestCalibrate(unittest.TestCase): # | # MatMul - A = helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 1, 5, 5]) - B = helper.make_tensor_value_info("B", TensorProto.FLOAT, [1, 1, 3, 3]) - E = helper.make_tensor_value_info("E", TensorProto.FLOAT, [1, 1, 5, 1]) - F = helper.make_tensor_value_info("F", TensorProto.FLOAT, [1, 1, 5, 1]) + vi_a = helper.make_tensor_value_info("A", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_b = helper.make_tensor_value_info("B", TensorProto.FLOAT, [1, 1, 3, 3]) + vi_e = helper.make_tensor_value_info("E", TensorProto.FLOAT, [1, 1, 5, 1]) + vi_f = helper.make_tensor_value_info("F", TensorProto.FLOAT, [1, 1, 5, 1]) conv_node = onnx.helper.make_node( - "Conv", - ["A", "B"], - ["C"], - name="Conv", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["A", "B"], ["C"], name="Conv", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) clip_node = onnx.helper.make_node("Clip", ["C"], ["D"], name="Clip") matmul_node = onnx.helper.make_node("MatMul", ["D", "E"], ["F"], name="MatMul") - graph = helper.make_graph([conv_node, clip_node, matmul_node], "test_graph_1", [A, B, E], [F]) + graph = helper.make_graph([conv_node, clip_node, matmul_node], "test_graph_1", [vi_a, vi_b, vi_e], [vi_f]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) test_model_path = "./test_model_1.onnx" @@ -86,23 +81,47 @@ class TestCalibrate(unittest.TestCase): augmented_model_node_names = [node.name for node in augmented_model.graph.node] augmented_model_outputs = [output.name for output in augmented_model.graph.output] added_node_names = [ + "A_ReduceMin", + "A_ReduceMax", + "B_ReduceMin", + "B_ReduceMax", "C_ReduceMin", "C_ReduceMax", "D_ReduceMin", "D_ReduceMax", + "E_ReduceMin", + "E_ReduceMax", "F_ReduceMin", "F_ReduceMax", + "A_ReduceMin_Reshape", + "A_ReduceMax_Reshape", + "B_ReduceMin_Reshape", + "B_ReduceMax_Reshape", + "C_ReduceMin_Reshape", + "C_ReduceMax_Reshape", + "D_ReduceMin_Reshape", + "D_ReduceMax_Reshape", + "E_ReduceMin_Reshape", + "E_ReduceMax_Reshape", + "F_ReduceMin_Reshape", + "F_ReduceMax_Reshape", ] added_outputs = [ + "A_ReduceMin", + "A_ReduceMax", + "B_ReduceMin", + "B_ReduceMax", "C_ReduceMin", "C_ReduceMax", "D_ReduceMin", "D_ReduceMax", + "E_ReduceMin", + "E_ReduceMax", "F_ReduceMin", "F_ReduceMax", ] # Original 3 nodes + added ReduceMin/Max nodes - self.assertEqual(len(augmented_model_node_names), 15) + self.assertEqual(len(augmented_model_node_names), 27) # Original 1 graph output + added outputs * 6 self.assertEqual(len(augmented_model_outputs), 13) for name in added_node_names: @@ -116,27 +135,17 @@ class TestCalibrate(unittest.TestCase): # | # Conv - G = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5]) - H = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3]) - J = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3]) - K = helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_g = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_h = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3]) + vi_j = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3]) + vi_k = helper.make_tensor_value_info("K", TensorProto.FLOAT, [1, 1, 5, 5]) conv_node_1 = onnx.helper.make_node( - "Conv", - ["G", "H"], - ["I"], - name="Conv1", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["G", "H"], ["I"], name="Conv1", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) conv_node_2 = onnx.helper.make_node( - "Conv", - ["I", "J"], - ["K"], - name="Conv2", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["I", "J"], ["K"], name="Conv2", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) - graph = helper.make_graph([conv_node_1, conv_node_2], "test_graph_2", [G, H, J], [K]) + graph = helper.make_graph([conv_node_1, conv_node_2], "test_graph_2", [vi_g, vi_h, vi_j], [vi_k]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) test_model_path = "./test_model_2.onnx" onnx.save(model, test_model_path) @@ -149,9 +158,9 @@ class TestCalibrate(unittest.TestCase): augmented_model_outputs = [output.name for output in augmented_model.graph.output] added_node_names = ["I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax"] added_outputs = ["I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax"] - # Original 2 nodes + added ReduceMin/Max nodes * 4 - self.assertEqual(len(augmented_model_node_names), 12) - # Original 1 graph output + added outputs * 4 + # Original 2 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 5 tensors + self.assertEqual(len(augmented_model_node_names), 22) + # Original 1 graph output + 5 tensors * 2 self.assertEqual(len(augmented_model_outputs), 11) for name in added_node_names: self.assertTrue(name in augmented_model_node_names) @@ -173,21 +182,16 @@ class TestCalibrate(unittest.TestCase): # | # (output) - L = helper.make_tensor_value_info("L", TensorProto.FLOAT, [1, 1, 5, 5]) - N = helper.make_tensor_value_info("N", TensorProto.FLOAT, [1, 1, 3, 3]) - Q = helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_l = helper.make_tensor_value_info("L", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_n = helper.make_tensor_value_info("N", TensorProto.FLOAT, [1, 1, 3, 3]) + vi_q = helper.make_tensor_value_info("Q", TensorProto.FLOAT, [1, 1, 5, 5]) relu_node = onnx.helper.make_node("Relu", ["L"], ["M"], name="Relu") conv_node = onnx.helper.make_node( - "Conv", - ["M", "N"], - ["O"], - name="Conv", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["M", "N"], ["O"], name="Conv", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) clip_node = onnx.helper.make_node("Clip", ["O"], ["P"], name="Clip") matmul_node = onnx.helper.make_node("MatMul", ["P", "M"], ["Q"], name="MatMul") - graph = helper.make_graph([relu_node, conv_node, clip_node, matmul_node], "test_graph_3", [L, N], [Q]) + graph = helper.make_graph([relu_node, conv_node, clip_node, matmul_node], "test_graph_3", [vi_l, vi_n], [vi_q]) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) test_model_path = "./test_model_3.onnx" onnx.save(model, test_model_path) @@ -202,16 +206,30 @@ class TestCalibrate(unittest.TestCase): added_node_names = [ "M_ReduceMin", "M_ReduceMax", + "N_ReduceMin", + "N_ReduceMax", "O_ReduceMin", "O_ReduceMax", "P_ReduceMin", "P_ReduceMax", "Q_ReduceMin", "Q_ReduceMax", + "M_ReduceMin_Reshape", + "M_ReduceMax_Reshape", + "N_ReduceMin_Reshape", + "N_ReduceMax_Reshape", + "O_ReduceMin_Reshape", + "O_ReduceMax_Reshape", + "P_ReduceMin_Reshape", + "P_ReduceMax_Reshape", + "Q_ReduceMin_Reshape", + "Q_ReduceMax_Reshape", ] added_outputs = [ "M_ReduceMin", "M_ReduceMax", + "N_ReduceMin", + "N_ReduceMax", "O_ReduceMin", "O_ReduceMax", "P_ReduceMin", @@ -219,9 +237,9 @@ class TestCalibrate(unittest.TestCase): "Q_ReduceMin", "Q_ReduceMax", ] - # Original 4 nodes + added ReduceMin/Max nodes - self.assertEqual(len(augmented_model_node_names), 14) - # Original 1 graph output + added outputs + # Original 4 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 5 tensors + self.assertEqual(len(augmented_model_node_names), 24) + # Original 1 graph output + 5 tensors * 2 self.assertEqual(len(augmented_model_outputs), 11) for name in added_node_names: self.assertTrue(name in augmented_model_node_names) @@ -243,18 +261,18 @@ class TestCalibrate(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]) - 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") - B3 = generate_input_initializer([3], np.float32, "B3") - W5 = generate_input_initializer([3, 3, 1, 1], np.float32, "W5") - B5 = generate_input_initializer([3], np.float32, "B5") + 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]) + 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") + b3 = generate_input_initializer([3], np.float32, "B3") + w5 = generate_input_initializer([3, 3, 1, 1], np.float32, "W5") + b5 = generate_input_initializer([3], np.float32, "B5") relu_node_1 = onnx.helper.make_node("Relu", ["input"], ["X1"], name="Relu1") conv_node_1 = onnx.helper.make_node("Conv", ["X1", "W1", "B1"], ["X2"], name="Conv1") relu_node_2 = onnx.helper.make_node("Relu", ["X2"], ["X3"], name="Relu2") @@ -265,14 +283,14 @@ class TestCalibrate(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], + [x1_output, x2_output, x3_output, x4_output, x5_output, x6_output], ) - graph.initializer.add().CopyFrom(W1) - graph.initializer.add().CopyFrom(B1) - graph.initializer.add().CopyFrom(W3) - graph.initializer.add().CopyFrom(B3) - graph.initializer.add().CopyFrom(W5) - graph.initializer.add().CopyFrom(B5) + graph.initializer.add().CopyFrom(w1) + graph.initializer.add().CopyFrom(b1) + graph.initializer.add().CopyFrom(w3) + graph.initializer.add().CopyFrom(b3) + graph.initializer.add().CopyFrom(w5) + graph.initializer.add().CopyFrom(b5) model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) onnx.save(model, test_model_path) @@ -289,9 +307,7 @@ class TestCalibrate(unittest.TestCase): sess_options = onnxruntime.SessionOptions() sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL infer_session = onnxruntime.InferenceSession( - test_model_path, - sess_options=sess_options, - providers=["CPUExecutionProvider"], + test_model_path, sess_options=sess_options, providers=["CPUExecutionProvider"], ) data_reader.rewind() rmin = np.array([np.inf, np.inf, np.inf, np.inf, np.inf, np.inf], dtype=np.float32) @@ -318,36 +334,23 @@ class TestCalibrate(unittest.TestCase): # | # Resize - G = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5]) - H = helper.make_tensor_value_info("H", TensorProto.FLOAT, [1, 1, 3, 3]) - J = helper.make_tensor_value_info("J", TensorProto.FLOAT, [1, 1, 3, 3]) - M = helper.make_tensor_value_info("M", TensorProto.FLOAT, [0]) - N = helper.make_tensor_value_info("N", TensorProto.FLOAT, [0]) - O = helper.make_tensor_value_info("O", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_g = helper.make_tensor_value_info("G", TensorProto.FLOAT, [1, 1, 5, 5]) + vi_m = helper.make_tensor_value_info("M", TensorProto.FLOAT, [0]) + vi_n = helper.make_tensor_value_info("N", TensorProto.FLOAT, [0]) + vi_o = helper.make_tensor_value_info("O", TensorProto.FLOAT, [1, 1, 5, 5]) # O = helper.make_tensor_value_info('O', TensorProto.FLOAT, None) conv_node_1 = onnx.helper.make_node( - "Conv", - ["G", "H"], - ["I"], - name="Conv1", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["G", "conv1_w"], ["I"], name="Conv1", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) conv_node_2 = onnx.helper.make_node( - "Conv", - ["I", "J"], - ["K"], - name="Conv2", - kernel_shape=[3, 3], - pads=[1, 1, 1, 1], + "Conv", ["I", "conv2_w"], ["K"], name="Conv2", kernel_shape=[3, 3], pads=[1, 1, 1, 1], ) resize_node_1 = onnx.helper.make_node("Resize", ["K", "M", "N"], ["O"], name="Reize1") - graph = helper.make_graph( - [conv_node_1, conv_node_2, resize_node_1], - "test_graph_5", - [G, H, J, M, N], - [O], - ) + graph = helper.make_graph([conv_node_1, conv_node_2, resize_node_1], "test_graph_5", [vi_g, vi_m, vi_n], [vi_o],) + conv1_w = generate_input_initializer([1, 1, 3, 3], np.float32, "conv1_w") + conv2_w = generate_input_initializer([1, 1, 3, 3], np.float32, "conv2_w") + graph.initializer.extend([conv1_w, conv2_w]) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) test_model_path = "./test_model_5.onnx" onnx.save(model, test_model_path) @@ -359,25 +362,49 @@ class TestCalibrate(unittest.TestCase): augmented_model_node_names = [node.name for node in augmented_model.graph.node] augmented_model_outputs = [output.name for output in augmented_model.graph.output] added_node_names = [ + "G_ReduceMin", + "G_ReduceMax", "I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax", + "M_ReduceMin", + "M_ReduceMax", + "N_ReduceMin", + "N_ReduceMax", "O_ReduceMin", "O_ReduceMax", + "G_ReduceMin_Reshape", + "G_ReduceMax_Reshape", + "I_ReduceMin_Reshape", + "I_ReduceMax_Reshape", + "K_ReduceMin_Reshape", + "K_ReduceMax_Reshape", + "M_ReduceMin_Reshape", + "M_ReduceMax_Reshape", + "N_ReduceMin_Reshape", + "N_ReduceMax_Reshape", + "O_ReduceMin_Reshape", + "O_ReduceMax_Reshape", ] added_outputs = [ + "G_ReduceMin", + "G_ReduceMax", "I_ReduceMin", "I_ReduceMax", "K_ReduceMin", "K_ReduceMax", + "M_ReduceMin", + "M_ReduceMax", + "N_ReduceMin", + "N_ReduceMax", "O_ReduceMin", "O_ReduceMax", ] - # Original 3 nodes + added ReduceMin/Max nodes * 8 - self.assertEqual(len(augmented_model_node_names), 19) - # Original 1 graph output + added outputs * 8 - self.assertEqual(len(augmented_model_outputs), 17) + # Original 3 nodes + (ReduceMin + Reshape, ReduceMax + Reshape) * 6 tensors + self.assertEqual(len(augmented_model_node_names), 27) + # Original 1 graph output + 6 tensors * 2 + self.assertEqual(len(augmented_model_outputs), 13) for name in added_node_names: self.assertTrue(name in augmented_model_node_names) for output in added_outputs: