codekingpro/portable-devtools

1import onnx
2from onnx import onnx_pb as onnx_proto  # noqa: F401
3
4from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, attribute_to_kwarg, ms_domain
5from .base_operator import QuantOperatorBase
6
7
8class QLinearBinaryOp(QuantOperatorBase):
9    def __init__(self, onnx_quantizer, onnx_node):
10        super().__init__(onnx_quantizer, onnx_node)
11
12    def quantize(self):
13        node = self.node
14
15        (
16            data_found,
17            output_scale_name,
18            output_zp_name,
19            _,
20            _,
21        ) = self.quantizer._get_quantization_params(node.output[0])
22        (
23            quantized_input_names,
24            zero_point_names,
25            scale_names,
26            nodes,
27        ) = self.quantizer.quantize_activation(node, [0, 1])
28        if not data_found or quantized_input_names is None:
29            return super().quantize()
30
31        qlinear_binary_math_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
32        qlinear_binary_math_name = node.name + "_quant" if node.name else ""
33
34        kwargs = {}
35        for attribute in node.attribute:
36            kwargs.update(attribute_to_kwarg(attribute))
37        kwargs["domain"] = ms_domain
38
39        qlinear_binary_math_inputs = []
40        # Input 0
41        qlinear_binary_math_inputs.append(quantized_input_names[0])
42        qlinear_binary_math_inputs.append(scale_names[0])
43        qlinear_binary_math_inputs.append(zero_point_names[0])
44        # Input 1
45        qlinear_binary_math_inputs.append(quantized_input_names[1])
46        qlinear_binary_math_inputs.append(scale_names[1])
47        qlinear_binary_math_inputs.append(zero_point_names[1])
48
49        # Output
50        qlinear_binary_math_inputs.append(output_scale_name)
51        qlinear_binary_math_inputs.append(output_zp_name)
52
53        qlinear_binary_math_node = onnx.helper.make_node(
54            "QLinear" + node.op_type,
55            qlinear_binary_math_inputs,
56            [qlinear_binary_math_output],
57            qlinear_binary_math_name,
58            **kwargs,
59        )
60        nodes.append(qlinear_binary_math_node)
61
62        # Create an entry for this quantized value
63        q_output = QuantizedValue(
64            node.output[0],
65            qlinear_binary_math_output,
66            output_scale_name,
67            output_zp_name,
68            QuantizedValueType.Input,
69        )
70        self.quantizer.quantized_value_map[node.output[0]] = q_output
71
72        self.quantizer.new_nodes += nodes
73