1286 lines
50 KiB
Python
1286 lines
50 KiB
Python
# -------------------------------------------------------------------------
|
|
# 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 copy
|
|
import logging
|
|
import os
|
|
import tempfile
|
|
from enum import Enum
|
|
from pathlib import Path
|
|
|
|
import numpy
|
|
import onnx
|
|
from ml_dtypes import float8_e4m3fn, int4, uint4
|
|
from onnx import ModelProto, TensorProto, external_data_helper
|
|
from onnx import onnx_pb as onnx_proto
|
|
from onnx.helper import make_graph, make_model, make_node, make_tensor_value_info
|
|
from onnx.reference import ReferenceEvaluator
|
|
|
|
from onnxruntime import GraphOptimizationLevel, InferenceSession, SessionOptions
|
|
|
|
try:
|
|
from onnx.reference.op_run import to_array_extended
|
|
except ImportError:
|
|
# old version of onnx.
|
|
to_array_extended = None
|
|
|
|
|
|
__producer__ = "onnx.quantize"
|
|
__version__ = "0.1.0"
|
|
onnx_domain = "ai.onnx"
|
|
ms_domain = "com.microsoft"
|
|
QUANT_OP_NAME = "QuantizeLinear"
|
|
QUANT_INPUT_SUFFIX = "_QuantizeLinear_Input"
|
|
DEQUANT_OP_NAME = "DequantizeLinear"
|
|
DEQUANT_OUTPUT_SUFFIX = "_DequantizeLinear_Output"
|
|
TENSOR_NAME_QUANT_SUFFIX = "_quantized"
|
|
MODEL_SIZE_THRESHOLD = 2147483648 # Quant model should use external data if >= 2GB
|
|
|
|
FLOAT8_DISTRIBUTIONS = {}
|
|
FLOAT8_TYPES = (
|
|
onnx_proto.TensorProto.FLOAT8E4M3FN,
|
|
onnx_proto.TensorProto.FLOAT8E4M3FNUZ,
|
|
onnx_proto.TensorProto.FLOAT8E5M2,
|
|
onnx_proto.TensorProto.FLOAT8E5M2FNUZ,
|
|
)
|
|
|
|
type_to_name = {getattr(TensorProto, k): k for k in dir(TensorProto) if isinstance(getattr(TensorProto, k), int)}
|
|
|
|
# Quantization mode
|
|
# IntegerOps: Use IntegerOps in quantized model. Only ConvInteger and MatMulInteger ops are supported now.
|
|
# QLinearOps: Use QLinearOps in quantized model. Only QLinearConv and QLinearMatMul ops are supported now.
|
|
|
|
|
|
class QuantizationMode(Enum):
|
|
IntegerOps = 0
|
|
QLinearOps = 1
|
|
|
|
def __str__(self):
|
|
return self.name
|
|
|
|
@staticmethod
|
|
def from_string(mode):
|
|
try:
|
|
return QuantizationMode[mode]
|
|
except KeyError:
|
|
raise ValueError() # noqa: B904
|
|
|
|
|
|
class QuantizedValueType(Enum):
|
|
Input = 0
|
|
Initializer = 1
|
|
|
|
def __str__(self):
|
|
return self.name
|
|
|
|
@staticmethod
|
|
def from_string(v):
|
|
try:
|
|
return QuantizedValueType[v]
|
|
except KeyError:
|
|
raise ValueError() # noqa: B904
|
|
|
|
|
|
class QuantType(Enum):
|
|
QInt8 = 0
|
|
QUInt8 = 1
|
|
QFLOAT8E4M3FN = 2
|
|
QInt16 = 3
|
|
QUInt16 = 4
|
|
QInt4 = 5
|
|
QUInt4 = 6
|
|
|
|
def __str__(self):
|
|
return self.name
|
|
|
|
@staticmethod
|
|
def from_string(t):
|
|
try:
|
|
return QuantType[t]
|
|
except KeyError:
|
|
raise ValueError() # noqa: B904
|
|
|
|
@property
|
|
def tensor_type(self):
|
|
if self == QuantType.QInt8:
|
|
return TensorProto.INT8
|
|
if self == QuantType.QUInt8:
|
|
return TensorProto.UINT8
|
|
if self == QuantType.QUInt16:
|
|
return TensorProto.UINT16
|
|
if self == QuantType.QInt16:
|
|
return TensorProto.INT16
|
|
if self == QuantType.QFLOAT8E4M3FN:
|
|
return TensorProto.FLOAT8E4M3FN
|
|
if self == QuantType.QUInt4:
|
|
return TensorProto.UINT4
|
|
if self == QuantType.QInt4:
|
|
return TensorProto.INT4
|
|
raise ValueError(f"Unexpected value qtype={self!r}.")
|
|
|
|
|
|
class QuantFormat(Enum):
|
|
QOperator = 0
|
|
QDQ = 1
|
|
|
|
def __str__(self):
|
|
return self.name
|
|
|
|
@staticmethod
|
|
def from_string(format):
|
|
try:
|
|
return QuantFormat[format]
|
|
except KeyError:
|
|
raise ValueError() # noqa: B904
|
|
|
|
|
|
ONNX_TYPE_TO_NP_TYPE = {
|
|
onnx_proto.TensorProto.INT8: numpy.dtype("int8"),
|
|
onnx_proto.TensorProto.UINT8: numpy.dtype("uint8"),
|
|
onnx_proto.TensorProto.INT16: numpy.dtype("int16"),
|
|
onnx_proto.TensorProto.UINT16: numpy.dtype("uint16"),
|
|
onnx_proto.TensorProto.FLOAT8E4M3FN: float8_e4m3fn,
|
|
onnx_proto.TensorProto.INT4: int4,
|
|
onnx_proto.TensorProto.UINT4: uint4,
|
|
}
|
|
|
|
ONNX_INT_TYPE_RANGE = {
|
|
onnx_proto.TensorProto.UINT8: (numpy.array(0, dtype=numpy.uint8), numpy.array(255, dtype=numpy.uint8)),
|
|
onnx_proto.TensorProto.INT8: (numpy.array(-128, dtype=numpy.int8), numpy.array(127, dtype=numpy.int8)),
|
|
onnx_proto.TensorProto.UINT16: (numpy.array(0, dtype=numpy.uint16), numpy.array(65535, dtype=numpy.uint16)),
|
|
onnx_proto.TensorProto.INT16: (numpy.array(-32768, dtype=numpy.int16), numpy.array(32767, dtype=numpy.int16)),
|
|
onnx_proto.TensorProto.UINT4: (numpy.array(0, dtype=uint4), numpy.array(15, dtype=uint4)),
|
|
onnx_proto.TensorProto.INT4: (numpy.array(-8, dtype=int4), numpy.array(7, dtype=int4)),
|
|
}
|
|
|
|
ONNX_INT_TYPE_SYMMETRIC_RANGE = {
|
|
onnx_proto.TensorProto.INT8: (numpy.array(-127, dtype=numpy.int8), numpy.array(127, dtype=numpy.int8)),
|
|
onnx_proto.TensorProto.INT16: (numpy.array(-32767, dtype=numpy.int16), numpy.array(32767, dtype=numpy.int16)),
|
|
}
|
|
|
|
ONNX_INT_TYPE_REDUCED_RANGE = {
|
|
onnx_proto.TensorProto.UINT8: (numpy.array(0, dtype=numpy.uint8), numpy.array(127, dtype=numpy.uint8)),
|
|
onnx_proto.TensorProto.INT8: (numpy.array(-64, dtype=numpy.int8), numpy.array(64, dtype=numpy.int8)),
|
|
onnx_proto.TensorProto.UINT16: (numpy.array(0, dtype=numpy.uint16), numpy.array(32767, dtype=numpy.uint16)),
|
|
onnx_proto.TensorProto.INT16: (numpy.array(-16384, dtype=numpy.int16), numpy.array(16384, dtype=numpy.int16)),
|
|
onnx_proto.TensorProto.UINT4: (numpy.array(0, dtype=uint4), numpy.array(7, dtype=uint4)),
|
|
onnx_proto.TensorProto.INT4: (numpy.array(-4, dtype=int4), numpy.array(3, dtype=int4)),
|
|
}
|
|
|
|
|
|
def _check_type(*args, zero_point_index=-1):
|
|
new_args = []
|
|
for i, a in enumerate(args):
|
|
if numpy.issubdtype(type(a), numpy.number):
|
|
new_args.append(numpy.array(a))
|
|
elif isinstance(a, numpy.ndarray):
|
|
new_args.append(a)
|
|
else:
|
|
raise TypeError(f"arg {i} is not an array: {a}")
|
|
if i == zero_point_index:
|
|
v = new_args[-1]
|
|
if v.dtype == numpy.float32 or v.dtype == numpy.float16:
|
|
raise TypeError(f"zero_point cannot be {v.dtype}")
|
|
return tuple(new_args) if len(new_args) > 1 else new_args[0]
|
|
|
|
|
|
def quantize_nparray(qType, arr, scale, zero_point, low=None, high=None):
|
|
assert qType in ONNX_TYPE_TO_NP_TYPE, (
|
|
f"Unexpected data type {qType} requested. Only INT8, UINT8, INT16, and UINT16 are supported."
|
|
)
|
|
if qType in (
|
|
onnx_proto.TensorProto.FLOAT8E4M3FN,
|
|
onnx_proto.TensorProto.FLOAT8E4M3FNUZ,
|
|
onnx_proto.TensorProto.FLOAT8E5M2,
|
|
onnx_proto.TensorProto.FLOAT8E5M2FNUZ,
|
|
):
|
|
if zero_point != 0:
|
|
raise NotImplementedError(f"zero_point is expected to be null for float 8 not {zero_point!r}.")
|
|
if arr.dtype == numpy.float32:
|
|
onnx_type = TensorProto.FLOAT
|
|
elif arr.dtype == numpy.float16:
|
|
onnx_type = TensorProto.FLOAT16
|
|
else:
|
|
raise ValueError(f"Unexpected dtype {arr.dtype}.")
|
|
onnx_model = make_model(
|
|
make_graph(
|
|
[
|
|
make_node(
|
|
"Constant", [], ["zero_point"], value=onnx.helper.make_tensor("zero_point", qType, [], [0])
|
|
),
|
|
make_node("QuantizeLinear", ["X", "scale", "zero_point"], ["Y"]),
|
|
],
|
|
"qu",
|
|
[
|
|
make_tensor_value_info("X", onnx_type, None),
|
|
make_tensor_value_info("scale", onnx_type, None),
|
|
],
|
|
[make_tensor_value_info("Y", qType, None)],
|
|
)
|
|
)
|
|
ref = ReferenceEvaluator(onnx_model)
|
|
return _check_type(ref.run(None, {"X": arr, "scale": scale})[0])
|
|
else:
|
|
# Quantizes data for all integer types.
|
|
#
|
|
# For int4 types, the quantized data is returned as either np.int8 or np.uint8,
|
|
# which matches the python reference ONNX implementation of QuantizeLinear.
|
|
# This data can be packed into 4-bit elements by using pack_bytes_to_4bit().
|
|
dtype = ONNX_TYPE_TO_NP_TYPE[qType]
|
|
qmin, qmax = get_qmin_qmax_for_qType(qType, reduce_range=False, symmetric=False)
|
|
|
|
cliplow = max(qmin, low) if low is not None else qmin
|
|
cliphigh = min(qmax, high) if high is not None else qmax
|
|
arr_fp32 = numpy.asarray((arr.astype(numpy.float32) / scale).round() + zero_point)
|
|
numpy.clip(arr_fp32, cliplow, cliphigh, out=arr_fp32)
|
|
return _check_type(arr_fp32.astype(dtype))
|
|
|
|
|
|
def compute_scale_zp(rmin, rmax, qmin, qmax, symmetric=False, min_real_range=None):
|
|
"""Calculate the scale s and zero point z for the quantization relation
|
|
r = s(q-z), where r are the original values and q are the corresponding
|
|
quantized values.
|
|
|
|
r and z are calculated such that every value within [rmin,rmax] has an
|
|
approximate representation within [qmin,qmax]. In addition, qmin <= z <=
|
|
qmax is enforced. If the symmetric flag is set to True, the interval
|
|
[rmin,rmax] is symmetrized to [-absmax, +absmax], where
|
|
absmax = max(abs(rmin), abs(rmax)).
|
|
|
|
:parameter rmin: minimum value of r
|
|
:parameter rmax: maximum value of r
|
|
:parameter qmin: minimum value representable by the target quantization data type
|
|
:parameter qmax: maximum value representable by the target quantization data type
|
|
:parameter symmetric: True if the floating-point range should be made symmetric. Defaults to False.
|
|
:parameter min_real_range: Minimum floating-point range (i.e., rmax - rmin) to enforce. Defaults to None.
|
|
:return: zero and scale [z, s]
|
|
|
|
"""
|
|
if qmin > 0 or qmax < 0:
|
|
raise ValueError(f"qmin and qmax must meet requirement: qmin <= 0 <= qmax while qmin:{qmin}, qmmax:{qmax}")
|
|
|
|
# Adjust rmin and rmax such that 0 is included in the range. This is
|
|
# required to make sure zero can be represented by the quantization data
|
|
# type (i.e. to make sure qmin <= zero_point <= qmax)
|
|
rmin = numpy.minimum(rmin, numpy.array(0, dtype=rmin.dtype))
|
|
rmax = numpy.maximum(rmax, numpy.array(0, dtype=rmax.dtype))
|
|
|
|
# Ensure a minimum float-point range if specified.
|
|
if min_real_range is not None:
|
|
rmax = max(rmax, rmin + numpy.asarray(min_real_range, dtype=rmin.dtype))
|
|
|
|
if symmetric:
|
|
absmax = numpy.maximum(numpy.abs(rmin), numpy.abs(rmax))
|
|
rmin = -absmax
|
|
rmax = +absmax
|
|
|
|
assert qmin <= qmax, f"qmin={rmin} > qmax={rmax}"
|
|
dr = numpy.array(rmax - rmin, dtype=numpy.float64)
|
|
dq = numpy.array(qmax, dtype=numpy.float64) - numpy.array(qmin, dtype=numpy.float64)
|
|
scale = numpy.array(dr / dq)
|
|
assert scale >= 0, "scale issue"
|
|
if scale < numpy.finfo(rmax.dtype).tiny:
|
|
scale = numpy.array(1.0, dtype=rmax.dtype)
|
|
zero_point = numpy.array(0, dtype=qmin.dtype)
|
|
else:
|
|
if symmetric:
|
|
# When symmetric (i.e., rmax == -rmin), the zero_point formula reduces to round((qmax + qmin) / 2.0).
|
|
# This simpler formula doesn't depend on scale and guarantees that the zero point values
|
|
# for int8, uint8, int16, and uint16 are always 0, 128, 0, and 32768, respectively.
|
|
# This is important for per-channel/symmetric QLinearConv on CPU EP, which requires all channels to have
|
|
# the exact same zero_point values.
|
|
zero_point = numpy.array(
|
|
numpy.round((qmin + qmax) / numpy.array(2.0, dtype=numpy.float64)), dtype=qmin.dtype
|
|
)
|
|
else:
|
|
zero_point = numpy.array(numpy.round(qmin - rmin / scale), dtype=qmin.dtype)
|
|
scale = scale.astype(rmax.dtype)
|
|
|
|
return [zero_point, scale]
|
|
|
|
|
|
def snap_zero_point_to_uint8(rmin, rmax, qmin: int = 0, qmax: int = 255, min_real_range: float | None = None):
|
|
"""Snap a uint8 activation zero-point to qmin (when rmin >= 0) or mid (when rmin < 0).
|
|
|
|
Used by the ActivationRestrictedAsymmetric quantization option. Recomputes scale so the
|
|
dequantized range still covers [rmin, rmax] without clipping.
|
|
|
|
:parameter rmin: calibrated minimum activation value (numpy scalar)
|
|
:parameter rmax: calibrated maximum activation value (numpy scalar)
|
|
:parameter qmin: minimum quantized value (int, default 0)
|
|
:parameter qmax: maximum quantized value (int, default 255)
|
|
:parameter min_real_range: minimum floating-point range to enforce (same semantics as compute_scale_zp).
|
|
When not None and > 0, rmax is adjusted to max(rmax, rmin + min_real_range) before scale computation.
|
|
:return: (zero_point, scale) with zero_point dtype uint8 and scale dtype float32
|
|
"""
|
|
qmin_val = int(qmin)
|
|
qmax_val = int(qmax)
|
|
mid = (qmin_val + qmax_val + 1) // 2
|
|
|
|
rmin = float(numpy.squeeze(rmin))
|
|
rmax = float(numpy.squeeze(rmax))
|
|
|
|
# Expand the range to include zero, mirroring compute_scale_zp's ordering.
|
|
rmin = min(rmin, 0.0)
|
|
rmax = max(rmax, 0.0)
|
|
|
|
# Apply minimum real range after zero-inclusion, mirroring compute_scale_zp behaviour.
|
|
if min_real_range is not None and min_real_range > 0:
|
|
rmax = max(rmax, rmin + float(min_real_range))
|
|
|
|
if rmax <= rmin:
|
|
# Degenerate range - apply the same snap logic as the normal path, then
|
|
# compute a meaningful scale rather than a hardcoded 1.0.
|
|
degenerate_zp = qmin_val if rmin >= 0.0 else mid
|
|
abs_max = max(abs(rmin), abs(rmax))
|
|
# Use full range when zp snaps to qmin (all-positive), half range for mid snap.
|
|
denom = (qmax_val - qmin_val) if degenerate_zp == qmin_val else max(1, (qmax_val - qmin_val) // 2)
|
|
scale_val = (abs_max if abs_max > 0 else 1.0) / max(1, denom)
|
|
if min_real_range is not None and scale_val < min_real_range / (qmax_val - qmin_val):
|
|
scale_val = min_real_range / (qmax_val - qmin_val)
|
|
return numpy.array(degenerate_zp, dtype=numpy.uint8), numpy.array(scale_val, dtype=numpy.float32)
|
|
|
|
if rmin >= 0.0:
|
|
zero_point = numpy.array(qmin_val, dtype=numpy.uint8)
|
|
scale = numpy.array(rmax / (qmax_val - qmin_val), dtype=numpy.float32)
|
|
else:
|
|
# Snap zero-point to the midpoint of the quantized range.
|
|
zero_point = numpy.array(mid, dtype=numpy.uint8)
|
|
# Choose scale that covers both halves without clipping.
|
|
scale_neg = -rmin / (mid - qmin_val) # scale needed to represent rmin at q=qmin
|
|
scale_pos = rmax / (qmax_val - mid) # scale needed to represent rmax at q=qmax
|
|
scale = numpy.array(max(scale_neg, scale_pos), dtype=numpy.float32)
|
|
|
|
# Enforce minimum real range floor on scale.
|
|
if min_real_range is not None and float(scale) < min_real_range / (qmax_val - qmin_val):
|
|
scale = numpy.array(min_real_range / (qmax_val - qmin_val), dtype=numpy.float32)
|
|
|
|
return zero_point, scale
|
|
|
|
|
|
def compute_scale_zp_float8(element_type, std):
|
|
"""Calculate the scale s for a float8 type (E4M3FN).
|
|
The function assumes the coefficient distribution and the float 8
|
|
distribution are similar to two gaussian laws.
|
|
|
|
:return: zero and scale [z, s]
|
|
|
|
More details in notebook `quantization_fp8.ipynb
|
|
<https://github.com/microsoft/onnxruntime/blob/main/docs/python/notebooks/quantization_fp8.ipynb>`_.
|
|
"""
|
|
zp_dtype = None
|
|
if element_type not in FLOAT8_DISTRIBUTIONS:
|
|
if element_type == TensorProto.FLOAT8E4M3FN:
|
|
from ml_dtypes import float8_e4m3fn # noqa: PLC0415
|
|
|
|
zp_dtype = float8_e4m3fn
|
|
all_values = numpy.arange(256, dtype=numpy.uint8).view(float8_e4m3fn).astype(numpy.float32)
|
|
values = numpy.array(
|
|
[f for f in all_values if not numpy.isnan(f) and not numpy.isinf(f)], dtype=numpy.float32
|
|
)
|
|
else:
|
|
raise ValueError(f"Quantization to element_type={element_type} not implemented.")
|
|
FLOAT8_DISTRIBUTIONS[element_type] = values
|
|
elif element_type == TensorProto.FLOAT8E4M3FN:
|
|
from ml_dtypes import float8_e4m3fn # noqa: PLC0415
|
|
|
|
zp_dtype = float8_e4m3fn
|
|
|
|
if zp_dtype is None:
|
|
raise TypeError(f"Unexpected element_type {element_type}.")
|
|
std_f8 = numpy.std(FLOAT8_DISTRIBUTIONS[element_type])
|
|
zero = numpy.array(0, dtype=zp_dtype)
|
|
scale = numpy.array(std / std_f8, dtype=std.dtype)
|
|
return [zero, scale]
|
|
|
|
|
|
def compute_scale_zp_blocked(
|
|
weight: numpy.ndarray,
|
|
quant_type: int,
|
|
axis: int,
|
|
block_size: int,
|
|
symmetric: bool,
|
|
) -> tuple[numpy.ndarray, numpy.ndarray]:
|
|
"""Compute per-block scale and zero-point for a weight tensor.
|
|
|
|
The weight is sliced along *axis* into blocks of *block_size* elements.
|
|
Per the ONNX opset-21 spec, QuantizeLinear/DequantizeLinear require the
|
|
scale and zero_point tensors to have the **same rank** as the input tensor.
|
|
Only rank-2 weight tensors are supported; rank > 2 is explicitly rejected.
|
|
|
|
Returns arrays with the same rank and dimensions as *weight*, except
|
|
``shape[axis] == ceil(weight.shape[axis] / block_size)``. This matches
|
|
the ONNX opset-21 QuantizeLinear/DequantizeLinear blocked-quantization spec.
|
|
|
|
:param weight: Float32/float16 weight array (must be rank-2).
|
|
:param quant_type: ONNX tensor data type for quantization.
|
|
:param axis: Axis along which to apply block-wise quantization.
|
|
:param block_size: Number of elements per block along *axis*.
|
|
:param symmetric: Whether to use symmetric quantization per block.
|
|
:return: Tuple of (zero_point, scale), each with shape matching *weight*
|
|
except ``shape[axis] == n_blocks``, where ``n_blocks == ceil(weight.shape[axis] / block_size)``.
|
|
:raises NotImplementedError: If weight rank is not 2 (opset-21 constraint).
|
|
"""
|
|
if weight.ndim != 2:
|
|
raise NotImplementedError(
|
|
f"Per-block (opset-21) quantization is only supported for rank-2 weight tensors. "
|
|
f"Got rank-{weight.ndim} tensor with shape {weight.shape}. "
|
|
"For rank > 2 tensors, reshape to 2-D before quantizing or use per-channel quantization."
|
|
)
|
|
|
|
k = weight.shape[axis]
|
|
n_blocks = (k + block_size - 1) // block_size
|
|
|
|
# Flatten all non-axis dims into a single "other" dimension.
|
|
other = int(numpy.prod([d for i, d in enumerate(weight.shape) if i != axis]))
|
|
# Move the quantized axis to position 0 for easy slicing.
|
|
moved = numpy.moveaxis(weight, axis, 0) # shape: [k, other]
|
|
moved = moved.reshape(k, other)
|
|
|
|
qmin, qmax = get_qmin_qmax_for_qType(quant_type, reduce_range=False, symmetric=symmetric)
|
|
zp_dtype = ONNX_INT_TYPE_RANGE[quant_type][0].dtype
|
|
|
|
# Pad along axis-0 to a multiple of block_size for vectorised reduction.
|
|
pad_len = n_blocks * block_size - k
|
|
if pad_len > 0:
|
|
pad = numpy.zeros((pad_len, other), dtype=moved.dtype)
|
|
moved_padded = numpy.concatenate([moved, pad], axis=0)
|
|
else:
|
|
moved_padded = moved
|
|
|
|
# Reshape to [n_blocks, block_size, other] for axis-based min/max.
|
|
blocks = moved_padded.reshape(n_blocks, block_size, other)
|
|
|
|
# Compute per-block min/max vectorised (no Python loop over blocks or cols).
|
|
rmin = blocks.min(axis=1) # [n_blocks, other]
|
|
rmax = blocks.max(axis=1) # [n_blocks, other]
|
|
|
|
# Clamp to include zero, matching compute_scale_zp semantics.
|
|
rmin = numpy.minimum(rmin, numpy.zeros_like(rmin))
|
|
rmax = numpy.maximum(rmax, numpy.zeros_like(rmax))
|
|
|
|
if symmetric:
|
|
absmax = numpy.maximum(numpy.abs(rmin), numpy.abs(rmax))
|
|
rmin = -absmax
|
|
rmax = absmax
|
|
|
|
qmin_val = numpy.float64(qmin)
|
|
qmax_val = numpy.float64(qmax)
|
|
dr = (rmax - rmin).astype(numpy.float64)
|
|
dq = qmax_val - qmin_val
|
|
raw_scale = (dr / dq).astype(weight.dtype)
|
|
|
|
tiny = numpy.finfo(weight.dtype).tiny
|
|
degenerate = raw_scale < tiny # blocks where the float range is essentially zero
|
|
|
|
scales = numpy.where(degenerate, numpy.ones_like(raw_scale), raw_scale)
|
|
|
|
if symmetric:
|
|
zp_val = int(numpy.round((qmin_val + qmax_val) / 2.0))
|
|
zero_points = numpy.full((n_blocks, other), zp_val, dtype=zp_dtype)
|
|
else:
|
|
raw_zp = numpy.round(qmin_val - rmin.astype(numpy.float64) / scales.astype(numpy.float64))
|
|
raw_zp = numpy.clip(raw_zp, qmin_val, qmax_val)
|
|
zero_points = raw_zp.astype(zp_dtype)
|
|
zero_points[degenerate] = 0
|
|
|
|
# Move the block axis back to its original position so the returned arrays have
|
|
# the same rank as the weight and shape[axis] == n_blocks, as required by the
|
|
# ONNX opset-21 QuantizeLinear/DequantizeLinear spec.
|
|
scales = numpy.moveaxis(scales, 0, axis)
|
|
zero_points = numpy.moveaxis(zero_points, 0, axis)
|
|
|
|
return zero_points, scales
|
|
|
|
|
|
def compute_data_quant_params(
|
|
data: numpy.ndarray,
|
|
quant_type: onnx.TensorProto.DataType,
|
|
symmetric: bool,
|
|
reduce_range: bool = False,
|
|
min_real_range: float | None = None,
|
|
rmin_override: float | None = None,
|
|
rmax_override: float | None = None,
|
|
) -> tuple[numpy.ndarray, numpy.ndarray]:
|
|
"""
|
|
Returns the zero_point and scale for the given data.
|
|
|
|
:param data: The data for which to compute quantization parameters.
|
|
:param quant_type: The quantization data type.
|
|
:param symmetric: whether symmetric quantization is used or not.
|
|
:parameter reduce_range: True if the quantization range should be reduced. Defaults to False.
|
|
:parameter min_real_range: Minimum floating-point range (i.e., rmax - rmin) to enforce. Defaults to None.
|
|
:parameter rmin_override: The value of rmin to use if not None. Otherwise, uses min(data).
|
|
:parameter rmax_override: The value of rmax to use if not None. Otherwise, uses max(data).
|
|
:return: zero point and scale
|
|
"""
|
|
if not isinstance(data, numpy.ndarray):
|
|
raise TypeError(f"Weight must be given as an array not {type(data)}.")
|
|
if rmin_override is not None:
|
|
rmin = rmin_override
|
|
else:
|
|
rmin = data.min() if len(data) else 0.0
|
|
|
|
if rmax_override is not None:
|
|
rmax = rmax_override
|
|
else:
|
|
rmax = data.max() if len(data) else 0.0
|
|
|
|
rmin = numpy.array(rmin, dtype=data.dtype)
|
|
rmax = numpy.array(rmax, dtype=data.dtype)
|
|
scale = numpy.array(1.0, dtype=data.dtype)
|
|
|
|
if quant_type == TensorProto.FLOAT8E4M3FN:
|
|
if reduce_range:
|
|
raise RuntimeError("Unsupported option reduce_range=True for float 8.")
|
|
std = numpy.std(data)
|
|
zero_point, scale = compute_scale_zp_float8(quant_type, std)
|
|
return _check_type(zero_point, scale, zero_point_index=0)
|
|
|
|
if quant_type in (
|
|
TensorProto.INT8,
|
|
TensorProto.UINT8,
|
|
TensorProto.INT16,
|
|
TensorProto.UINT16,
|
|
TensorProto.INT4,
|
|
TensorProto.UINT4,
|
|
):
|
|
qmin, qmax = get_qmin_qmax_for_qType(quant_type, reduce_range, symmetric=symmetric)
|
|
if len(data):
|
|
zero_point, scale = compute_scale_zp(rmin, rmax, qmin, qmax, symmetric, min_real_range)
|
|
else:
|
|
zero_point = numpy.array(0, dtype=qmin.dtype)
|
|
return _check_type(zero_point, scale, zero_point_index=0)
|
|
|
|
raise ValueError(f"Unexpected value for quant_type={quant_type}.")
|
|
|
|
|
|
def quantize_data(
|
|
data, qType, symmetric, reduce_range=False, min_real_range=None, rmin_override=None, rmax_override=None
|
|
) -> tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
|
|
"""
|
|
:param data: data to quantize
|
|
:param qType: data type to quantize to.
|
|
:param symmetric: whether symmetric quantization is used or not.
|
|
:parameter reduce_range: True if the quantization range should be reduced. Defaults to False.
|
|
:parameter min_real_range: Minimum floating-point range (i.e., rmax - rmin) to enforce. Defaults to None.
|
|
:parameter rmin_override: The value of rmin to use if not None. Otherwise, uses min(data).
|
|
:parameter rmax_override: The value of rmax to use if not None. Otherwise, uses max(data).
|
|
:return: minimum, maximum, zero point, scale, and quantized weights
|
|
|
|
To pack weights, we compute a linear transformation
|
|
|
|
- when data `type == uint8` mode, from `[rmin, rmax]` -> :math:`[0, 2^{b-1}]` and
|
|
- when data `type == int8`, from `[-m , m]` -> :math:`[-(2^{b-1}-1), 2^{b-1}-1]` where
|
|
`m = max(abs(rmin), abs(rmax))`
|
|
|
|
and add necessary intermediate nodes to transform quantized weight to full weight using the equation
|
|
|
|
:math:`r = S(q-z)`, where
|
|
|
|
- *r*: real original value
|
|
- *q*: quantized value
|
|
- *S*: scale
|
|
- *z*: zero point
|
|
"""
|
|
zero_point, scale = compute_data_quant_params(
|
|
data,
|
|
qType,
|
|
symmetric,
|
|
reduce_range,
|
|
min_real_range,
|
|
rmin_override,
|
|
rmax_override,
|
|
)
|
|
if qType == TensorProto.FLOAT8E4M3FN:
|
|
quantized_data = quantize_nparray(qType, data, scale, zero_point)
|
|
if any((quantized_data.view(numpy.uint8).ravel() & 127) == 127):
|
|
np_data = numpy.asarray(data)
|
|
raise RuntimeError(
|
|
f"One of the quantized value is NaN data in [{np_data.min()}, {np_data.max()}], "
|
|
f"quantized_data in [{quantized_data.min()}, {quantized_data.max()}]."
|
|
)
|
|
return zero_point, scale, quantized_data
|
|
|
|
if qType in (
|
|
TensorProto.INT8,
|
|
TensorProto.UINT8,
|
|
TensorProto.INT16,
|
|
TensorProto.UINT16,
|
|
TensorProto.INT4,
|
|
TensorProto.UINT4,
|
|
):
|
|
quantized_data = quantize_nparray(qType, data, scale, zero_point)
|
|
return zero_point, scale, quantized_data
|
|
|
|
raise ValueError(f"Unexpected value for qType={qType}.")
|
|
|
|
|
|
def quantize_onnx_initializer(
|
|
weight: onnx.TensorProto,
|
|
quant_type: onnx.TensorProto.DataType,
|
|
zero_point: numpy.ndarray,
|
|
scale: numpy.ndarray,
|
|
axis: int | None = None,
|
|
quant_weight_name: str | None = None,
|
|
block_size: int = 0,
|
|
) -> onnx.TensorProto:
|
|
"""
|
|
Returns a quantized version of the given ONNX initializer.
|
|
|
|
:param weight: The ONNX initializer to quantize.
|
|
:param quant_type: The final quantized data type.
|
|
:param zero_point: The zero-point value to use for quantization.
|
|
:param scale: The scale value to use for quantization.
|
|
:param axis: The quantization axis if quantizing per-channel or per-block. Defaults to None.
|
|
:param quant_weight_name: The name of the quantized initializer.
|
|
If not specified, the quantized name is generated.
|
|
:param block_size: Block size for opset-21 block-wise quantization. 0 means disabled.
|
|
:return: The quantized ONNX initializer.
|
|
"""
|
|
weight_data = tensor_proto_to_array(weight)
|
|
q_weight_data: numpy.ndarray | None = None
|
|
|
|
if axis is not None and block_size > 0: # Per-block quantization
|
|
k = weight_data.shape[axis]
|
|
other = int(numpy.prod([d for i, d in enumerate(weight_data.shape) if i != axis]))
|
|
moved = numpy.moveaxis(weight_data, axis, 0).reshape(k, other)
|
|
# scale/zero_point are in spec shape (shape[axis] == n_blocks); move the block
|
|
# axis to position 0 locally so the loop can index [blk, col] uniformly.
|
|
scale_moved = numpy.moveaxis(scale, axis, 0)
|
|
zp_moved = numpy.moveaxis(zero_point, axis, 0)
|
|
n_blocks = scale_moved.shape[0]
|
|
quant_np_dtype = onnx.helper.tensor_dtype_to_np_dtype(quant_type)
|
|
q_moved = numpy.empty_like(moved, dtype=quant_np_dtype)
|
|
for blk in range(n_blocks):
|
|
start = blk * block_size
|
|
end = min(start + block_size, k)
|
|
for col in range(other):
|
|
q_moved[start:end, col] = quantize_nparray(
|
|
quant_type, moved[start:end, col].ravel(), scale_moved[blk, col], zp_moved[blk, col]
|
|
)
|
|
q_weight_data = numpy.moveaxis(
|
|
q_moved.reshape([k] + [d for i, d in enumerate(weight_data.shape) if i != axis]), 0, axis
|
|
)
|
|
elif axis is None: # Per-tensor quantization
|
|
q_weight_data = quantize_nparray(quant_type, weight_data.ravel(), scale, zero_point)
|
|
else: # Per-channel quantization
|
|
channel_count = weight_data.shape[axis]
|
|
channel_dims = list(weight_data.shape) # deep copy
|
|
channel_dims[axis] = 1 # only one per channel for reshape
|
|
quantized_channel_data_list = []
|
|
|
|
for i in range(channel_count):
|
|
channel_data = weight_data.take(i, axis)
|
|
channel_scale = scale[i]
|
|
channel_zero_point = zero_point[i]
|
|
quantized_channel_data = quantize_nparray(
|
|
quant_type, channel_data.ravel(), channel_scale, channel_zero_point
|
|
)
|
|
quantized_channel_data_list.append(numpy.asarray(quantized_channel_data).reshape(channel_dims))
|
|
|
|
q_weight_data = numpy.concatenate(quantized_channel_data_list, axis)
|
|
|
|
q_weight_name = quant_weight_name if quant_weight_name else f"{weight.name}{TENSOR_NAME_QUANT_SUFFIX}"
|
|
|
|
if quant_type == onnx.TensorProto.FLOAT8E4M3FN:
|
|
q_weight_initializer = onnx.TensorProto()
|
|
q_weight_initializer.data_type = quant_type
|
|
q_weight_initializer.dims.extend(weight.dims)
|
|
q_weight_initializer.name = q_weight_name
|
|
# Do not remove .flatten().copy() numpy is not clear about data persistence.
|
|
q_weight_initializer.raw_data = q_weight_data.flatten().copy().tobytes()
|
|
if to_array_extended is not None:
|
|
# This test should not be needed but it helped catch some issues
|
|
# with data persistence and tobytes.
|
|
check = to_array_extended(q_weight_initializer)
|
|
if check.shape != weight_data.shape or check.tobytes() != q_weight_data.tobytes():
|
|
raise RuntimeError(
|
|
f"The initializer of shape {weight_data.shape} could not be created, expecting "
|
|
f"{q_weight_data.tobytes()[:10]}, got {check.tobytes()[:10]} and shape={weight.shape}"
|
|
f"\nraw={str(q_weight_initializer)[:200]}."
|
|
)
|
|
elif quant_type in (onnx.TensorProto.INT4, onnx.TensorProto.UINT4):
|
|
if q_weight_data.dtype not in (int4, uint4):
|
|
raise RuntimeError(f"Quantized weights for {q_weight_name} must be 8-bit before packing as 4-bit values.")
|
|
|
|
# We do not use onnx.helper.pack_float32_to_4bit() due to performance.
|
|
# This can be the difference between a large model taking 30 minutes to quantize vs 5 minutes.
|
|
packed_data = bytes(pack_bytes_to_4bit(q_weight_data.tobytes()))
|
|
|
|
# We only use onnx.helper.make_tensor with raw data due to bug: https://github.com/onnx/onnx/pull/6161
|
|
q_weight_initializer = onnx.helper.make_tensor(q_weight_name, quant_type, weight.dims, packed_data, raw=True)
|
|
else:
|
|
quant_np_dtype = onnx.helper.tensor_dtype_to_np_dtype(quant_type)
|
|
q_weight_data = numpy.asarray(q_weight_data, dtype=quant_np_dtype).reshape(weight.dims)
|
|
q_weight_initializer = onnx.numpy_helper.from_array(q_weight_data, q_weight_name)
|
|
|
|
return q_weight_initializer
|
|
|
|
|
|
def get_qmin_qmax_for_qType(qType, reduce_range=False, symmetric=False): # noqa: N802
|
|
"""
|
|
Return qmin and qmax, the minimum and maximum value representable by the given qType
|
|
:parameter qType: onnx.onnx_pb.TensorProto.UINT8 or onnx.onnx_pb.TensorProto.UINT8
|
|
:return: qmin, qmax
|
|
"""
|
|
if qType == onnx_proto.TensorProto.FLOAT8E4M3FN:
|
|
raise NotImplementedError("This function is not implemented for float 8 as not needed.")
|
|
|
|
qrange = None
|
|
|
|
if reduce_range:
|
|
qrange = ONNX_INT_TYPE_REDUCED_RANGE.get(qType)
|
|
elif symmetric and qType in ONNX_INT_TYPE_SYMMETRIC_RANGE:
|
|
qrange = ONNX_INT_TYPE_SYMMETRIC_RANGE[qType]
|
|
else:
|
|
qrange = ONNX_INT_TYPE_RANGE.get(qType)
|
|
|
|
if not qrange:
|
|
raise ValueError(f"Unexpected data type {qType} requested. Only INT8, UINT8, INT16, and UINT16 are supported.")
|
|
|
|
qmin, qmax = qrange
|
|
if qmin > 0 or qmax < 0:
|
|
raise ValueError(
|
|
f"qmin and qmax must meet requirement: qmin <= 0 <= qmax while "
|
|
f"qmin:{qmin}, qmmax:{qmax}, dtype={qmin.dtype}, reduce_range={reduce_range}, "
|
|
f"symmetric={symmetric}, qType={qType}"
|
|
)
|
|
|
|
return qrange
|
|
|
|
|
|
def get_qrange_for_qType(qType, reduce_range=False, symmetric=False): # noqa: N802
|
|
"""
|
|
Helper function to get the quantization range for a type.
|
|
parameter qType: quantization type.
|
|
return: quantization range.
|
|
"""
|
|
qmin, qmax = get_qmin_qmax_for_qType(qType, reduce_range, symmetric=symmetric)
|
|
return qmax - qmin
|
|
|
|
|
|
def normalize_axis(axis: int, rank: int) -> tuple[bool, int]:
|
|
"""
|
|
Helper function that tries to return a normalized axis in the range [0, rank - 1].
|
|
:parameter axis: The axis to normalize.
|
|
:parameter rank: The tensor rank (number of dimensions).
|
|
:return (is_valid, axis_norm)
|
|
"""
|
|
axis_norm = axis + rank if axis < 0 else axis
|
|
is_valid = axis_norm >= 0 and axis_norm < rank
|
|
return is_valid, axis_norm
|
|
|
|
|
|
def pack_bytes_to_4bit(src_8bit: bytes) -> bytearray:
|
|
"""
|
|
Copies a source array of 8-bit values into a destination bytearray of packed 4-bit values.
|
|
Assumes that the source values are already in the appropriate int4 range.
|
|
:parameter src_8bit: The 8-bit element values to pack.
|
|
:return A bytearray with every two 8-bit src elements packed into a single byte.
|
|
"""
|
|
num_elems = len(src_8bit)
|
|
if num_elems == 0:
|
|
return bytearray()
|
|
|
|
dst_size = (num_elems + 1) // 2 # Ex: 5 8-bit elems packed into 3 bytes
|
|
dst = bytearray(dst_size)
|
|
|
|
src_i: int = 0
|
|
dst_i: int = 0
|
|
|
|
# Pack two 8-bit elements into a single byte in each iteration.
|
|
while src_i < num_elems - 1:
|
|
dst[dst_i] = ((src_8bit[src_i + 1] & 0xF) << 4) | (src_8bit[src_i] & 0xF)
|
|
dst_i += 1
|
|
src_i += 2
|
|
|
|
if src_i < num_elems:
|
|
# Odd number of elements.
|
|
dst[dst_i] = src_8bit[src_i] & 0xF
|
|
|
|
return dst
|
|
|
|
|
|
class QuantizedInitializer:
|
|
"""
|
|
Represents a linearly quantized weight input from ONNX operators
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name,
|
|
initializer,
|
|
rmins,
|
|
rmaxs,
|
|
zero_points,
|
|
scales,
|
|
data=[], # noqa: B006
|
|
quantized_data=[], # noqa: B006
|
|
axis=None,
|
|
):
|
|
self.name = name
|
|
self.initializer = initializer # TensorProto initializer in ONNX graph
|
|
self.rmins = rmins # List of minimum range for each axis
|
|
self.rmaxs = rmaxs # List of maximum range for each axis
|
|
# 1D tensor of zero points computed for each axis. scalar if axis is empty
|
|
self.zero_points = zero_points
|
|
self.scales = scales # 1D tensor of scales computed for each axis. scalar if axis is empty
|
|
self.data = data # original data from initializer TensorProto
|
|
self.quantized_data = quantized_data # weight-packed data from data
|
|
# Scalar to specify which dimension in the initializer to weight pack.
|
|
self.axis = axis
|
|
# If empty, single zero point and scales computed from a single rmin and rmax
|
|
|
|
|
|
class QuantizedValue:
|
|
"""
|
|
Represents a linearly quantized value (input\\output\\intializer)
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
name,
|
|
new_quantized_name,
|
|
scale_name,
|
|
zero_point_name,
|
|
quantized_value_type,
|
|
axis=None,
|
|
node_type=None,
|
|
node_qtype=None,
|
|
scale_type=None,
|
|
):
|
|
self.original_name = name
|
|
self.q_name = new_quantized_name
|
|
self.scale_name = scale_name
|
|
self.zp_name = zero_point_name
|
|
self.value_type = quantized_value_type
|
|
self.axis = axis
|
|
self.node_type = node_type
|
|
self.node_qtype = node_qtype
|
|
self.scale_type = scale_type
|
|
|
|
|
|
class BiasToQuantize:
|
|
"""
|
|
Represents a bias to be quantized
|
|
"""
|
|
|
|
def __init__(self, bias_name, input_name, weight_name):
|
|
self.bias_name = bias_name
|
|
self.input_name = input_name
|
|
self.weight_name = weight_name
|
|
|
|
|
|
def attribute_to_kwarg(attribute):
|
|
"""
|
|
Convert attribute to kwarg format for use with onnx.helper.make_node.
|
|
:parameter attribute: attribute in AttributeProto format.
|
|
:return: attribute in {key: value} format.
|
|
"""
|
|
if attribute.type == 0:
|
|
raise ValueError(f"attribute {attribute.name} does not have type specified.")
|
|
|
|
# Based on attribute type definitions from AttributeProto
|
|
# definition in https://github.com/onnx/onnx/blob/main/onnx/onnx.proto
|
|
if attribute.type == 1:
|
|
value = attribute.f
|
|
elif attribute.type == 2:
|
|
value = attribute.i
|
|
elif attribute.type == 3:
|
|
value = attribute.s
|
|
elif attribute.type == 4:
|
|
value = attribute.t
|
|
elif attribute.type == 5:
|
|
value = attribute.g
|
|
elif attribute.type == 6:
|
|
value = attribute.floats
|
|
elif attribute.type == 7:
|
|
value = attribute.ints
|
|
elif attribute.type == 8:
|
|
value = attribute.strings
|
|
elif attribute.type == 9:
|
|
value = attribute.tensors
|
|
elif attribute.type == 10:
|
|
value = attribute.graphs
|
|
else:
|
|
raise ValueError(f"attribute {attribute.name} has unsupported type {attribute.type}.")
|
|
|
|
return {attribute.name: value}
|
|
|
|
|
|
def find_by_name(item_name, item_list):
|
|
"""
|
|
Helper function to find item by name in a list.
|
|
parameter item_name: name of the item.
|
|
parameter item_list: list of items.
|
|
return: item if found. None otherwise.
|
|
"""
|
|
items = [item for item in item_list if item.name == item_name]
|
|
return items[0] if len(items) > 0 else None
|
|
|
|
|
|
def get_elem_index(elem_name, elem_list):
|
|
"""
|
|
Helper function to return index of an item in a node list
|
|
"""
|
|
elem_idx = -1
|
|
for i in range(len(elem_list)):
|
|
if elem_list[i] == elem_name:
|
|
elem_idx = i
|
|
return elem_idx
|
|
|
|
|
|
def get_mul_node(inputs, output, name):
|
|
"""
|
|
Helper function to create a Mul node.
|
|
parameter inputs: list of input names.
|
|
parameter output: output name.
|
|
parameter name: name of the node.
|
|
return: Mul node in NodeProto format.
|
|
"""
|
|
return onnx.helper.make_node("Mul", inputs, [output], name)
|
|
|
|
|
|
def generate_identified_filename(filename: Path, identifier: str) -> Path:
|
|
"""
|
|
Helper function to generate a identifiable filepath by concatenating the given identifier as a suffix.
|
|
"""
|
|
return filename.parent.joinpath(filename.stem + identifier + filename.suffix)
|
|
|
|
|
|
def apply_plot(hist, hist_edges):
|
|
import sys # noqa: PLC0415
|
|
|
|
import matplotlib.pyplot as plt # noqa: PLC0415
|
|
import numpy # noqa: PLC0415
|
|
|
|
numpy.set_printoptions(threshold=sys.maxsize)
|
|
print("Histogram:")
|
|
print(hist)
|
|
print("Histogram Edges:")
|
|
print(hist_edges)
|
|
plt.stairs(hist, hist_edges, fill=True)
|
|
plt.xlabel("Tensor value")
|
|
plt.ylabel("Counts")
|
|
plt.title("Tensor value V.S. Counts")
|
|
plt.show()
|
|
|
|
|
|
def write_calibration_table(calibration_cache, dir="."):
|
|
"""
|
|
Helper function to write calibration table to files.
|
|
"""
|
|
|
|
import json # noqa: PLC0415
|
|
|
|
import flatbuffers # noqa: PLC0415
|
|
import numpy as np # noqa: PLC0415
|
|
|
|
import onnxruntime.quantization.CalTableFlatBuffers.KeyValue as KeyValue # noqa: PLC0415
|
|
import onnxruntime.quantization.CalTableFlatBuffers.TrtTable as TrtTable # noqa: PLC0415
|
|
|
|
# Use the shared encoder from calibrate.py so write_calibration_table and
|
|
# save_tensors_data produce identical JSON for numpy scalar/array values.
|
|
from onnxruntime.quantization.calibrate import CalibrationCacheEncoder # noqa: PLC0415
|
|
|
|
logging.info(f"calibration cache: {calibration_cache}")
|
|
|
|
json_data = json.dumps(calibration_cache, cls=CalibrationCacheEncoder)
|
|
|
|
with open(os.path.join(dir, "calibration.json"), "w") as file:
|
|
file.write(json_data) # use `json.loads` to do the reverse
|
|
|
|
# Serialize data using FlatBuffers
|
|
zero = np.array(0)
|
|
builder = flatbuffers.Builder(1024)
|
|
key_value_list = []
|
|
for key in sorted(calibration_cache.keys()):
|
|
values = calibration_cache[key]
|
|
d_values = values.to_dict()
|
|
floats = [
|
|
float(d_values.get("highest", zero).item()),
|
|
float(d_values.get("lowest", zero).item()),
|
|
]
|
|
value = str(max(floats))
|
|
|
|
flat_key = builder.CreateString(key)
|
|
flat_value = builder.CreateString(value)
|
|
|
|
KeyValue.KeyValueStart(builder)
|
|
KeyValue.KeyValueAddKey(builder, flat_key)
|
|
KeyValue.KeyValueAddValue(builder, flat_value)
|
|
key_value = KeyValue.KeyValueEnd(builder)
|
|
|
|
key_value_list.append(key_value)
|
|
|
|
TrtTable.TrtTableStartDictVector(builder, len(key_value_list))
|
|
for key_value in key_value_list:
|
|
builder.PrependUOffsetTRelative(key_value)
|
|
main_dict = builder.EndVector()
|
|
|
|
TrtTable.TrtTableStart(builder)
|
|
TrtTable.TrtTableAddDict(builder, main_dict)
|
|
cal_table = TrtTable.TrtTableEnd(builder)
|
|
|
|
builder.Finish(cal_table)
|
|
buf = builder.Output()
|
|
|
|
with open(os.path.join(dir, "calibration.flatbuffers"), "wb") as file:
|
|
file.write(buf)
|
|
|
|
# Deserialize data (for validation)
|
|
if os.environ.get("QUANTIZATION_DEBUG", "0") in (1, "1"):
|
|
cal_table = TrtTable.TrtTable.GetRootAsTrtTable(buf, 0)
|
|
dict_len = cal_table.DictLength()
|
|
for i in range(dict_len):
|
|
key_value = cal_table.Dict(i)
|
|
logging.info(key_value.Key())
|
|
logging.info(key_value.Value())
|
|
|
|
# write plain text
|
|
with open(os.path.join(dir, "calibration.cache"), "w") as file:
|
|
for key in sorted(calibration_cache.keys()):
|
|
values = calibration_cache[key]
|
|
d_values = values.to_dict()
|
|
floats = [
|
|
float(d_values.get("highest", zero).item()),
|
|
float(d_values.get("lowest", zero).item()),
|
|
]
|
|
value = key + " " + str(max(floats))
|
|
file.write(value)
|
|
file.write("\n")
|
|
|
|
|
|
def smooth_distribution(p, eps=0.0001):
|
|
"""Given a discrete distribution (may have not been normalized to 1),
|
|
smooth it by replacing zeros with eps multiplied by a scaling factor
|
|
and taking the corresponding amount off the non-zero values.
|
|
Ref: http://web.engr.illinois.edu/~hanj/cs412/bk3/KL-divergence.pdf
|
|
https://github.com//apache/incubator-mxnet/blob/master/python/mxnet/contrib/quantization.py
|
|
"""
|
|
is_zeros = (p == 0).astype(numpy.float32)
|
|
is_nonzeros = (p != 0).astype(numpy.float32)
|
|
n_zeros = is_zeros.sum()
|
|
n_nonzeros = p.size - n_zeros
|
|
|
|
if not n_nonzeros:
|
|
# raise ValueError('The discrete probability distribution is malformed. All entries are 0.')
|
|
return None
|
|
eps1 = eps * float(n_zeros) / float(n_nonzeros)
|
|
assert eps1 < 1.0, f"n_zeros={n_zeros}, n_nonzeros={n_nonzeros}, eps1={eps1}"
|
|
|
|
hist = p.astype(numpy.float32)
|
|
hist += eps * is_zeros + (-eps1) * is_nonzeros
|
|
assert (hist <= 0).sum() == 0
|
|
|
|
return hist
|
|
|
|
|
|
def model_has_external_data(model_path: Path):
|
|
model = onnx.load(model_path.as_posix(), load_external_data=False)
|
|
return any(external_data_helper.uses_external_data(intializer) for intializer in model.graph.initializer)
|
|
|
|
|
|
def optimize_model(model_path: Path, opt_model_path: Path):
|
|
"""
|
|
Generate model that applies graph optimization (constant folding, etc.)
|
|
parameter model_path: path to the original onnx model
|
|
parameter opt_model_path: path to the optimized onnx model
|
|
:return: optimized onnx model
|
|
"""
|
|
sess_option = SessionOptions()
|
|
sess_option.optimized_model_filepath = opt_model_path.as_posix()
|
|
sess_option.graph_optimization_level = GraphOptimizationLevel.ORT_ENABLE_BASIC
|
|
kwargs = {}
|
|
# This will rename constant initializer names, disable it to make test pass.
|
|
kwargs["disabled_optimizers"] = ["ConstantSharing"]
|
|
_ = InferenceSession(model_path.as_posix(), sess_option, providers=["CPUExecutionProvider"], **kwargs)
|
|
|
|
|
|
def add_pre_process_metadata(model: ModelProto):
|
|
"""Tag the model that it went through quantization pre-processing"""
|
|
metadata_props = {"onnx.quant.pre_process": "onnxruntime.quant"}
|
|
if model.metadata_props:
|
|
for prop in model.metadata_props:
|
|
metadata_props.update({prop.key: prop.value})
|
|
onnx.helper.set_model_props(model, metadata_props)
|
|
|
|
|
|
def model_has_pre_process_metadata(model: ModelProto) -> bool:
|
|
"""Check the model whether it went through quantization pre-processing"""
|
|
if model.metadata_props:
|
|
for prop in model.metadata_props:
|
|
if prop.key == "onnx.quant.pre_process" and prop.value == "onnxruntime.quant":
|
|
return True
|
|
return False
|
|
|
|
|
|
def add_infer_metadata(model: ModelProto):
|
|
metadata_props = {"onnx.infer": "onnxruntime.quant"}
|
|
if model.metadata_props:
|
|
for p in model.metadata_props:
|
|
metadata_props.update({p.key: p.value})
|
|
onnx.helper.set_model_props(model, metadata_props)
|
|
|
|
|
|
def model_has_infer_metadata(model: ModelProto) -> bool:
|
|
if model.metadata_props:
|
|
for p in model.metadata_props:
|
|
if p.key == "onnx.infer" and p.value == "onnxruntime.quant":
|
|
return True
|
|
return False
|
|
|
|
|
|
def get_opset_version(model: ModelProto) -> int:
|
|
ai_onnx_domain = [opset for opset in model.opset_import if not opset.domain or opset.domain == "ai.onnx"]
|
|
if len(ai_onnx_domain) != 1:
|
|
raise ValueError("Failed to find proper ai.onnx domain")
|
|
opset_version = ai_onnx_domain[0].version
|
|
|
|
return opset_version
|
|
|
|
|
|
def update_opset_version(
|
|
model: ModelProto,
|
|
weight_type: QuantType,
|
|
activation_type: QuantType | None = None,
|
|
tensor_quant_overrides: dict | None = None,
|
|
block_size: int = 0,
|
|
) -> ModelProto:
|
|
opset_version = get_opset_version(model)
|
|
target_opset_version = opset_version
|
|
weight_quant_type = getattr(weight_type, "tensor_type", weight_type)
|
|
activation_quant_type = (
|
|
getattr(activation_type, "tensor_type", activation_type) if activation_type is not None else None
|
|
)
|
|
|
|
_int16_types = (onnx.TensorProto.UINT16, onnx.TensorProto.INT16)
|
|
needs_opset21_for_16bit = weight_quant_type in _int16_types or activation_quant_type in _int16_types
|
|
|
|
# Also check TensorQuantOverrides for any 16-bit types, including per-override convert.quant_type.
|
|
# Validation of structure is deferred to TensorQuantOverridesHelper.is_valid(); skip bump heuristic on malformed input.
|
|
if not needs_opset21_for_16bit and tensor_quant_overrides:
|
|
_int16_quant_types = {QuantType.QInt16, QuantType.QUInt16}
|
|
try:
|
|
for overrides_list in tensor_quant_overrides.values():
|
|
for override in overrides_list:
|
|
qt = override.get("quant_type")
|
|
if qt in _int16_quant_types:
|
|
needs_opset21_for_16bit = True
|
|
break
|
|
convert = override.get("convert")
|
|
if convert is not None:
|
|
convert_qt = convert.get("quant_type")
|
|
if convert_qt in _int16_quant_types:
|
|
needs_opset21_for_16bit = True
|
|
break
|
|
if needs_opset21_for_16bit:
|
|
break
|
|
except (AttributeError, TypeError):
|
|
# Malformed overrides; structural validation is deferred to
|
|
# TensorQuantOverridesHelper.is_valid(). Skip bump heuristic.
|
|
logging.debug("Skipping 16-bit opset bump heuristic for TensorQuantOverrides: structure not as expected.")
|
|
|
|
if opset_version < 21 and block_size > 0:
|
|
logging.warning(
|
|
f"The original model opset version is {opset_version}, which does not support block-wise "
|
|
"quantization natively. "
|
|
"Please update the model to opset >= 21. Automatically updating the model to opset 21. "
|
|
"Please verify the quantized model."
|
|
)
|
|
target_opset_version = 21
|
|
|
|
elif opset_version < 19 and weight_quant_type == onnx.TensorProto.FLOAT8E4M3FN:
|
|
logging.warning(
|
|
f"The original model opset version is {opset_version}, which does not support quantization to float 8. "
|
|
"Please update the model to opset >= 19. Automatically update the model to opset 19. "
|
|
"Please verify the quantized model."
|
|
)
|
|
target_opset_version = 19
|
|
|
|
elif opset_version < 21 and needs_opset21_for_16bit:
|
|
logging.warning(
|
|
f"The original model opset version is {opset_version}, which does not support 16-bit integer "
|
|
"quantization natively. "
|
|
"Please update the model to opset >= 21. Automatically update the model to opset 21. "
|
|
"Please verify the quantized model."
|
|
)
|
|
target_opset_version = 21
|
|
|
|
elif opset_version == 10:
|
|
logging.warning(
|
|
f"The original model opset version is {opset_version}, which does not support node fusions. "
|
|
"Please update the model to opset >= 11 for better performance."
|
|
)
|
|
|
|
elif opset_version < 10:
|
|
logging.warning(
|
|
f"The original model opset version is {opset_version}, which does not support quantization. "
|
|
"Please update the model to opset >= 11. Automatically update the model to opset 11. "
|
|
"Please verify the quantized model."
|
|
)
|
|
target_opset_version = 11
|
|
|
|
if target_opset_version != opset_version:
|
|
model = onnx.version_converter.convert_version(model, target_opset_version)
|
|
# Additional nodes may be added to the model during the opset version conversion. Run shape inference
|
|
# to ensure all nodes are included in model.graph.value_info.
|
|
model = save_and_reload_model_with_shape_infer(model)
|
|
|
|
return model
|
|
|
|
|
|
def load_model_with_shape_infer(model_path: Path) -> ModelProto:
|
|
inferred_model_path = generate_identified_filename(model_path, "-inferred")
|
|
onnx.shape_inference.infer_shapes_path(str(model_path), str(inferred_model_path))
|
|
model = onnx.load(inferred_model_path.as_posix())
|
|
add_infer_metadata(model)
|
|
inferred_model_path.unlink()
|
|
return model
|
|
|
|
|
|
def save_and_reload_model_with_shape_infer(model: ModelProto) -> ModelProto:
|
|
with tempfile.TemporaryDirectory(prefix="ort.quant.") as quant_tmp_dir:
|
|
model_copy = copy.deepcopy(model)
|
|
model_path = Path(quant_tmp_dir).joinpath("model.onnx")
|
|
onnx.save_model(model_copy, model_path.as_posix(), save_as_external_data=True)
|
|
return load_model_with_shape_infer(model_path)
|
|
|
|
|
|
def tensor_proto_to_array(initializer: TensorProto) -> numpy.ndarray:
|
|
if initializer.data_type in (onnx_proto.TensorProto.FLOAT, onnx_proto.TensorProto.FLOAT16):
|
|
return onnx.numpy_helper.to_array(initializer)
|
|
|
|
raise ValueError(
|
|
f"Only float type is supported. Weights {initializer.name} is {type_to_name[initializer.data_type]}"
|
|
)
|
|
|
|
|
|
def add_quant_suffix(tensor_name: str) -> str:
|
|
return tensor_name + "_QuantizeLinear"
|
|
|
|
|
|
def add_quant_input_suffix(tensor_name: str) -> str:
|
|
return tensor_name + QUANT_INPUT_SUFFIX
|
|
|
|
|
|
def add_quant_output_suffix(tensor_name) -> str:
|
|
return tensor_name + "_QuantizeLinear_Output"
|
|
|
|
|
|
def add_dequant_suffix(tensor_name) -> str:
|
|
return tensor_name + "_DequantizeLinear"
|
|
|
|
|
|
def add_dequant_input_suffix(tensor_name) -> str:
|
|
return tensor_name + "_DequantizeLinear_Input"
|
|
|
|
|
|
def add_dequant_output_suffix(tensor_name) -> str:
|
|
return tensor_name + DEQUANT_OUTPUT_SUFFIX
|