File size: 2,166 Bytes
610d9c2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType
from .base_operator import QuantOperatorBase
from .qdq_base_operator import QDQOperatorBase
"""
Quantize Gather
"""
class GatherQuant(QuantOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def should_quantize(self):
if not self.quantizer.should_quantize_node(self.node):
return False
return self.quantizer.is_valid_quantize_weight(self.node.input[0])
def quantize(self):
node = self.node
assert node.op_type == "Gather"
(
quantized_input_names,
zero_point_names,
scale_names,
nodes,
) = self.quantizer.quantize_activation(node, [0])
if quantized_input_names is None:
return super().quantize()
gather_new_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
# Create an entry for this quantized value
q_output = QuantizedValue(
node.output[0],
gather_new_output,
scale_names[0],
zero_point_names[0],
QuantizedValueType.Input,
)
self.quantizer.quantized_value_map[node.output[0]] = q_output
node.output[0] = gather_new_output
node.input[0] = quantized_input_names[0]
nodes.append(node)
self.quantizer.new_nodes += nodes
class QDQGather(QDQOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
assert node.op_type == "Gather" or node.op_type == "GatherElements"
if self.quantizer.is_valid_quantize_weight(node.input[0]) or self.quantizer.force_quantize_no_input_check:
self.quantizer.quantize_activation_tensor(node.input[0])
self.quantizer.quantize_output_same_as_input(node.output[0], node.input[0], node.name)
elif self.quantizer.is_tensor_quantized(node.input[0]):
self.quantizer.quantize_output_same_as_input(node.output[0], node.input[0], node.name)
|