quantization tools support qlinear average pool (#7309)

onnxruntime/python/tools/quantization/operators/pooling.py

@@ -0,0 +1,42 @@
+import onnx
+from .base_operator import QuantOperatorBase
+from ..quant_utils import attribute_to_kwarg, ms_domain, QuantizedValue, QuantizedValueType
+
+class QLinearPool(QuantOperatorBase):
+    def __init__(self, onnx_quantizer, onnx_node):
+        super().__init__(onnx_quantizer, onnx_node)
+
+    def quantize(self):
+        node = self.node
+
+        # only try to quantize when given quantization parameters for it
+        data_found, output_scale_name, output_zp_name, _, _ = \
+            self.quantizer._get_quantization_params(node.output[0])
+        if (not data_found):
+            return super().quantize()
+
+        # get quantized input tensor names, quantize input if needed
+        quantized_input_names, input_zero_point_names, input_scale_names, nodes = self.quantizer.quantize_inputs(node, [0])
+
+        # Create an entry for output quantized value.
+        qlinear_output_name = node.output[0] + "_quantized"
+        quantized_output_value = QuantizedValue(
+            node.output[0], qlinear_output_name, output_scale_name, output_zp_name, QuantizedValueType.Input)
+        self.quantizer.quantized_value_map[node.output[0]] = quantized_output_value
+
+        # Create qlinear pool node for given type (AveragePool, etc)
+        kwargs = {}
+        for attribute in node.attribute:
+            kwargs.update(attribute_to_kwarg(attribute))
+        kwargs["domain"] = ms_domain
+        qlinear_node_name = node.name + "_quant" if node.name != "" else ""
+        qnode = onnx.helper.make_node(
+            "QLinear" + node.op_type,
+            [quantized_input_names[0], input_scale_names[0], input_zero_point_names[0], output_scale_name, output_zp_name],
+            [qlinear_output_name],
+            qlinear_node_name,
+            **kwargs)
+
+        # add all newly created nodes
+        nodes.append(qnode)
+        self.quantizer.new_nodes += nodes

onnxruntime/python/tools/quantization/registry.py

@@ -15,6 +15,7 @@ from .operators.split import QSplit
 from .operators.pad import QPad
 from .operators.direct_q8 import Direct8BitOp, QDQDirect8BitOp
 from .operators.resize import QResize, QDQResize
+from .operators.pooling import QLinearPool
 
 CommonOpsRegistry = {
     "Gather": GatherQuant,
@@ -45,6 +46,7 @@ QLinearOpsRegistry = {
     "Reshape": Direct8BitOp,
     "Transpose" : Direct8BitOp,
     "Resize": QResize,
+    "AveragePool" : QLinearPool,
 }
 QLinearOpsRegistry.update(CommonOpsRegistry)
 
@@ -56,6 +58,7 @@ QDQRegistry = {
     "Transpose" : QDQDirect8BitOp,
     "Resize": QDQResize,
     "MaxPool": QDQMaxPool,
+    "AveragePool" : QDQDirect8BitOp,
 }
 
 

onnxruntime/test/python/quantization/test_op_pooling.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+# coding: utf-8
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import unittest
+import onnx
+import numpy as np
+from onnx import helper, TensorProto
+from onnxruntime.quantization import quantize_static, QuantFormat
+from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count, check_op_nodes
+
+
+class TestOpAveragePool(unittest.TestCase):
+    def input_feeds(self, n, name2shape):
+        input_data_list = []
+        for i in range(n):
+            inputs = {}
+            for name, shape in name2shape.items():
+                inputs.update({name: np.random.randint(-1, 2, shape).astype(np.float32)})
+            input_data_list.extend([inputs])
+        dr = TestDataFeeds(input_data_list)
+        return dr
+
+    def construct_model_conv_avgpool(self, output_model_path,
+                                     conv_input_shape, conv_weight_shape,
+                                     avgpool_input_shape, avgpool_attributes,
+                                     output_shape,
+                                     ):
+        #      (input)
+        #          \
+        #         Conv
+        #        /    \
+        #   Identity  AveragePool
+        #    /            \
+        # (identity_out)  (output)
+        input_tensor = helper.make_tensor_value_info('input', TensorProto.FLOAT, conv_input_shape)
+
+        conv_weight_arr = np.random.randint(-1, 2, conv_weight_shape).astype(np.float32)
+        conv_weight_initializer = onnx.numpy_helper.from_array(conv_weight_arr, name='conv1_weight')
+        conv_node = onnx.helper.make_node('Conv', ['input', 'conv1_weight'], ['conv_output'], name='conv_node')
+
+        identity_out = helper.make_tensor_value_info('identity_out', TensorProto.FLOAT, avgpool_input_shape)
+        identity_node = helper.make_node('Identity', ['conv_output'], ['identity_out'], name='IdentityNode')
+
+        initializers = [conv_weight_initializer]
+
+        output_tensor = helper.make_tensor_value_info('output', TensorProto.FLOAT, output_shape)
+        avgpool_node = helper.make_node('AveragePool', ['conv_output'], ['output'], name='avgpool_node', **avgpool_attributes)
+
+        graph = helper.make_graph([conv_node, identity_node, avgpool_node], 'TestOpQuantizerAveragePool_test_model',
+                                  [input_tensor], [identity_out, output_tensor], initializer=initializers)
+        model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 12)])
+        model.ir_version = onnx.IR_VERSION
+        onnx.save(model, output_model_path)
+
+    def test_quantize_avgpool(self):
+        np.random.seed(1)
+
+        model_fp32_path = 'avgpool_fp32.onnx'
+        model_uint8_path = 'avgpool_uint8.onnx'
+        model_uint8_qdq_path = 'avgpool_uint8_qdq.onnx'
+
+        self.construct_model_conv_avgpool(model_fp32_path,
+                                          [1, 2, 26, 42], [3, 2, 3, 3],
+                                          [1, 3, 24, 40], {'kernel_shape': [3, 3]},
+                                          [1, 3, 22, 38])
+
+        # Verify QOperator mode
+        data_reader = self.input_feeds(1, {'input': [1, 2, 26, 42]})
+        quantize_static(model_fp32_path, model_uint8_path, data_reader)
+        qnode_counts = {'QLinearConv': 1, 'QuantizeLinear': 1, 'DequantizeLinear': 2, 'QLinearAveragePool': 1}
+        check_op_type_count(self, model_uint8_path, **qnode_counts)
+        data_reader.rewind()
+        check_model_correctness(self, model_fp32_path, model_uint8_path, data_reader.get_next())
+
+        # Verify QDQ mode
+        data_reader.rewind()
+        quantize_static(model_fp32_path, model_uint8_qdq_path, data_reader, quant_format=QuantFormat.QDQ)
+        qdqnode_counts = {'Conv': 1, 'QuantizeLinear': 2, 'DequantizeLinear': 3, 'AveragePool': 1}
+        check_op_type_count(self, model_uint8_qdq_path, **qdqnode_counts)
+        data_reader.rewind()
+        check_model_correctness(self, model_fp32_path, model_uint8_qdq_path, data_reader.get_next())
+
+
+if __name__ == '__main__':
+    unittest.main()