codekingpro/portable-devtools

1from .base_operator import QuantOperatorBase
2
3
4# Use the quantized tensor as input without DQ.
5class QArgMax(QuantOperatorBase):
6    def __init__(self, onnx_quantizer, onnx_node):
7        super().__init__(onnx_quantizer, onnx_node)
8
9    def quantize(self):
10        node = self.node
11
12        quantized_input_value = self.quantizer.find_quantized_value(node.input[0])
13        if quantized_input_value is None:
14            self.quantizer.new_nodes += [node]
15            return
16
17        node.input[0] = quantized_input_value.q_name
18        self.quantizer.new_nodes += [node]
19