File size: 7,779 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 | # --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# --------------------------------------------------------------------------
from __future__ import annotations
from typing import Any
import numpy as np
import onnx
from ..quant_utils import (
TENSOR_NAME_QUANT_SUFFIX,
QuantizedValue,
QuantizedValueType,
attribute_to_kwarg,
quantize_nparray,
)
from .base_operator import QuantOperatorBase
from .qdq_base_operator import QDQOperatorBase
class QPad(QuantOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
assert node.op_type == "Pad"
# Only after version 11, it has the optional constant_value
# If input[0] is not quantized, do not quanitize this node
if (self.quantizer.opset_version < 11) or (node.input[0] not in self.quantizer.quantized_value_map):
super().quantize()
return
quantized_input_value = self.quantizer.quantized_value_map[node.input[0]]
kwargs = {}
for attribute in node.attribute:
kv = attribute_to_kwarg(attribute)
kwargs.update(kv)
if "mode" not in kwargs or kwargs["mode"] == b"constant":
if len(node.input) > 2 and node.input[2] != "": # There is 3rd input 'constant_value'
zp_tensor = self.quantizer.model.get_initializer(quantized_input_value.zp_name)
scale_tensor = self.quantizer.model.get_initializer(quantized_input_value.scale_name)
if zp_tensor is None or scale_tensor is None:
super().quantize()
return
padding_constant_initializer = self.quantizer.model.get_initializer(node.input[2])
if padding_constant_initializer is not None:
zp_array = onnx.numpy_helper.to_array(zp_tensor)
zp_value = zp_array.item() if zp_array.ndim == 0 else zp_array[0]
scale_array = onnx.numpy_helper.to_array(scale_tensor)
scale_value = scale_array.item() if scale_array.ndim == 0 else scale_array[0]
padding_constant_array = onnx.numpy_helper.to_array(padding_constant_initializer)
quantized_padding_constant_array = quantize_nparray(
self.quantizer.activation_qType,
padding_constant_array,
scale_value,
zp_value,
)
quantized_padding_constant_name = node.input[2] + TENSOR_NAME_QUANT_SUFFIX
quantized_padding_constant_initializer = onnx.numpy_helper.from_array(
quantized_padding_constant_array,
quantized_padding_constant_name,
)
# Suppose this padding constant initializer only used by the node
self.quantizer.model.remove_initializer(padding_constant_initializer)
self.quantizer.model.add_initializer(quantized_padding_constant_initializer)
node.input[2] = quantized_padding_constant_name
else:
# TODO: check quantize_inputs after sub graph is supported
pad_value_qnodes = self.quantizer._get_quantize_input_nodes(
node,
2,
self.quantizer.activation_qType,
quantized_input_value.scale_name,
quantized_input_value.zp_name,
initial_type=scale_tensor.data_type,
)
self.quantizer.new_nodes.extend(pad_value_qnodes)
node.input[2] = pad_value_qnodes[0].output[0]
else:
# In quantized format, the `zero` before quantization is mapped
# to quantized_input_value.zp_name. Thus, padding 0 to
# original tensor should become padding zero point to quantized
# tensor.
if len(node.input) == 2:
# Feed quantization's zero point to padding node.
node.input.append(quantized_input_value.zp_name)
else:
# Assign quantization's zero point to padding node.
assert node.input[2] == ""
node.input[2] = quantized_input_value.zp_name
# Create an entry for output quantized value
quantized_output_value = QuantizedValue(
node.output[0],
node.output[0] + TENSOR_NAME_QUANT_SUFFIX,
quantized_input_value.scale_name,
quantized_input_value.zp_name,
QuantizedValueType.Input,
)
self.quantizer.quantized_value_map[node.output[0]] = quantized_output_value
node.input[0] = quantized_input_value.q_name
node.output[0] = quantized_output_value.q_name
self.quantizer.new_nodes += [node]
class QDQPad(QDQOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def _get_pad_const_val(self, attrs_dict: dict[str, Any]) -> np.ndarray | None:
"""
Returns the Pad's constant padding value. Returns `None` if the padding value is
not constant (i.e., comes from a dynamic input).
"""
const_val = None
onnx_tensor_type = self.quantizer.model.get_tensor_type(self.node.input[0])
if onnx_tensor_type is None:
return None
np_dtype = onnx.helper.tensor_dtype_to_np_dtype(onnx_tensor_type.elem_type)
if self.quantizer.opset_version < 11:
const_val = np.array(attrs_dict.get("value", 0), dtype=np_dtype)
elif len(self.node.input) >= 3 and self.node.input[2]:
const_val = self.quantizer.model.get_constant_value(self.node.input[2])
else:
const_val = np.array(0, dtype=np_dtype)
return const_val
def _should_quantize_output_same_as_input(self) -> bool:
"""
Returns true if Pad's output should use the same quantization parameters as input[0]
"""
attrs_dict = {}
for attribute in self.node.attribute:
kv = attribute_to_kwarg(attribute)
attrs_dict.update(kv)
pad_mode = attrs_dict.get("mode", b"constant")
if pad_mode in (b"reflect", b"edge", b"wrap"):
# These modes pad the output with a value that already exists in the input.
# So, we can quantize the output the same as the input.
return True
# For 'constant' mode, if padding with 0, we can also quantize the output the same as the input
# because our quantization floating-point range always includes 0.
if pad_mode == b"constant":
pad_val = self._get_pad_const_val(attrs_dict)
if pad_val is not None and pad_val.dtype in (np.float32, np.float16):
return float(pad_val.item()) == 0
return False
def quantize(self):
assert self.node.op_type == "Pad"
for input_name in self.node.input:
if input_name:
self.quantizer.quantize_activation_tensor(input_name)
if not self.disable_qdq_for_node_output:
if self._should_quantize_output_same_as_input():
self.quantizer.quantize_output_same_as_input(self.node.output[0], self.node.input[0], self.node.name)
else:
self.quantizer.quantize_activation_tensor(self.node.output[0])
|