codekingpro/portable-devtools

1import onnx
2
3from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, attribute_to_kwarg, ms_domain
4from .base_operator import QuantOperatorBase
5from .qdq_base_operator import QDQOperatorBase
6
7
8class QLinearActivation(QuantOperatorBase):
9    def __init__(self, onnx_quantizer, onnx_node):
10        super().__init__(onnx_quantizer, onnx_node)
11
12    def QuantizeClipRelu(self):  # noqa: N802
13        node = self.node
14        assert node.op_type == "Relu" or node.op_type == "Clip"
15
16        # When mode is QLinearOps, the output quantization params are calculated based on outputs from
17        # activation nodes, therefore these nodes can be removed from the graph if they follow a quantized op.
18        # If input to this node is not quantized then keep this node
19        # If activation is symmetric, not quantize the op and simply return
20        if node.input[0] not in self.quantizer.quantized_value_map or self.quantizer.is_activation_symmetric:
21            return super().quantize()
22
23        quantized_value = self.quantizer.quantized_value_map[node.input[0]]
24        self.quantizer.quantized_value_map[node.output[0]] = quantized_value
25
26    def quantize(self):
27        node = self.node
28        if node.op_type == "Relu" or node.op_type == "Clip":
29            self.QuantizeClipRelu()
30            return
31
32        nnapi_sigmoid_option = "extra.Sigmoid.nnapi"
33        sigmoid_nnapi_mode = (
34            node.op_type == "Sigmoid"
35            and nnapi_sigmoid_option in self.quantizer.extra_options
36            and self.quantizer.extra_options[nnapi_sigmoid_option]
37        )
38        use_scale = 1 / 256.0 if sigmoid_nnapi_mode else None
39        use_zeropoint = 0 if sigmoid_nnapi_mode else None
40
41        # No assert on op_type as it is controlled by registry
42        # only try to quantize when given quantization parameters for it
43        (
44            data_found,
45            output_scale_name,
46            output_zp_name,
47            _,
48            _,
49        ) = self.quantizer._get_quantization_params(node.output[0], use_scale, use_zeropoint)
50        (
51            quantized_input_names,
52            zero_point_names,
53            scale_names,
54            nodes,
55        ) = self.quantizer.quantize_activation(node, [0])
56        if not data_found or quantized_input_names is None:
57            return super().quantize()
58
59        qlinear_activation_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
60        qlinear_activation_name = ""
61        if node.name:
62            qlinear_activation_name = node.name + "_quant"
63        kwargs = {}
64        for attribute in node.attribute:
65            kwargs.update(attribute_to_kwarg(attribute))
66        kwargs["domain"] = ms_domain
67
68        qlinear_activation_inputs = [
69            quantized_input_names[0],
70            scale_names[0],
71            zero_point_names[0],
72            output_scale_name,
73            output_zp_name,
74        ]
75
76        qlinear_activation_node = onnx.helper.make_node(
77            "QLinear" + node.op_type,
78            qlinear_activation_inputs,
79            [qlinear_activation_output],
80            qlinear_activation_name,
81            **kwargs,
82        )
83
84        # Create an entry for this quantized value
85        q_output = QuantizedValue(
86            node.output[0],
87            qlinear_activation_output,
88            output_scale_name,
89            output_zp_name,
90            QuantizedValueType.Input,
91        )
92        self.quantizer.quantized_value_map[node.output[0]] = q_output
93
94        nodes.append(qlinear_activation_node)
95        self.quantizer.new_nodes += nodes
96
97
98class QDQRemovableActivation(QDQOperatorBase):
99    def __init__(self, onnx_quantizer, onnx_node):
100        super().__init__(onnx_quantizer, onnx_node)
101
102    def quantize(self):
103        node = self.node
104
105        # If input to this node is not quantized then keep this node
106        if not self.quantizer.is_tensor_quantized(node.input[0]):
107            return
108
109        if (
110            not self.quantizer.is_activation_symmetric
111            and not self.quantizer.qdq_keep_removable_activations
112            and self.quantizer.try_replacing_upstream_output(node.input[0], node.output[0])
113        ):
114            self.quantizer.remove_node(self.node)
115        else:
116            self.quantizer.quantize_activation_tensor(node.input[0])
117
118        if not self.disable_qdq_for_node_output:
119            self.quantizer.quantize_activation_tensor(node.output[0])
120