generate quantization parameter for outputs

onnxruntime/python/tools/quantization/calibrate.py

@@ -67,23 +67,18 @@ def __init__(
         self.symmetric = symmetric
         self.use_external_data_format = use_external_data_format
 
-        # augment graph
         self.augment_model = None
-        self.augment_graph()
-
-        # Create InferenceSession
         self.infer_session = None
         self.execution_providers = ["CPUExecutionProvider"]
-        self._create_inference_session()
 
     def set_execution_providers(self, execution_providers=["CPUExecutionProvider"]):
         """
         reset the execution providers to execute the collect_data. It triggers to re-creating inference session.
         """
         self.execution_providers = execution_providers
-        self._create_inference_session()
+        self.create_inference_session()
 
-    def _create_inference_session(self):
+    def create_inference_session(self):
         """
         create an OnnxRuntime InferenceSession.
         """
@@ -358,6 +353,7 @@ def __init__(
         self.num_bins = num_bins
         self.num_quantized_bins = num_quantized_bins
         self.percentile = percentile
+        self.tensors_to_calibrate = None
 
     def augment_graph(self):
         """
@@ -366,15 +362,11 @@ def augment_graph(self):
         """
         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:
-            added_outputs.append(value_infos[tensor])
+        self.tensors_to_calibrate, value_infos = self.select_tensors_to_calibrate(model)
+        for tensor in self.tensors_to_calibrate:
+            if tensor not in self.model_original_outputs:
+                model.graph.output.append(value_infos[tensor])
 
-        model.graph.node.extend(added_nodes)
-        model.graph.output.extend(added_outputs)
         onnx.save(
             model,
             self.augmented_model_path,
@@ -408,7 +400,7 @@ def collect_data(self, data_reader: CalibrationDataReader):
             for k, v in d.items():
                 merged_dict.setdefault(k, []).append(v)
 
-        clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict if i not in self.model_original_outputs)
+        clean_merged_dict = dict((i, merged_dict[i]) for i in merged_dict if i in self.tensors_to_calibrate)
 
         if not self.collector:
             self.collector = HistogramCollector(
@@ -825,12 +817,13 @@ def create_calibrator(
     extra_options={},
 ):
 
+    calibrator = None
     if calibrate_method == CalibrationMethod.MinMax:
         # default settings for min-max algorithm
         symmetric = False if "symmetric" not in extra_options else extra_options["symmetric"]
         moving_average = False if "moving_average" not in extra_options else extra_options["moving_average"]
         averaging_constant = 0.01 if "averaging_constant" not in extra_options else extra_options["averaging_constant"]
-        return MinMaxCalibrater(
+        calibrator = MinMaxCalibrater(
             model,
             op_types_to_calibrate,
             augmented_model_path,
@@ -844,7 +837,7 @@ def create_calibrator(
         num_bins = 128 if "num_bins" not in extra_options else extra_options["num_bins"]
         num_quantized_bins = 128 if "num_quantized_bins" not in extra_options else extra_options["num_quantized_bins"]
         symmetric = False if "symmetric" not in extra_options else extra_options["symmetric"]
-        return EntropyCalibrater(
+        calibrator = EntropyCalibrater(
             model,
             op_types_to_calibrate,
             augmented_model_path,
@@ -858,7 +851,7 @@ def create_calibrator(
         num_bins = 2048 if "num_bins" not in extra_options else extra_options["num_bins"]
         percentile = 99.999 if "percentile" not in extra_options else extra_options["percentile"]
         symmetric = True if "symmetric" not in extra_options else extra_options["symmetric"]
-        return PercentileCalibrater(
+        calibrator = PercentileCalibrater(
             model,
             op_types_to_calibrate,
             augmented_model_path,
@@ -868,4 +861,9 @@ def create_calibrator(
             percentile=percentile,
         )
 
+    if calibrator:
+        calibrator.augment_graph()
+        calibrator.create_inference_session()
+        return calibrator
+
     raise ValueError("Unsupported calibration method {}".format(calibrate_method))

onnxruntime/python/tools/quantization/qdq_quantizer.py

@@ -229,7 +229,7 @@ def quantize_tensors(self):
 
                 if data_found == False:
                     raise ValueError(
-                        "Quantization parameters are not specified for param {}."
+                        "Quantization parameters are not specified for param {}. "
                         "In static mode quantization params for inputs and outputs of nodes to be quantized are required.".format(
                             tensor_name
                         )

onnxruntime/test/python/quantization/test_calibration.py

@@ -6,14 +6,16 @@
 # license information.
 # --------------------------------------------------------------------------
 
+import tempfile
 import unittest
+from pathlib import Path
 
 import numpy as np
 import onnx
 from onnx import TensorProto, helper, numpy_helper
 
 import onnxruntime
-from onnxruntime.quantization.calibrate import CalibrationDataReader, MinMaxCalibrater
+from onnxruntime.quantization.calibrate import CalibrationDataReader, create_calibrator
 
 
 def generate_input_initializer(tensor_shape, tensor_dtype, input_name):
@@ -47,7 +49,15 @@ def rewind(self):
         self.preprocess_flag = True
 
 
-class TestCalibrate(unittest.TestCase):
+class TestCalibrateMinMaxCalibrator(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls._tmp_model_dir = tempfile.TemporaryDirectory(prefix="test_calibration.")
+
+    @classmethod
+    def tearDownClass(cls):
+        cls._tmp_model_dir.cleanup()
+
     def test_augment_graph_config_1(self):
         """TEST_CONFIG_1"""
 
@@ -74,12 +84,12 @@ def test_augment_graph_config_1(self):
         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"
-        onnx.save(model, test_model_path)
+        test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_1.onnx")
+        onnx.save(model, test_model_path.as_posix())
 
         # Augmenting graph
-        augmented_model_path = "./augmented_test_model_1.onnx"
-        calibrater = MinMaxCalibrater(test_model_path, ["Conv", "MatMul"], augmented_model_path)
+        augmented_model_path = Path(self._tmp_model_dir.name).joinpath("./augmented_test_model_1.onnx")
+        calibrater = create_calibrator(test_model_path, ["Conv", "MatMul"], augmented_model_path.as_posix())
         augmented_model = calibrater.get_augment_model()
 
         # Checking if each added ReduceMin and ReduceMax node and its output exists
@@ -162,11 +172,11 @@ def test_augment_graph_config_2(self):
         )
         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)
+        test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_2.onnx")
+        onnx.save(model, test_model_path.as_posix())
 
-        augmented_model_path = "./augmented_test_model_2.onnx"
-        calibrater = MinMaxCalibrater(test_model_path, ["Conv", "MatMul"], augmented_model_path)
+        augmented_model_path = Path(self._tmp_model_dir.name).joinpath("./augmented_test_model_2.onnx")
+        calibrater = create_calibrator(test_model_path, ["Conv", "MatMul"], augmented_model_path.as_posix())
         augmented_model = calibrater.get_augment_model()
 
         augmented_model_node_names = [node.name for node in augmented_model.graph.node]
@@ -213,12 +223,12 @@ def test_augment_graph_config_3(self):
         matmul_node = helper.make_node("MatMul", ["P", "M"], ["Q"], name="MatMul")
         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)
+        test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_3.onnx")
+        onnx.save(model, test_model_path.as_posix())
 
         # Augmenting graph
-        augmented_model_path = "./augmented_test_model_3.onnx"
-        calibrater = MinMaxCalibrater(test_model_path, ["Conv", "MatMul"], augmented_model_path)
+        augmented_model_path = Path(self._tmp_model_dir.name).joinpath("./augmented_test_model_3.onnx")
+        calibrater = create_calibrator(test_model_path, ["Conv", "MatMul"], augmented_model_path.as_posix())
         augmented_model = calibrater.get_augment_model()
 
         augmented_model_node_names = [node.name for node in augmented_model.graph.node]
@@ -315,19 +325,19 @@ def construct_test_compute_range_model(self, test_model_path):
         onnx.save(model, test_model_path)
 
     def test_compute_range(self):
-        test_model_path = "./test_model_4.onnx"
-        self.construct_test_compute_range_model(test_model_path)
+        test_model_path = Path(self._tmp_model_dir.name).joinpath("./test_model_4.onnx")
+        self.construct_test_compute_range_model(test_model_path.as_posix())
 
-        augmented_model_path = "./augmented_test_model_4.onnx"
-        calibrater = MinMaxCalibrater(test_model_path, augmented_model_path=augmented_model_path)
+        augmented_model_path = Path(self._tmp_model_dir.name).joinpath("./augmented_test_model_4.onnx")
+        calibrater = create_calibrator(test_model_path, augmented_model_path=augmented_model_path.as_posix())
         data_reader = TestDataReader()
         calibrater.collect_data(data_reader)
         tensors_range = calibrater.compute_range()
 
         sess_options = onnxruntime.SessionOptions()
         sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
         infer_session = onnxruntime.InferenceSession(
-            test_model_path,
+            test_model_path.as_posix(),
             sess_options=sess_options,
             providers=["CPUExecutionProvider"],
         )
@@ -392,8 +402,8 @@ def test_augment_graph_with_zero_value_dimension(self):
         test_model_path = "./test_model_5.onnx"
         onnx.save(model, test_model_path)
 
-        augmented_model_path = "./augmented_test_model_5.onnx"
-        calibrater = MinMaxCalibrater(test_model_path, [], augmented_model_path)
+        augmented_model_path = Path(self._tmp_model_dir.name).joinpath("./augmented_test_model_5.onnx")
+        calibrater = create_calibrator(test_model_path, [], augmented_model_path.as_posix())
         augmented_model = calibrater.get_augment_model()
 
         augmented_model_node_names = [node.name for node in augmented_model.graph.node]