# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations
import json
from collections.abc import MutableMapping
from dataclasses import dataclass
from typing import Any
import onnx
from .quant_utils import QuantType
@dataclass
class QuantTypeInfo:
"""
The quantization type information for a tensor override.
"""
quant_type: QuantType
symmetric: bool | None = None # If None, assumes default is used.
reduce_range: bool | None = None # If None, assumes default is used.
axis: int | None = None # If None, assumes per-tensor quantization
def __eq__(self, other: object):
if isinstance(other, QuantTypeInfo):
return (
self.quant_type == other.quant_type
and (self.symmetric is None or other.symmetric is None or self.symmetric == other.symmetric)
and (self.reduce_range is None or other.reduce_range is None or self.reduce_range == other.reduce_range)
and (self.axis == other.axis)
)
return NotImplemented
@staticmethod
def load_from_dict(
raw_dict: dict[str, Any],
default_qtype: QuantType | None = None,
default_symmetric: bool | None = None,
default_reduce_range: bool | None = None,
) -> QuantTypeInfo:
return QuantTypeInfo(
raw_dict.get("quant_type", default_qtype),
raw_dict.get("symmetric", default_symmetric),
raw_dict.get("reduce_range", default_reduce_range),
raw_dict.get("axis"),
)
def save_to_dict(self, raw_dict: dict[str, Any]):
raw_dict["quant_type"] = self.quant_type
if self.symmetric is not None:
raw_dict["symmetric"] = self.symmetric
if self.reduce_range is not None:
raw_dict["reduce_range"] = self.reduce_range
if self.axis is not None:
raw_dict["axis"] = self.axis
class TensorQuantOverridesHelper(MutableMapping):
"""
Utility wrapper over the tensor quantization overrides passed via extra_options.
"""
def __init__(self, raw_overrides: dict[str, list[dict[str, Any]]]):
self.overrides = raw_overrides
self.quant_types = None
self.keys_unsupported_with_scale_zp = {"symmetric", "reduce_range", "rmax", "rmin"}
def has_per_tensor_overrides(self, tensor_name: str) -> bool:
overrides_list = self.overrides.get(tensor_name)
return overrides_list and "axis" not in overrides_list[0]
def has_per_channel_overrides(self, tensor_name: str) -> bool:
overrides_list = self.overrides.get(tensor_name)
return overrides_list and "axis" in overrides_list[0]
def overrides_scale_zp(self, tensor_name: str) -> bool:
overrides_list = self.overrides.get(tensor_name)
return overrides_list and ("scale" in overrides_list[0]) and ("zero_point" in overrides_list[0])
def get_per_tensor_overrides(
self,
tensor_name: str,
default_val: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
default_list_val = [default_val] if default_val is not None else None
overrides_list = self.overrides.get(tensor_name, default_list_val)
if overrides_list and "axis" in overrides_list[0]:
raise ValueError(
f"Expected tensor '{tensor_name}' to use per-tensor quantization overrides, "
f"but found per-channel overrides."
)
return overrides_list[0] if overrides_list else None
def get_per_channel_overrides(
self,
tensor_name: str,
default_val: list[dict[str, Any]] | None = None,
) -> list[dict[str, Any]] | None:
overrides_list = self.overrides.get(tensor_name, default_val)
if not overrides_list:
return None
if "axis" not in overrides_list[0]:
raise ValueError(
f"Expected tensor '{tensor_name}' to have per-channel quantization overrides (axis value is missing).",
)
return overrides_list
def get_quant_types(self) -> set[QuantType]:
if self.quant_types is not None:
return self.quant_types
self.quant_types = set()
if self.overrides:
for quant_overrides_list in self.overrides.values():
for quant_overrides in quant_overrides_list:
if "quant_type" in quant_overrides:
self.quant_types.add(quant_overrides["quant_type"])
if "convert" in quant_overrides and "quant_type" in quant_overrides["convert"]:
self.quant_types.add(quant_overrides["convert"]["quant_type"])
return self.quant_types
def _is_valid_per_tensor(
self,
initializers,
default_activation_qtype,
tensor_name: str,
quant_overrides: dict[str, Any],
) -> tuple[bool, str | None]:
if not isinstance(quant_overrides, dict):
return (
False,
f"Tensor quantization overrides for '{tensor_name}' are not in a dict",
)
is_initializer = tensor_name in initializers
quant_type = quant_overrides.get("quant_type")
if quant_type:
self.quant_types.add(quant_type)
has_scale = "scale" in quant_overrides
has_zero_point = "zero_point" in quant_overrides
if (has_scale and not has_zero_point) or (has_zero_point and not has_scale):
return (
False,
"Must provide both 'scale' and 'zero_point' if one of the overrides is provided",
)
if has_scale:
keys = self.keys_unsupported_with_scale_zp.intersection(set(quant_overrides))
if keys:
return (
False,
f"Tensor override option(s) [{', '.join(keys)}] are invalid with 'scale' and 'zero_point'",
)
if "reduce_range" in quant_overrides and not is_initializer:
return (
False,
f"Option 'reduce_range' is only supported for initializers, not for activation {tensor_name}",
)
if "convert" in quant_overrides:
if is_initializer:
return False, "Cannot use 'convert' override for initializers"
if "quant_type" not in quant_overrides["convert"]:
return False, f"'convert' options (tensor '{tensor_name}') must specify a 'quant_type'"
if "reduce_range" in quant_overrides["convert"]:
return (
False,
f"Option 'reduce_range' is only supported for initializers, not for activation {tensor_name}",
)
convert_quant_type = quant_overrides["convert"]["quant_type"]
original_quant_type = quant_type if quant_type is not None else default_activation_qtype
if convert_quant_type == original_quant_type:
return (
False,
f"'convert' quant_type must differ from original quant_type (tensor '{tensor_name}')",
)
convert_has_scale = "scale" in quant_overrides["convert"]
convert_has_zero_point = "zero_point" in quant_overrides["convert"]
if (convert_has_scale and not convert_has_zero_point) or (convert_has_zero_point and not convert_has_scale):
return (
False,
f"Must provide both 'scale' and 'zero_point' if one of the overrides is provided (tensor '{tensor_name}')",
)
if convert_has_scale:
keys = self.keys_unsupported_with_scale_zp.intersection(set(quant_overrides["convert"]))
if keys:
return (
False,
f"Tensor override option(s) [{', '.join(keys)}] are invalid with 'scale' and 'zero_point' "
f"(tensor '{tensor_name}')",
)
self.quant_types.add(convert_quant_type)
return True, None
def _is_valid_per_channel(
self,
initializers,
tensor_name: str,
quant_overrides_list: list[dict[str, Any]],
) -> tuple[bool, str | None]:
is_initializer = tensor_name in initializers
if not is_initializer:
return (
False,
f"Tensor '{tensor_name}' has per-channel overrides, but is not an initializer",
)
axis = quant_overrides_list[0].get("axis")
if axis is None:
return (
False,
f"Per-channel overrides for tensor {tensor_name} is missing an 'axis' value in "
"the first channel dictionary.",
)
weight_shape = list(initializers[tensor_name].dims)
weight_rank = len(weight_shape)
norm_axis = axis
if norm_axis < 0:
norm_axis += weight_rank
if norm_axis < 0 or norm_axis >= len(weight_shape):
return (
False,
f"Axis override value is out-of-bounds for tensor {tensor_name} (rank {len(weight_shape)})",
)
if len(quant_overrides_list) > 1 and len(quant_overrides_list) != weight_shape[norm_axis]:
return (
False,
f"Incorrect number of channel overrides for tensor {tensor_name} (axis {axis}), "
f"expected {weight_shape[axis]}, but found {len(quant_overrides_list)}.",
)
if "convert" in quant_overrides_list[0]:
return False, f"Cannot use 'convert' override for initializers, such as {tensor_name}."
quant_type = quant_overrides_list[0].get("quant_type")
if quant_type:
self.quant_types.add(quant_type)
symmetric = quant_overrides_list[0].get("symmetric")
reduce_range = quant_overrides_list[0].get("reduce_range")
has_scale = "scale" in quant_overrides_list[0]
has_zero_point = "zero_point" in quant_overrides_list[0]
has_scale_zp = has_scale and has_zero_point
if (has_scale and not has_zero_point) or (has_zero_point and not has_scale):
return (
False,
"Must provide both 'scale' and 'zero_point' if one of the overrides is provided",
)
if has_scale_zp:
keys = self.keys_unsupported_with_scale_zp.intersection(set(quant_overrides_list[0]))
if keys:
return (
False,
f"Tensor override option(s) [{', '.join(keys)}] are invalid with 'scale' and 'zero_point'",
)
has_rmin = "rmin" in quant_overrides_list[0]
has_rmax = "rmax" in quant_overrides_list[0]
has_rmin_rmax = has_rmin and has_rmax
if (has_rmin and not has_rmax) or (not has_rmin and has_rmax):
return (
False,
"Must provide both 'rmin' and 'rmax' if one is provided",
)
for index, quant_overrides in enumerate(quant_overrides_list[1:]):
if not isinstance(quant_overrides, dict):
return (
False,
f"Tensor quantization overrides at index {index} for '{tensor_name}' are not in a dict",
)
if "convert" in quant_overrides:
return False, f"Cannot use 'convert' override for initializers, such as {tensor_name}."
# For per-channel quantization, all channels must use the same quantization type, axis, symmetric
# and reduce_range values. And, if specified, they must be present in the first channel dict
# (i.e., quant_overrides_list[0]).
if "quant_type" in quant_overrides and quant_type != quant_overrides["quant_type"]:
return (
False,
"Channel quantization types for tensor '{tensor_name}' do not match at index {index}.",
)
if "axis" in quant_overrides and axis != quant_overrides["axis"] and norm_axis != quant_overrides["axis"]:
return (
False,
"Channel axis for tensor '{tensor_name}' does not match at index {index}.",
)
if "symmetric" in quant_overrides and symmetric != quant_overrides["symmetric"]:
return (
False,
"Channel symmetric value for tensor '{tensor_name}' does not match at index {index}.",
)
if "reduce_range" in quant_overrides and reduce_range != quant_overrides["reduce_range"]:
return (
False,
"Channel reduce_range value for tensor '{tensor_name}' does not match at index {index}.",
)
# If override scale/zp, must do so for all channels.
chan_has_scale_zp = "scale" in quant_overrides and "zero_point" in quant_overrides
if has_scale_zp and not chan_has_scale_zp:
return (
False,
"Per-channel overrides that specify scale/zero_point must do so for all channels, "
f"but tensor '{tensor_name}' is missing them at index {index}.",
)
if chan_has_scale_zp:
keys = self.keys_unsupported_with_scale_zp.intersection(set(quant_overrides))
if keys:
return (
False,
f"Tensor override option(s) [{', '.join(keys)}] are invalid with 'scale' and 'zero_point'",
)
# If override rmin/rmax, must do so for all channels.
chan_has_rmin_rmax = "rmin" in quant_overrides and "rmax" in quant_overrides
if has_rmin_rmax and not chan_has_rmin_rmax:
return (
False,
"Per-channel overrides that specify rmin/rmax must do so for all channels, "
f"but tensor '{tensor_name}' is missing them at index {index}.",
)
return True, None
def is_valid(
self,
initializers: dict[str, onnx.TensorProto],
activation_names: set[str],
default_activation_qtype,
) -> tuple[bool, str | None]:
self.quant_types = set()
# Validate that compatible/valid overrides are provided.
if self.overrides:
for tensor_name, quant_overrides_list in self.overrides.items():
if tensor_name not in initializers and tensor_name not in activation_names:
return False, f"Tensor '{tensor_name}' in TensorQuantOverrides is not present in the model"
if not isinstance(quant_overrides_list, list):
return False, f"Tensor quantization overrides for '{tensor_name}' are not in a list"
if not quant_overrides_list:
continue
if not isinstance(quant_overrides_list[0], dict):
return False, f"Tensor quantization overrides at index 0 for '{tensor_name}' are not in a dict"
if not quant_overrides_list[0]:
continue
axis = quant_overrides_list[0].get("axis")
is_per_channel = len(quant_overrides_list) > 1 or axis is not None
if is_per_channel:
return self._is_valid_per_channel(initializers, tensor_name, quant_overrides_list)
return self._is_valid_per_tensor(
initializers, default_activation_qtype, tensor_name, quant_overrides_list[0]
)
return True, None
def update_tensor_overrides(
self,
tensor_name: str,
new_vals: dict[str, Any],
channels: list[int] | None = None,
overwrite: bool = True,
) -> bool:
if not new_vals:
return False
channels = set(channels) if channels is not None else None
have_overrides = self.overrides.get(tensor_name)
# If `overwrite` is False, check if we would overwrite anything.
do_update = True
if not overwrite and have_overrides:
for channel, overrides in enumerate(self.overrides[tensor_name]):
if channels is not None and channel not in channels:
continue
if set(new_vals).intersection(set(overrides)):
do_update = False
break
# Do the update if `overwrite` is True or if nothing is overwritten (do not want partial overwrites).
if do_update:
if not have_overrides:
self.overrides[tensor_name] = [{}]
for channel, overrides in enumerate(self.overrides[tensor_name]):
if channels is not None and channel not in channels:
continue
overrides.update(new_vals)
return do_update
def get_node_output_qtype_info(
self,
output_name: str,
default_qtype: QuantType | None,
default_symmetric: bool | None = None,
) -> QuantTypeInfo:
# Outputs are activations, which do not support 'reduce_range' or 'axis'
if output_name not in self.overrides:
return QuantTypeInfo(default_qtype, default_symmetric)
tensor_overrides = self.overrides[output_name][0]
return QuantTypeInfo(
tensor_overrides.get("quant_type", default_qtype),
tensor_overrides.get("symmetric", default_symmetric),
)
def get_node_input_qtype_info(
self,
input_name: str,
node_name: str,
default_qtype: QuantType | None,
default_symmetric: bool | None = None,
default_reduce_range: bool | None = None,
) -> QuantTypeInfo:
if input_name not in self.overrides or not self.overrides[input_name]:
return QuantTypeInfo(default_qtype, default_symmetric, default_reduce_range)
# Get the first overrides dict in the list. This works for both per-tensor and per-channel
# quantization because all channels must use the same quant type.
tensor_overrides = self.overrides[input_name][0]
producer_type = tensor_overrides.get("quant_type", default_qtype)
if "convert" not in tensor_overrides:
return QuantTypeInfo(
producer_type,
tensor_overrides.get("symmetric", default_symmetric),
tensor_overrides.get("reduce_range", default_reduce_range),
tensor_overrides.get("axis"),
)
# This tensor is converted. Check if the node gets the original qtype or the converted qtype.
convert_dict = tensor_overrides["convert"]
qtype_info = QuantTypeInfo(
producer_type,
convert_dict.get("symmetric", default_symmetric),
# Converted tensors are not initializers, so do not have 'axis' or 'reduce_range'.
)
# Check if all nodes receive the converted type (i.e., recv_nodes is None) or this node
# is in the list of consumers (recv_nodes).
if ("recv_nodes" not in convert_dict) or (node_name in convert_dict["recv_nodes"]):
qtype_info.quant_type = convert_dict["quant_type"]
return qtype_info
def pprint_str(self, indent=None) -> str:
return json.dumps(self.overrides, default=str, indent=indent)
def empty(self) -> bool:
return not self.overrides
def get_dict(self) -> dict[str, list[dict[str, Any]]]:
return self.overrides
# Required implementations of abstract methods in collections.abc.MutableMapping
# so that this class can be used like a dict.
def __setitem__(self, key: str, value: list[dict]):
self.overrides[key] = value
def __getitem__(self, key: str) -> list[dict]:
return self.overrides[key]
def __delitem__(self, key: str):
del self.overrides[key]
def __iter__(self):
return iter(self.overrides)
def __len__(self):
return len(self.overrides)
def __str__(self) -> str:
return str(self.overrides)
def __repr__(self) -> str:
return f"{super().__repr__()}, TensorQuantOverridesHelper({self.overrides})"