File size: 1,609 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 | # -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from .qdq_base_operator import QDQOperatorBase
class QDQNormalization(QDQOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
assert node.op_type in {"InstanceNormalization", "LayerNormalization", "BatchNormalization"}
# Input
self.quantizer.quantize_activation_tensor(node.input[0])
# Scale
scale_is_initializer = self.quantizer.is_input_a_initializer(node.input[1])
scale_is_per_channel, scale_channel_axis = self.quantizer.is_tensor_per_channel(
node.input[1], default_axis=1, op_type=node.op_type
)
if scale_is_per_channel:
self.quantizer.quantize_weight_tensor_per_channel(node.input[1], axis=scale_channel_axis)
elif scale_is_initializer:
self.quantizer.quantize_weight_tensor(node.input[1])
else:
self.quantizer.quantize_activation_tensor(node.input[1])
# Bias
if len(node.input) > 2 and node.input[2]:
self.quantizer.quantize_bias_tensor(node.name, node.input[2], node.input[0], node.input[1])
# Output
if not self.disable_qdq_for_node_output:
for output_name in node.output:
self.quantizer.quantize_activation_tensor(output_name)
|