codekingpro/portable-devtools
114k
1# -------------------------------------------------------------------------
2# Copyright (c) Microsoft Corporation. All rights reserved.
3# Licensed under the MIT License. See License.txt in the project root for
4# license information.
5# --------------------------------------------------------------------------
6from __future__ import annotations
7
8import logging
9from dataclasses import dataclass
10from enum import Enum
11from typing import Any
12
13import numpy as np
14import onnx
15from onnx import TensorProto
16from onnx import onnx_pb as onnx_proto
17
18from .base_quantizer import BaseQuantizer, QuantizationParams
19from .calibrate import TensorData
20from .quant_utils import (
21 DEQUANT_OP_NAME,
22 ONNX_TYPE_TO_NP_TYPE,
23 QUANT_OP_NAME,
24 QuantizedValue,
25 QuantizedValueType,
26 __producer__,
27 __version__,
28 add_dequant_output_suffix,
29 add_dequant_suffix,
30 add_quant_input_suffix,
31 add_quant_output_suffix,
32 add_quant_suffix,
33 compute_data_quant_params,
34 compute_scale_zp,
35 compute_scale_zp_float8,
36 find_by_name,
37 get_qmin_qmax_for_qType,
38 ms_domain,
39 normalize_axis,
40 quantize_onnx_initializer,
41 tensor_proto_to_array,
42)
43from .registry import CreateQDQQuantizer
44
45
46class QDQQuantTensorType(Enum):
47 ACTIVATION = 0
48 WEIGHT = 1
49 BIAS = 2
50
51
52# Holds the name of the node input from which a node output will share the
53# same quantization param initializers (zero-point and scale initializers).
54# Ex: A Transpose node's output will use the same quant param initializers used at the input.
55@dataclass
56class QDQQuantParamProvider:
57 input_name: str
58 node_name: str
59
60
61# Holds information for tensors that have been marked for quantization by operator quantizers.
62# Does not hold information for bias tensors.
63class QDQTensorQuantInfo:
64 def __init__(self, tensor_type=QDQQuantTensorType.ACTIVATION, quant_para_provider=None, axis=None, data_type=None):
65 self.tensor_type = tensor_type
66 self.quant_para_provider = quant_para_provider
67 self.axis = axis
68 self.is_shared = quant_para_provider is not None
69 assert data_type is not None
70 self.data_type = data_type
71
72
73# Holds information for bias tensors that have been marked for quantization by operator quantizers.
74@dataclass
75class QDQBiasQuantInfo:
76 node_name: str
77 input_name: str
78 weight_name: str
79 beta: float
80
81
82# Holds quantization parameter values (scale, zp) for a tensor.
83# A tensor typically has a one set of quantization parameters, unless the tensor is
84# at a "mixed-precision" boundary where the activation quantization type changes (e.g., from uint8 to uint16).
85@dataclass
86class QDQTensorQuantParams:
87 original: QuantizationParams # Generated by producer node.
88 converted: QuantizationParams | None # Converted type consumed by some (or all/none) consumer nodes.
89 converted_recv_nodes: set[str] | None # The name of nodes that consume the converted type.
90
91 def get_for_consumer(self, consumer_node_name) -> QuantizationParams:
92 if self.converted is None: # Quantized value is not converted, return original
93 return self.original
94
95 if self.converted_recv_nodes is None: # All consumers receive the converted value
96 return self.converted
97
98 # Check if consumer node name is in the list of nodes that
99 # receive the converted quantization value. If not, return the original value generated
100 # by the tensor's producer.
101 return self.converted if (consumer_node_name in self.converted_recv_nodes) else self.original
102
103
104# Holds scale and zero_point initializer TensorProtos.
105@dataclass
106class QDQScaleZpInitializers:
107 scale: TensorProto
108 zero_point: TensorProto
109
110
111# Holds all scale and zero-point initializers for a tensor.
112# A tensor typically has a one set of quantization parameters, unless the tensor is
113# at a "mixed-precision" boundary where the activation quantization type changes (e.g., from uint8 to uint16).
114@dataclass
115class QDQTensorScaleZpInitializers:
116 original: QDQScaleZpInitializers
117 converted: QDQScaleZpInitializers | None
118 converted_recv_nodes: set[str] | None
119
120
121# Holds cached information of a tensor's quantized values (types, zp/scale initializer names, etc.).
122# A tensor typically has a one set of quantization parameters, unless the tensor is
123# at a "mixed-precision" boundary where the activation quantization type changes (e.g., from uint8 to uint16).
124@dataclass
125class QDQTensorQuantizedValue:
126 original: QuantizedValue
127 converted: QuantizedValue | None
128 converted_recv_nodes: set[str] | None
129
130 def get_for_consumer(self, consumer_node_name) -> QuantizedValue:
131 if self.converted is None: # Quantized value is not converted, return original
132 return self.original
133
134 if self.converted_recv_nodes is None: # All consumers receive the converted value
135 return self.converted
136
137 # Check if consumer node name is in the list of nodes that
138 # receive the converted quantization value. If not, return the original value generated
139 # by the tensor's producer.
140 return self.converted if (consumer_node_name in self.converted_recv_nodes) else self.original
141
142
143class QDQQuantizer(BaseQuantizer):
144 def __init__(
145 self,
146 model,
147 per_channel,
148 reduce_range,
149 weight_qType,
150 activation_qType,
151 tensors_range,
152 nodes_to_quantize,
153 nodes_to_exclude,
154 op_types_to_quantize,
155 extra_options=None,
156 ):
157 BaseQuantizer.__init__(
158 self,
159 model,
160 per_channel,
161 reduce_range,
162 weight_qType,
163 activation_qType,
164 tensors_range,
165 nodes_to_quantize,
166 nodes_to_exclude,
167 op_types_to_quantize,
168 extra_options,
169 )
170 self.tensors_to_quantize: dict[str, QDQTensorQuantInfo] = {}
171 self.bias_to_quantize: dict[str, QDQBiasQuantInfo] = {}
172
173 self.nodes_to_remove = []
174
175 # Specific op types to exclude qdq quantization for their outputs.
176 # In TRT, it's not recommended to quantize outputs for weighted ops such as Conv, Matmul, Gemm
177 # because those ops may be followed by nodes that require high resolution inputs.
178 # Adding QDQ for those ops' output may end up with worse accuracy.
179 # So, we don't recommend to add QDQ to node's output under such condition.
180 self.op_types_to_exclude_output_quantization = extra_options.get("OpTypesToExcludeOutputQuantization", [])
181
182 # We do quantization on Dequantizelinear's input to remove Quantizelinear for weight as an optimization.
183 # In some cases, for example QDQ BERT model for TensorRT, QDQ should always appear as a pair.
184 # Therefore, we need to disable this optimization and add qdq pair to weight.
185 self.add_qdq_pair_to_weight = extra_options.get("AddQDQPairToWeight", False)
186
187 # Some scenarios do not need the bias quantized. For example, in the case of Quantization Aware Training,
188 # quantizing the bias is not needed. This is because in QAT, all model parameters are expected to be in
189 # floating point format. To that end, we can use the FakeQuant operator for weights and activations that
190 # can always have QDQ pairs (by using AddQDQPairToWeight). But for biases in a quantized model, we can't use
191 # FakeQuant because it only ever appears before a DQ (since it is quantized as int32).
192 self.quantize_bias = extra_options.get("QuantizeBias", True)
193
194 # The default behavior is that multiple nodes can share a QDQ pair as their inputs.
195 # In TRT, QDQ pair can`t be shared between nodes, so it will create dedicated QDQ pairs for each node.
196 self.dedicated_qdq_pair = extra_options.get("DedicatedQDQPair", False)
197 self.tensor_to_its_receiving_nodes: dict[str, list[onnx.NodeProto]] = {}
198
199 # Maps a tensor to the DequantizeLinear node (in the original input model) that outputs the tensor.
200 # Populated for input models with some pre-quantized weights (typically via a different tool).
201 self.tensor_to_producing_dq: dict[str, onnx.NodeProto] = {}
202
203 # Let user set channel axis for specific op type and it's effective only when per channel quantization is supported and per_channel is True.
204 self.qdq_op_type_per_channel_support_to_axis = extra_options.get("QDQOpTypePerChannelSupportToAxis", {})
205
206 self.qdq_op_domain = ms_domain if extra_options.get("UseQDQContribOps", False) else None
207
208 # User can specify if removable activations, like Clip/Relu, should be kept in the graph.
209 # Used in the QDQRemovableActivation class.
210 self.qdq_keep_removable_activations = extra_options.get("QDQKeepRemovableActivations", False)
211
212 # Let user disable adjustment of weight scales for bias inputs that are quantized to int32.
213 self.qdq_disable_weight_adjust_for_int32_bias = extra_options.get("QDQDisableWeightAdjustForInt32Bias", False)
214
215 # The ONNX spec did not support 16-bit Q/DQ ops before opset 21.
216 # So, may have to override the Q/DQ op domain to 'com.microsoft' if the activation or weight types
217 # are 16-bit or 4-bit integers.
218 if self.opset_version < 21:
219 opset21_types = (TensorProto.UINT16, TensorProto.INT16, TensorProto.UINT4, TensorProto.INT4)
220 overrides_have_opset21_types = any(
221 t.tensor_type in opset21_types for t in self.tensor_quant_override_qtypes
222 )
223 if not self.qdq_op_domain and (
224 self.activation_qType in opset21_types
225 or self.weight_qType in opset21_types
226 or overrides_have_opset21_types
227 ):
228 logging.warning(
229 "ONNX QuantizeLinear and DequantizeLinear operators do not support "
230 "16-bit/4-bit integer quantization types prior to opset 21. "
231 f"The domain of QuantizeLinear and DequantizeLinear operators will be set to '{ms_domain}' to "
232 "enable support."
233 )
234 self.qdq_op_domain = ms_domain
235
236 self.quantization_params = self.calc_graph_quant_params()
237 self.initializer_quant_params: dict[str, QuantizationParams] = {}
238
239 # Map of all original value names to quantized value names
240 self.quantized_value_map = {}
241
242 def _get_tensor_type(self, tensor_name):
243 """
244 Check if tensor can be quantized
245 """
246 weight = find_by_name(tensor_name, self.model.initializer())
247 if weight is not None:
248 return weight.data_type
249 elif tensor_name in self.value_infos:
250 vi = self.value_infos[tensor_name]
251 if vi.type.HasField("tensor_type"):
252 return vi.type.tensor_type.elem_type
253 return None
254
255 def _is_tensor_quantizable(self, tensor_name):
256 """
257 Check if tensor can be quantized
258 """
259 weight = find_by_name(tensor_name, self.model.initializer())
260 if weight is not None:
261 if weight.data_type in (onnx_proto.TensorProto.FLOAT, onnx_proto.TensorProto.FLOAT16):
262 return True
263 elif tensor_name in self.value_infos:
264 vi = self.value_infos[tensor_name]
265 if vi.type.HasField("tensor_type") and vi.type.tensor_type.elem_type in (
266 TensorProto.FLOAT,
267 TensorProto.FLOAT16,
268 ):
269 return True
270 else:
271 logging.warning(
272 f"failed to infer the type of tensor: {tensor_name}. Skip to quantize it. Please check if it is expected."
273 )
274
275 return False
276
277 def __quantize_tensor(self, tensor_name, quant_sharing_provider=None, tensor_type=QDQQuantTensorType.ACTIVATION):
278 """
279 Adds a tensor to the list (actually a dict) of tensors to quantize. Called indirectly by op quantizers that
280 want to quantize a tensor (i.e., "mark" a tensor for quantization).
281
282 If quant_sharing_provider is not None, tensor with name tensor_name will be quantized with the same
283 quantization parameters as the node input specified in quant_sharing_provider. Ex: A Tranpose node's output
284 will typically use the same quantization parameter initializers used at the Transpose node's input.
285
286 Args:
287 tensor_name: name of the tensor to quantize
288 quant_sharing_provider: name of the tensor and node that provides quantization parameter
289 tensor_type: QDQQuantTensorType default ACTIVATION
290 """
291 if self._is_tensor_quantizable(tensor_name):
292 if quant_sharing_provider:
293 if not isinstance(quant_sharing_provider, QDQQuantParamProvider):
294 raise TypeError(
295 f"quant_sharing_provider must be of type QDQQuantParamProvider, not {type(quant_sharing_provider)}."
296 )
297
298 data_type = self._get_tensor_type(tensor_name)
299 self.tensors_to_quantize[tensor_name] = QDQTensorQuantInfo(
300 tensor_type=tensor_type, quant_para_provider=quant_sharing_provider, data_type=data_type
301 )
302 elif tensor_name not in self.tensors_to_quantize:
303 data_type = self._get_tensor_type(tensor_name)
304 self.tensors_to_quantize[tensor_name] = QDQTensorQuantInfo(tensor_type=tensor_type, data_type=data_type)
305
306 def quantize_activation_tensor(self, tensor_name: str):
307 """
308 Adds a tensor to the list of tensors to quantize. Called by op quantizers that
309 want to quantize a tensor (i.e., "mark" a tensor for quantization).
310
311 Args:
312 tensor_name: name of the tensor to quantize
313 """
314 return self.__quantize_tensor(tensor_name, None, QDQQuantTensorType.ACTIVATION)
315
316 def quantize_output_same_as_input(self, output_name: str, input_name: str, node_name: str):
317 """
318 Adds a tensor to the list of tensors to quantize. Called by op quantizers that
319 want to quantize an output tensor using the same quantization parameters as one of the node's inputs.
320
321 Ex: A Tranpose node's output will typically use the same quantization parameter initializers used at
322 the Transpose node's input.
323
324 Args:
325 output_name: name of the node output to quantize so that it uses the same quantization params as an input.
326 input_name: name of the node input from which the output tensor will get its quantization params.
327 node_name: name of the node that consumes `input_name`.
328 """
329 return self.__quantize_tensor(
330 output_name, QDQQuantParamProvider(input_name, node_name), QDQQuantTensorType.ACTIVATION
331 )
332
333 def quantize_weight_tensor(self, tensor_name: str):
334 """
335 Adds a tensor to the list of weight tensors to quantize. Called by op quantizers that
336 want to quantize a weight (i.e., "mark" a weight for quantization).
337
338 Args:
339 tensor_name: name of the weight to quantize
340 """
341 return self.__quantize_tensor(tensor_name, None, QDQQuantTensorType.WEIGHT)
342
343 def quantize_weight_tensor_per_channel(self, tensor_name, axis):
344 weight = find_by_name(tensor_name, self.model.initializer())
345 if weight:
346 if weight.data_type in (onnx_proto.TensorProto.FLOAT, onnx_proto.TensorProto.FLOAT16):
347 self.tensors_to_quantize[tensor_name] = QDQTensorQuantInfo(
348 tensor_type=QDQQuantTensorType.WEIGHT, axis=axis, data_type=weight.data_type
349 )
350 else:
351 logging.warning(f"only support per-channel quantization on weight. Tensor: {tensor_name} is not quantized.")
352
353 def _dup_initializer(self, initializer: onnx.TensorProto) -> onnx.TensorProto:
354 """
355 Duplicates an existing initializer and adds it to the model. Returns the new initializer.
356 """
357 name_suffix: int = self.model.get_largest_initializer_name_suffix(initializer.name) + 1
358 new_initializer_name = f"{initializer.name}{name_suffix}"
359 new_initializer = onnx.TensorProto()
360 new_initializer.CopyFrom(initializer)
361 new_initializer.name = new_initializer_name
362 self.model.add_initializer(new_initializer)
363 return new_initializer
364
365 def quantize_bias_tensor(self, node_name, bias_name, input_name, weight_name, beta=1.0):
366 """
367 Adds a bias tensor to the list of bias tensors to quantize. Called by op quantizers that
368 want to quantize a bias with bias_zero_point = 0 and bias_scale = input_scale * weight_scale * beta.
369 TODO: Explain the reasoning for using this formula.
370
371 Args:
372 node_name: name of the node that consumes the bias, input, and weight tensors.
373 bias_name: name of the bias tensor to quantize.
374 input_name: name of the input tensor whose scale is used to compute the bias's scale.
375 weight_name: name of the weight tensor whose scale is used to compute the bias's scale.
376 beta: Multiplier used to compute the bias's scale.
377 """
378 # If the user provided quantization overrides for this tensor, treat it as a regular weight.
379 if self.tensor_quant_overrides.get(bias_name):
380 logging.info(
381 f"Quantizing bias tensor '{bias_name}' as a weight due to the presence of user-specified overrides"
382 )
383 is_per_channel, axis = self.is_tensor_per_channel(bias_name, default_axis=0)
384 if is_per_channel:
385 self.quantize_weight_tensor_per_channel(bias_name, axis)
386 else:
387 self.quantize_weight_tensor(bias_name)
388 return
389
390 bias_initializer = find_by_name(bias_name, self.model.initializer())
391 if bias_initializer is None:
392 logging.warning(f"Expected bias '{bias_name}' to be an initializer")
393 return
394
395 if bias_initializer.data_type not in (onnx_proto.TensorProto.FLOAT, onnx_proto.TensorProto.FLOAT16):
396 logging.info(f"Expected bias '{bias_name}' to be an floating-point initializer")
397 return
398
399 actual_bias_name = bias_name
400 if bias_name in self.bias_to_quantize:
401 # This bias input is consumed by two different nodes. We need to duplicate the bias so that
402 # each node has its own bias input. This is necessary because the bias's scale is computed
403 # from the node's other input scales.
404 new_bias_initializer = self._dup_initializer(bias_initializer)
405 actual_bias_name = new_bias_initializer.name
406
407 # Replace this node's bias input
408 self.model.replace_input_of_nodes(bias_name, actual_bias_name, {node_name})
409 logging.info(f"Created a copy of bias input '{bias_name}' called '{actual_bias_name}'")
410
411 # Add this to our list of biases to quantize.
412 self.bias_to_quantize[actual_bias_name] = QDQBiasQuantInfo(node_name, input_name, weight_name, beta)
413
414 def _adjust_weight_scale_for_int32_bias(
415 self,
416 input_scale: np.ndarray,
417 weight_scale: np.ndarray,
418 weight_name: str,
419 bias_tp: onnx.TensorProto,
420 is_per_channel: bool,
421 ) -> tuple[bool, np.ndarray | None]:
422 """
423 Checks if the bias scale (input_scale * weight_scale) that we intend to use is too small.
424 A bias scale that is too small leads to quantized bias values that fall outside the range of a int32 and have to
425 be clipped, which decreases accuracy. If this function detects such a scenario, the weight_scale value will be
426 increased to prevent this from happening.
427
428 Although the adjustment method and amount differs, the idea to adjust the weight's scale came from the following
429 reference:
430 https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/tools/optimize/quantization_utils.cc#L252
431
432 :param input_scale: The input's scale.
433 :param weight_scale: The weight scale to potentially adjust.
434 :param weight_name: The weight initializer's name. Used for logging.
435 :param bias_tp: The bias ONNX initializer.
436 :param is_per_channel: True if the bias and weight are quantized per-channel.
437 :return: A tuple with a bool indicating if the weight's scale was adjusted and the new weight scale.
438 """
439 if not weight_scale.size:
440 return False, None
441
442 bias_float_data = tensor_proto_to_array(bias_tp)
443
444 int32_info = np.iinfo(np.int32)
445 multiplicative_epsilon = 1.0001
446 qrange = np.array(int32_info.max, dtype=np.float64) - np.array(int32_info.min + 1, dtype=np.float64)
447 weight_scale_dtype = weight_scale.dtype
448 updated_an_elem = False
449
450 if not is_per_channel:
451 rmin = np.minimum(bias_float_data.min(), np.array(0, dtype=np.float64))
452 rmax = np.maximum(bias_float_data.max(), np.array(0, dtype=np.float64))
453 absmax = np.maximum(np.abs(rmin), np.abs(rmax))
454 bias_smallest_valid_scale = multiplicative_epsilon * (2.0 * absmax) / qrange
455
456 input_scale_fp64 = np.array(input_scale.item(), dtype=np.float64)
457 weight_scale_fp64 = np.array(weight_scale.item(), dtype=np.float64)
458 bias_candidate_scale = input_scale_fp64 * weight_scale_fp64
459
460 if (bias_candidate_scale < bias_smallest_valid_scale) and (bias_candidate_scale > 0.0):
461 # The candidate bias scale would be too small, so increase the weight_scale by the necessary ratio.
462 ratio = bias_smallest_valid_scale / bias_candidate_scale
463 logging.info(
464 f"Increasing scale for weight `{weight_name}` by the ratio {ratio} to "
465 f"ensure bias input `{bias_tp.name}` has a valid scale."
466 )
467 new_scale = weight_scale_fp64 * ratio
468 weight_scale = new_scale.astype(weight_scale_dtype)
469 updated_an_elem = True
470 elif weight_scale.shape and len(weight_scale.shape) == 1:
471 # per-channel case
472 num_elems = weight_scale.shape[0]
473
474 for i in range(num_elems):
475 bias_rmax = np.abs(bias_float_data[i])
476 bias_smallest_valid_scale = multiplicative_epsilon * (2.0 * bias_rmax) / qrange
477
478 input_scale_fp64 = np.array(input_scale.item(), dtype=np.float64)
479 weight_scale_fp64 = np.array(weight_scale[i].item(), dtype=np.float64)
480 bias_candidate_scale = input_scale_fp64 * weight_scale_fp64
481 if (bias_candidate_scale < bias_smallest_valid_scale) and (bias_candidate_scale > 0.0):
482 # The candidate bias scale would be too small, so increase the weight_scale by the necessary ratio.
483 ratio = bias_smallest_valid_scale / bias_candidate_scale
484 logging.info(
485 f"Increased scale[{i}] for weight `{weight_name}` by ratio {ratio} "
486 f"to ensure bias input `{bias_tp.name}` has a valid scale."
487 )
488 new_scale = weight_scale_fp64 * ratio
489 weight_scale[i] = new_scale.astype(weight_scale_dtype)
490 updated_an_elem = True
491
492 return updated_an_elem, weight_scale
493
494 def _adjust_weight_quant_params_for_bias_tensors(self):
495 """
496 Iterates through all bias inputs that should be quantized to int32. If the intended
497 bias scale (equal to input_scale * weight_scale) is too small, this function will increase
498 the associated weight's scale to ensure the bias does not overflow the int32 range when quantized.
499 """
500
501 if self.qdq_disable_weight_adjust_for_int32_bias:
502 # User passed an extra_option to disable this adjustment.
503 return
504
505 for bias_name, bias_info in self.bias_to_quantize.items():
506 if (
507 bias_info.input_name not in self.quantization_params
508 or bias_info.input_name not in self.tensors_to_quantize
509 or bias_info.weight_name not in self.initializer_quant_params
510 ):
511 continue
512
513 # Get the associated input's scale.
514 input_qparams = self.quantization_params[bias_info.input_name].get_for_consumer(bias_info.node_name)
515 input_info = self.tensors_to_quantize[bias_info.input_name]
516 input_scale = np.asarray(
517 input_qparams["scale"], dtype=onnx.helper.tensor_dtype_to_np_dtype(input_info.data_type)
518 )
519
520 weight_quant_params = self.initializer_quant_params[bias_info.weight_name]
521 weight_quant_type = weight_quant_params["quant_type"]
522 if weight_quant_type not in (onnx.TensorProto.INT8, onnx.TensorProto.INT16):
523 continue
524
525 weight_zero_point: np.ndarray = weight_quant_params["zero_point"]
526 if weight_zero_point.any():
527 # Skip if zero_point(s) are not all zero (i.e., symmetric quant)
528 continue
529
530 weight_scale: np.ndarray = weight_quant_params["scale"]
531 is_per_channel = weight_quant_params.get("axis", None) is not None
532
533 # Get adjusted weight scales.
534 did_update_weight_scale, new_weight_scale = self._adjust_weight_scale_for_int32_bias(
535 input_scale,
536 weight_scale,
537 bias_info.weight_name,
538 find_by_name(bias_name, self.model.initializer()),
539 is_per_channel,
540 )
541
542 if did_update_weight_scale:
543 weight_quant_params["scale"] = new_weight_scale
544
545 def remove_node(self, node):
546 self.nodes_to_remove.append(node)
547
548 def remove_nodes(self):
549 self.model.remove_nodes(self.nodes_to_remove)
550
551 def quantize_model(self):
552 for node in self.model.nodes():
553 if self.should_quantize_node(node):
554 op_quantizer = CreateQDQQuantizer(self, node)
555 op_quantizer.quantize()
556
557 for tensor_name in node.input:
558 if tensor_name not in self.tensor_to_its_receiving_nodes:
559 self.tensor_to_its_receiving_nodes[tensor_name] = []
560 self.tensor_to_its_receiving_nodes[tensor_name].append(node)
561 if node.op_type == DEQUANT_OP_NAME:
562 for tensor_name in node.output:
563 self.tensor_to_producing_dq[tensor_name] = node
564
565 self.initializer_quant_params = self._calc_initializer_quant_params()
566 self._adjust_weight_quant_params_for_bias_tensors()
567 self._quantize_normal_tensors()
568 self._quantize_sharing_param_tensors()
569 if self.quantize_bias:
570 self._quantize_bias_tensors()
571 self.remove_nodes()
572 if not self.add_qdq_pair_to_weight:
573 self.model.clean_initializers()
574
575 self.model.model.producer_name = __producer__
576 self.model.model.producer_version = __version__
577 if self.qdq_op_domain == ms_domain:
578 self.model.set_opset_import(ms_domain, 1)
579
580 return self.model.model
581
582 def try_replacing_upstream_output(self, upstream_output_name, output_name):
583 if (
584 output_name in self.quantization_params
585 and self.quantization_params[output_name].converted is None
586 and self.quantization_params[upstream_output_name].converted is None
587 and len(self.model.input_name_to_nodes()[upstream_output_name]) == 1
588 and not self.model.is_graph_output(upstream_output_name)
589 and not self.model.is_graph_input(upstream_output_name)
590 ):
591 self.model.replace_output_of_all_nodes(upstream_output_name, output_name)
592 if upstream_output_name in self.tensors_to_quantize:
593 del self.tensors_to_quantize[upstream_output_name]
594 return True
595 return False
596
597 def _create_q_node(
598 self,
599 q_input: str,
600 q_output: str,
601 quant_node_name: str,
602 scale_name: str,
603 zp_name: str,
604 axis: int | None = None,
605 ):
606 """
607 Creates a QuantizeLinear node and adds it to the model.
608 """
609 qlinear_node = onnx.helper.make_node(
610 QUANT_OP_NAME,
611 [q_input, scale_name, zp_name],
612 [q_output],
613 quant_node_name,
614 axis=axis,
615 domain=self.qdq_op_domain,
616 )
617 self.model.add_nodes([qlinear_node])
618
619 def _create_dq_node(
620 self,
621 dq_input: str,
622 dq_output: str,
623 dequant_node_name: str,
624 scale_name: str,
625 zp_name: str,
626 axis: int | None = None,
627 ):
628 """
629 Creates a DequantizeLinear node and adds it to the model.
630 """
631 dequant_node = onnx.helper.make_node(
632 DEQUANT_OP_NAME,
633 [dq_input, scale_name, zp_name],
634 [dq_output],
635 dequant_node_name,
636 axis=axis,
637 domain=self.qdq_op_domain,
638 )
639 self.model.add_nodes([dequant_node])
640
641 def _create_qdq_nodes(
642 self, q_input, q_output, quant_node_name, dq_input, dq_output, dequant_node_name, scale_name, zp_name, axis=None
643 ):
644 qlinear_node = onnx.helper.make_node(
645 QUANT_OP_NAME,
646 [q_input, scale_name, zp_name],
647 [q_output],
648 quant_node_name,
649 axis=axis,
650 domain=self.qdq_op_domain,
651 )
652 dequant_node = onnx.helper.make_node(
653 DEQUANT_OP_NAME,
654 [dq_input, scale_name, zp_name],
655 [dq_output],
656 dequant_node_name,
657 axis=axis,
658 domain=self.qdq_op_domain,
659 )
660 self.model.add_nodes([qlinear_node, dequant_node])
661
662 def _add_qdq_nodes_for_initializer(self, weight_proto: onnx.TensorProto):
663 """
664 Adds Q/DQ nodes for an initializer. If `self.add_qdq_pair_to_weight` is true, creates
665 the sequence (weight_f32 -> Q -> DQ -> ). Otherwise, this function quantizes the initializer
666 and adds the sequence (weight_quant -> DQ ->).
667 """
668 weight_name = weight_proto.name
669 if weight_name in self.quantized_value_map:
670 return
671
672 quant_params: QuantizationParams = self.initializer_quant_params[weight_name]
673 axis: int = quant_params.get("axis")
674 scale_zp_initializers = self._make_scale_zp_initializers(weight_name, quant_params)
675 q_weight_name: str | None = None
676 weight_dequant_output = add_dequant_output_suffix(weight_name)
677 self.model.replace_input_of_all_nodes(weight_name, weight_dequant_output)
678
679 if self.add_qdq_pair_to_weight:
680 # Don't actually quantize the weight. Instead, keep floating-point weight and create the node
681 # sequence (weight_f32 -> Q -> DQ -> weight_dequant)
682 weight_quant_output = add_quant_output_suffix(weight_name)
683
684 self._create_qdq_nodes(
685 weight_name,
686 weight_quant_output,
687 add_quant_suffix(weight_name),
688 weight_quant_output,
689 weight_dequant_output,
690 add_dequant_suffix(weight_name),
691 scale_zp_initializers.scale.name,
692 scale_zp_initializers.zero_point.name,
693 axis,
694 )
695 else:
696 # Quantize the weight and create the node sequence:
697 # (weight_quantized -> DQ -> weight_dequant)
698 quant_weight = quantize_onnx_initializer(
699 weight_proto,
700 quant_params["quant_type"],
701 quant_params["zero_point"],
702 quant_params["scale"],
703 axis,
704 )
705 self.model.add_initializer(quant_weight)
706
707 q_weight_name = quant_weight.name
708 dequant_node = onnx.helper.make_node(
709 DEQUANT_OP_NAME,
710 [quant_weight.name, scale_zp_initializers.scale.name, scale_zp_initializers.zero_point.name],
711 [weight_dequant_output],
712 add_dequant_suffix(weight_name),
713 axis=axis,
714 domain=self.qdq_op_domain,
715 )
716 self.model.add_node(dequant_node)
717
718 # Log entry for this quantized weight
719 quantized_value = QuantizedValue(
720 weight_name,
721 q_weight_name,
722 scale_zp_initializers.scale.name,
723 scale_zp_initializers.zero_point.name,
724 QuantizedValueType.Initializer,
725 axis=axis,
726 )
727 self.quantized_value_map[weight_name] = QDQTensorQuantizedValue(quantized_value, None, None)
728
729 def _add_qdq_pair_for_activation(self, tensor_name, scale_name, zp_name, data_type=None):
730 if (
731 self.dedicated_qdq_pair
732 and tensor_name in self.tensor_to_its_receiving_nodes
733 and len(self.tensor_to_its_receiving_nodes[tensor_name]) > 1
734 ):
735 num_dedicated_qdq_pair = len(self.tensor_to_its_receiving_nodes[tensor_name])
736 for i in range(num_dedicated_qdq_pair):
737 postfix = f"_{i + 1}"
738 tensor_name_quant_output_postfix = add_quant_output_suffix(tensor_name) + postfix
739 tensor_name_dequant_output_postfix = add_dequant_output_suffix(tensor_name) + postfix
740 quant_node_name_postfix = add_quant_suffix(tensor_name) + postfix
741 dequant_node_name_postfix = add_dequant_suffix(tensor_name) + postfix
742 self._create_qdq_nodes(
743 tensor_name,
744 tensor_name_quant_output_postfix,
745 quant_node_name_postfix,
746 tensor_name_quant_output_postfix,
747 tensor_name_dequant_output_postfix,
748 dequant_node_name_postfix,
749 scale_name,
750 zp_name,
751 )
752
753 node = self.tensor_to_its_receiving_nodes[tensor_name][i]
754 self.model.replace_node_input(node, tensor_name, tensor_name_dequant_output_postfix)
755 if i == 0:
756 quantized_value = QuantizedValue(
757 tensor_name,
758 tensor_name_dequant_output_postfix,
759 scale_name,
760 zp_name,
761 QuantizedValueType.Input,
762 scale_type=data_type,
763 )
764 self.quantized_value_map[tensor_name] = QDQTensorQuantizedValue(quantized_value, None, None)
765 else:
766 q_input = tensor_name
767 dq_output = add_dequant_output_suffix(tensor_name)
768 if self.model.is_graph_output(tensor_name):
769 q_input = add_quant_input_suffix(tensor_name)
770 dq_output = tensor_name
771 self.model.replace_output_of_all_nodes(tensor_name, q_input)
772 else:
773 self.model.replace_input_of_all_nodes(tensor_name, dq_output)
774
775 self._create_qdq_nodes(
776 q_input,
777 add_quant_output_suffix(tensor_name),
778 add_quant_suffix(tensor_name),
779 add_quant_output_suffix(tensor_name),
780 dq_output,
781 add_dequant_suffix(tensor_name),
782 scale_name,
783 zp_name,
784 )
785
786 quantized_value = QuantizedValue(
787 tensor_name,
788 dq_output,
789 scale_name,
790 zp_name,
791 QuantizedValueType.Input,
792 scale_type=data_type,
793 )
794 self.quantized_value_map[tensor_name] = QDQTensorQuantizedValue(quantized_value, None, None)
795
796 def _add_qdq_ops_for_converted_activation(
797 self,
798 tensor_name,
799 first_scale_name,
800 first_zp_name,
801 scale_data_type,
802 convert_scale_name,
803 convert_zp_name,
804 convert_recv_nodes,
805 ):
806 """
807 Adds Q and DQ ops to a tensor whose quantized data type is converted. That is, some consumers may use the
808 original data type from the producer, while other consumers use the converted data type.
809 This is generally done by adding a sequence of ops that convert from one data type (e.g., uint8) to another (e.g., uint16).
810
811 T_float ---> Quant(to u8) ---> Convert(to u16) ---> Dequant(to float) ---> T_float'
812 where Convert(to u16) is equivalent to: ---> Dequant(to float) ---> Quant(to u16) --->
813
814 This function handles the following scenarios:
815
816 1) Tensor T is not a graph output; all consumers use the converted type
817
818 <Producer> ---> Q1 ---> DQ1 ---> Q2 ---> DQ2 ---> <Consumers>
819
820 2) Tensor T is not a graph output; some consumers use the original type, others use the converted type
821
822 <Producer> ---> Q1 -+-> DQ1 ---> <Consumers of original type>
823 |
824 +-> DQ1' ---> Q2 ---> DQ2 ---> <Consumers of converted type>
825
826 3) Tensor T is a graph output; all consumers use the converted type
827
828 <Producer> ---> Q1 ---> DQ1 ---> Q2 ---> DQ2 -+-> <Consumers>
829 |
830 +-> <Graph output>
831
832 4) Tensor T is a graph output; some consumers use the original type, others use the converted type
833
834 <Producer> ---> Q1 -+-> DQ1 -+-> <Consumers of original type>
835 | |
836 | +-> <Graph output>
837 |
838 +-> DQ1' ---> Q2 ---> DQ2 ---> <Consumers of converted type>
839
840 5) Tensor T is a graph output that is not consumed by any other nodes.
841
842 <Producer> ---> Q1 ---> DQ1 ---> Q2 ---> DQ2 ---> <Graph output>
843 """
844 tensor_recv_nodes = {node.name for node in self.tensor_to_its_receiving_nodes.get(tensor_name, [])}
845
846 if (
847 self.dedicated_qdq_pair
848 and tensor_name in self.tensor_to_its_receiving_nodes
849 and len(self.tensor_to_its_receiving_nodes[tensor_name]) > 1
850 ):
851 # TODO: Add support for dedicated_qdq_pair if/when needed.
852 raise ValueError(
853 "Do not currently support converted quant_types in TensorQuantOverrides when the `dedicated_qdq_pair` extra_option is enabled"
854 )
855
856 # Determine which nodes consume the original quantized type and which nodes
857 # consume the converted quantized type.
858 original_recv_nodes = tensor_recv_nodes
859 if convert_recv_nodes is None: # In this case, all consumers receive the converted type.
860 convert_recv_nodes = tensor_recv_nodes
861 original_recv_nodes = set()
862 else:
863 original_recv_nodes = original_recv_nodes - convert_recv_nodes
864
865 all_use_converted = len(convert_recv_nodes) == len(tensor_recv_nodes)
866 is_graph_output = self.model.is_graph_output(tensor_name)
867
868 # Create first Q op.
869 first_q_input = tensor_name
870 if is_graph_output:
871 first_q_input = add_quant_input_suffix(tensor_name)
872 self.model.replace_output_of_all_nodes(tensor_name, first_q_input)
873
874 first_q_output = add_quant_output_suffix(tensor_name)
875 self._create_q_node(
876 first_q_input, first_q_output, add_quant_suffix(tensor_name), first_scale_name, first_zp_name
877 )
878
879 # Create first DQ op.
880 first_dq_output = add_dequant_output_suffix(tensor_name)
881 if is_graph_output and not all_use_converted:
882 first_dq_output = tensor_name
883 if original_recv_nodes and first_dq_output != tensor_name:
884 self.model.replace_input_of_nodes(tensor_name, first_dq_output, original_recv_nodes)
885
886 self._create_dq_node(
887 first_q_output, first_dq_output, add_dequant_suffix(tensor_name), first_scale_name, first_zp_name
888 )
889
890 # Create parallel clone of first DQ op if _not all_ consumers use the converted type.
891 # --> DQ1' --> Q2 --> DQ2 --> <Consumers of converted type>
892 #
893 # This DQ clone would only have one consumer Q node (Q2) and could be potentially fused with
894 # it by some EPs (e.g., QNN) without breaking other "node units".
895 # Ex QNN fusion:
896 # --> Convert (fused) --> DQ2 --> <Consumers of converted type>
897 second_q_input = first_dq_output
898 if not all_use_converted:
899 second_q_input = add_quant_input_suffix(f"{tensor_name}_convert")
900 self._create_dq_node(
901 first_q_output,
902 second_q_input,
903 add_dequant_suffix(f"{tensor_name}_convert_clone"),
904 first_scale_name,
905 first_zp_name,
906 )
907
908 # Create second Q op.
909 second_q_output = add_quant_output_suffix(f"{tensor_name}_convert")
910 self._create_q_node(
911 second_q_input,
912 second_q_output,
913 add_quant_suffix(f"{tensor_name}_convert"),
914 convert_scale_name,
915 convert_zp_name,
916 )
917
918 # Create second DQ op.
919 second_dq_output = add_dequant_output_suffix(f"{tensor_name}_convert")
920 if is_graph_output and all_use_converted:
921 second_dq_output = tensor_name
922 if convert_recv_nodes and second_dq_output != tensor_name:
923 self.model.replace_input_of_nodes(tensor_name, second_dq_output, convert_recv_nodes)
924 self._create_dq_node(
925 second_q_output,
926 second_dq_output,
927 add_dequant_suffix(f"{tensor_name}_convert"),
928 convert_scale_name,
929 convert_zp_name,
930 )
931
932 # Store in quantized_value_map
933 original_quantized_value = QuantizedValue(
934 tensor_name,
935 first_dq_output,
936 first_scale_name,
937 first_zp_name,
938 QuantizedValueType.Input,
939 scale_type=scale_data_type,
940 )
941 converted_quantized_value = QuantizedValue(
942 tensor_name,
943 second_dq_output,
944 convert_scale_name,
945 convert_zp_name,
946 QuantizedValueType.Input,
947 scale_type=scale_data_type,
948 )
949 self.quantized_value_map[tensor_name] = QDQTensorQuantizedValue(
950 original_quantized_value, converted_quantized_value, convert_recv_nodes
951 )
952
953 def _quantize_normal_tensors(self):
954 """
955 Adds Q/DQ ops to tensors (activations and weights) that have been marked for quantization by op quantizers.
956 """
957 for tensor_name, tensor_info in self.tensors_to_quantize.copy().items():
958 if tensor_name in self.quantized_value_map:
959 continue
960
961 if not tensor_info.is_shared:
962 # Quantize the input
963 initializer = find_by_name(tensor_name, self.model.initializer())
964 if initializer:
965 self._add_qdq_nodes_for_initializer(initializer)
966 else:
967 # Check if this tensor is already a dequantized value. If so, skip it.
968 # This happens if the original input model already has some pre-quantized weights
969 # generated by a different tool.
970 # Ex: (quantized_weight -> DequantizeLinear -> this_tensor)
971 if tensor_name in self.tensor_to_producing_dq:
972 del self.tensors_to_quantize[tensor_name]
973 continue
974
975 tensor_qparam_initializers = self._make_tensor_scale_zp_initializers(tensor_name)
976 if not tensor_qparam_initializers:
977 raise ValueError(
978 f"Quantization parameters are not specified for param {tensor_name}. "
979 "In static mode quantization params for inputs and outputs of nodes to be quantized are required."
980 )
981
982 if tensor_qparam_initializers.converted is None:
983 # Normal case: <producer> --> Q --> DQ --> <consumers>
984 self._add_qdq_pair_for_activation(
985 tensor_name,
986 tensor_qparam_initializers.original.scale.name,
987 tensor_qparam_initializers.original.zero_point.name,
988 data_type=tensor_info.data_type,
989 )
990 else:
991 # Conversion case: <producer> ---> Q1 -+-> DQ1 --> <consumers of original type>
992 # |
993 # +-> DQ1' --> Q2 --> DQ2 --> <consumers of converted type>
994 assert tensor_info.data_type == tensor_qparam_initializers.original.scale.data_type
995 self._add_qdq_ops_for_converted_activation(
996 tensor_name,
997 tensor_qparam_initializers.original.scale.name,
998 tensor_qparam_initializers.original.zero_point.name,
999 tensor_info.data_type,
1000 tensor_qparam_initializers.converted.scale.name,
1001 tensor_qparam_initializers.converted.zero_point.name,
1002 tensor_qparam_initializers.converted_recv_nodes,
1003 )
1004
1005 del self.tensors_to_quantize[tensor_name]
1006
1007 def _quantize_sharing_param_tensors(self):
1008 """
1009 Adds Q/DQ ops to tensors that have been marked for quantization by op quantizers.
1010 Only operates on tensors that want to use the quantization parameter initializers from an upstream tensor.
1011 For example, a Transpose node's output tensor will typically want to use the same quantization parameter
1012 initializers as the Transpose node's input.
1013 """
1014 while self.tensors_to_quantize:
1015 for tensor_name, tensor_info in self.tensors_to_quantize.copy().items():
1016 quant_provider = tensor_info.quant_para_provider
1017 if quant_provider and quant_provider.input_name in self.quantized_value_map:
1018 del self.tensors_to_quantize[tensor_name]
1019
1020 quantized_value = self.quantized_value_map[quant_provider.input_name].get_for_consumer(
1021 quant_provider.node_name
1022 )
1023 if self.is_input_a_initializer(tensor_name):
1024 raise ValueError("Quantization parameter shared mode is not supported for weight yet")
1025
1026 if tensor_name in self.tensor_to_producing_dq:
1027 raise ValueError(
1028 f"Quantization parameter sharing is invalid for tensor {tensor_name} "
1029 "because it has already been quantized"
1030 )
1031
1032 # Need to check if this tensor's quant_type is converted for some consumers.
1033 # If so, create new scale/zp initializers for these consumers.
1034 converted_qparam_inits = None
1035 converted_recv_nodes = None
1036 if tensor_name in self.quantization_params:
1037 tensor_params = self.quantization_params[tensor_name]
1038 if tensor_params.converted:
1039 converted_qparam_inits = self._make_scale_zp_initializers(
1040 tensor_name, tensor_params.converted, "_convert"
1041 )
1042 converted_recv_nodes = tensor_params.converted_recv_nodes
1043
1044 if converted_qparam_inits is None:
1045 # Normal case: <producer> --> Q_shared --> DQ_shared --> <consumers>
1046 self._add_qdq_pair_for_activation(
1047 tensor_name, quantized_value.scale_name, quantized_value.zp_name
1048 )
1049 else:
1050 # Conversion case: <producer> ---> Q_shared -+-> DQ_shared --> <consumers of original type>
1051 # |
1052 # +-> DQ_shared' --> Q2 --> DQ2 --> <consumers of converted type>
1053 self._add_qdq_ops_for_converted_activation(
1054 tensor_name,
1055 quantized_value.scale_name,
1056 quantized_value.zp_name,
1057 converted_qparam_inits.scale.data_type,
1058 converted_qparam_inits.scale.name,
1059 converted_qparam_inits.zero_point.name,
1060 converted_recv_nodes,
1061 )
1062
1063 def _quantize_bias_tensors(self):
1064 """
1065 Adds DQ ops (or Cast) for bias tensors that have been marked for quantization by op quantizers.
1066 """
1067 for bias_name, bias_info in self.bias_to_quantize.items():
1068 if bias_name in self.quantized_value_map:
1069 continue
1070 # Quantize the input
1071 self.quantize_bias_static(bias_name, bias_info)
1072 init = find_by_name(bias_name, self.model.initializer())
1073 self.model.remove_initializer(init)
1074 quant_value = self.quantized_value_map[bias_name].original
1075 if quant_value.node_type == "Cast":
1076 # simple cast to float 16 and not DequantizeLinear
1077 # cublasLtMatmul only supports (b)float16, float bias.
1078 if not isinstance(init.data_type, int):
1079 raise TypeError(f"Unexpected type {type(init.data_type)} for input={bias_info.input_name!r}")
1080 node_name = add_dequant_suffix(bias_name)
1081 dequant_node = onnx.helper.make_node(
1082 "Cast",
1083 [quant_value.q_name],
1084 [bias_name],
1085 name=node_name,
1086 to=init.data_type,
1087 )
1088 elif quant_value.node_type in (None, "DequantizeLinear"):
1089 if quant_value.node_qtype in {
1090 onnx.TensorProto.FLOAT16,
1091 onnx.TensorProto.BFLOAT16,
1092 onnx.TensorProto.FLOAT,
1093 }:
1094 raise RuntimeError(f"Unexpected quantize type {quant_value.node_qtype} for DequantizeLinear.")
1095 inputs = [quant_value.q_name, quant_value.scale_name, quant_value.zp_name]
1096 node_name = add_dequant_suffix(bias_name)
1097 if quant_value.axis is not None:
1098 dequant_node = onnx.helper.make_node(
1099 "DequantizeLinear",
1100 inputs,
1101 [bias_name],
1102 node_name,
1103 axis=quant_value.axis,
1104 domain=self.qdq_op_domain,
1105 )
1106 else:
1107 dequant_node = onnx.helper.make_node(
1108 "DequantizeLinear",
1109 inputs,
1110 [bias_name],
1111 node_name,
1112 domain=self.qdq_op_domain,
1113 )
1114 else:
1115 raise RuntimeError(f"Unexpected operator type {quant_value.node_type!r}.")
1116 self.model.add_node(dequant_node)
1117
1118 def is_tensor_quantized(self, tensor_name: str):
1119 return tensor_name in self.tensors_to_quantize or tensor_name in self.bias_to_quantize
1120
1121 def is_tensor_per_channel(
1122 self,
1123 tensor_name: str,
1124 default_axis: int,
1125 op_type: str | None = None,
1126 ) -> tuple[bool, int | None]:
1127 """
1128 Checks if a given tensor is configured to be quantized per-channel. If so, also returns the channel axis.
1129
1130 ORT only supports per-channel quantization on static weights (i.e., ONNX initializers). If the user did not provide
1131 tensor quantization overrides for this tensor, then the value of self.per_channel determines if the weight
1132 is to be quantized per-channel.
1133
1134 Params:
1135 tensor_name: The name of the tensor to check.
1136 default_axis: The default channel axis. This method checks if the normalized axis is within bounds.
1137 Can be overridden via the extra_options 'QDQOpTypePerChannelSupportToAxis'
1138 and 'TensorQuantOverrides'.
1139 op_type: Optional, defaults to None. The operator type that is the only consumer of this weight.
1140 Used to access the extra option 'QDQOpTypePerChannelSupportToAxis'.
1141 Returns:
1142 A tuple (is_per_channel, axis) in which the first element indicates whether the tensor is
1143 quantized per-channel and the second element is the channel axis.
1144 The returned axis is only None if the tensor is not per-channel or the axis is out of bounds.
1145 """
1146 weight_initializer = self.initializers.get(tensor_name)
1147 if weight_initializer is None:
1148 return False, None # Only support per-channel weights
1149
1150 if self.tensor_quant_overrides.has_per_tensor_overrides(tensor_name):
1151 return False, None # User provided per-tensor overrides for this initializer
1152
1153 has_per_chan_overrides = self.tensor_quant_overrides.has_per_channel_overrides(tensor_name)
1154 if not self.per_channel and not has_per_chan_overrides:
1155 return False, None # global self.per_channel is off and user did not provide per-channel overrides.
1156
1157 axis = self.qdq_op_type_per_channel_support_to_axis.get(op_type, default_axis) if op_type else default_axis
1158 if has_per_chan_overrides:
1159 per_chan_overrides = self.tensor_quant_overrides.get_per_channel_overrides(tensor_name)
1160 axis = per_chan_overrides[0]["axis"] # Prefer axis from user-specified tensor-level overrides if available
1161
1162 weight_rank = len(weight_initializer.dims)
1163 axis_valid, axis = normalize_axis(axis, weight_rank)
1164 if not axis_valid:
1165 logging.warning(f"Axis {axis} is out-of-range for weight '{tensor_name}' with rank {weight_rank}")
1166 return False, None
1167
1168 return True, axis
1169
1170 def _get_tensor_quantization_scale(self, tensor_name: str, consumer_node_name: str) -> np.ndarray | None:
1171 """
1172 Returns the quantization scale of a tensor that is consumed by the given node.
1173 :parameter tensor_name: The name of the tensor.
1174 :parameter consumer_node_name: The name of the node that consumes the tensor as input. Necessary in case
1175 the quantization type of the tensor was converted.
1176 Refer: QDQQuantizer::_add_qdq_ops_for_converted_activation.
1177 :returns: The quantization scale or None.
1178 """
1179 initializers = self.model.initializer()
1180 scale_initializer: onnx.TensorProto | None = None
1181
1182 if tensor_name in self.quantized_value_map:
1183 # Tensor was quantized by this tool, so get scale from initializer created by this tool run.
1184 scale_name = self.quantized_value_map[tensor_name].get_for_consumer(consumer_node_name).scale_name
1185 scale_initializer = find_by_name(scale_name, initializers)
1186 else:
1187 # Tensor was already quantized in original model, so get scale from DQ node that outputs the tensor.
1188 dq_node = self.tensor_to_producing_dq.get(tensor_name, None)
1189 if dq_node:
1190 scale_initializer = find_by_name(dq_node.input[1], initializers)
1191
1192 return tensor_proto_to_array(scale_initializer) if scale_initializer is not None else None
1193
1194 def quantize_bias_static(self, bias_name: str, bias_info: QDQBiasQuantInfo) -> str:
1195 """
1196 Quantized the bias. Zero Point == 0 and Scale == Input_Scale * Weight_Scale
1197 """
1198
1199 # Handle case where bias already in quantization map
1200 if bias_name in self.quantized_value_map: