aimet_onnx.quantsim

Note

It is recommended to use onnx-simplifier before creating quantsim model.

class aimet_onnx.QuantizationSimModel(model, *, param_type=int8, activation_type=int8, quant_scheme=QuantScheme.min_max, config_file=None, dummy_input=None, user_onnx_libs=None, providers=None, path=None, session_options_kwargs=None)[source]

Class that simulates the quantized model execution on a target hardware backend.

Parameters:

  • model (ModelProto) – ONNX ModelProto to quantize

  • param_type (Union[str, qtype]) – quantized type to use for parameter tensors. Can be { int2, int4, int8, int16, float16, float32, float8e4m3fn, float8e5m2 } or aimet_onnx.qtype

  • activation_type (Union[str, qtype]) – quantized type to use for activation tensors. Can be { int2, int4, int8, int16, float16, float32, float8e4m3fn, float8e5m2 } or aimet_onnx.qtype

  • quant_scheme (Union[str, QuantScheme]) – Quantization scheme to use for calibration. Can be { min_max, tf_enhanced } or QuantScheme

  • config_file (Optional[str]) – File path or alias of the configuration file. Alias can be one of { default, htp_v66, htp_v68, htp_v69, htp_v73, htp_v75, htp_v79, htp_v81, htp_v85, enpu_v6, enpu_v7 } (Default: “default”)

  • dummy_input (Optional[Dict[str, ndarray]]) – Sample input to the model. Only needed for non shape-inferable models with parameterized shapes

  • user_onnx_libs (Optional[List[str]]) – List of paths to all compiled ONNX custom ops libraries

  • providers (Optional[Sequence[Union[str, Tuple[str, Dict[Any, Any]]]]]) – Onnxruntime execution providers to use when building InferenceSession. If None, default provider is “CPUExecutionProvider”

  • path (Optional[str]) – Directory to save temporary artifacts.

  • session_options_kwargs (Optional[Dict[str, Any]]) – Attributes to set on the sim’s onnxruntime.SessionOptions (e.g. {"intra_op_num_threads": 2}). An unknown attribute name raises AttributeError.

compute_encodings(*args, **kwargs)[source]

Computes encodings for all quantizers in the model.

Overloads:

  • inputs (Iterable[Dict[str, ndarray]])

  • forward_pass_callback (Callable[[InferenceSession], Any])

  • forward_pass_callback (Callable[[ort.InferenceSession, T], Any]), forward_pass_callback_args (T)

This API will invoke forward_pass_callback, a function written by the user that runs forward pass(es) of the quantized model with a small, representative subset of the training dataset. By doing so, the quantizers in the quantized model will observe the inputs and initialize their quantization encodings according to the observed input statistics.

This function is overloaded with the following signatures:

compute_encodings(inputs)[source]
Parameters:

inputs (Iterable[Dict[str, np.ndarray]]) – The set of model input samples to use during calibration

compute_encodings(forward_pass_callback)[source]
Parameters:

forward_pass_callback (Callable[[ort.InferenceSession], Any]) – A function that takes a quantized model and runs forward passes with a small, representative subset of training dataset

compute_encodings(forward_pass_callback, forward_pass_callback_args)[source]
Parameters:

  • forward_pass_callback (Callable[[ort.InferenceSession, T], Any]) – A function that takes a quantized model and runs forward passes with a small, representative subset of training dataset

  • forward_pass_callback_args (T) – The second argument to forward_pass_callback.

Example

>>> sim = QuantizationSimModel(...)
>>> def run_forward_pass(session: ort.InferenceSession):
...     for input in dataset:
...         _ = sess.run(None, {"input": input})
...
>>> sim.compute_encodings(run_forward_pass)
export(path, filename_prefix, export_model=True, *, export_int32_bias=None, encoding_version=None, force_activation_as='unsigned')[source]

Compute encodings and export to files

Parameters:

  • path (str) – dir to save encoding files

  • filename_prefix (str) – filename to save encoding files

  • export_model (bool) – If True, then ONNX model is exported. When False, only encodings are exported.

  • export_int32_bias (Optional[bool]) – If true, generate and export int32 bias encoding on the fly. Default: True if encoding version is 2.0.0 or higher, otherwise False.

  • encoding_version (Optional[str]) – Version of the encoding format to use. (default: 1.0.0) Supported versions are: [‘0.6.1’, ‘1.0.0’, ‘2.0.0’, ‘2.1.0’]

  • force_activation_as (Union[Literal['unsigned'], Literal['signed'], None]) – Force representing quantized activations as signed or unsigned integers. This argument is only applicable for encoding version 2.0.0. For versions 0.6.1 and 1.0.0, this argument is ignored. (default: “unsigned”)

Example

>>> sim.export(path=".", filename_prefix="model", encoding_version="2.0.0")
classmethod from_onnx_qdq(model, strict=False, **kwargs)[source]

Create sim from an ONNX QDQ model that contains QuantizeLinear/DequantizeLinear nodes.

This method is designed to construct a fully quantized model based on a partially quantized ONNX QDQ model exported from aimet-torch (aimet_torch.onnx.export()) or other 3rd party tools.

Parameters:

  • model (ModelProto) – ONNX model that contains QuantizeLinear/DequantizeLinear

  • strict (bool) – If True, raises error if there are QuantizeLinear/DequantizeLinear nodes in the model that cannot be mapped to any quantizer in the sim. If False, ignores incompatible QuantizeLinear/DequantizeLinear nodes and only loads the rest of the encodings. (Default: False)

  • **kwargs – same as QuantizationSimModel.__init__

Returns:

QuantizationSimModel created from ONNX QDQ model

Return type:

QuantizationSimModel

Example

>>> sim = aimet_onnx.QuantizationSimModel.from_onnx_qdq(
...     onnx.load("model_qdq.onnx"),
...     config_file="htp_v81",
... )
Quant - INFO - Loaded 26 out of 63 encodings from QuantizeLinear/DequantizeLinear nodes
set_tensor_precision(names, precision, *, strict=True)[source]

Set the quantization precision for each tensor in names to the specified precision

If quantization is not performed directly on the specified tensor (e.g., for outputs of data movement ops), propagates upwards to the closest enabled quantizer.

Example

sim = QuantizationSimModel(…) sim.set_tensor_precision(“/Conv_output_0”, aimet_onnx.int16)

Parameters:

  • names (str | list[str]) – List of tensor names to set the precision for

  • precision (qtype | str) – Precision to quantize to. If string, must be a valid alias of a qtype

  • strict (bool) – If True, throws an error if the tensor does not exist or is not quantized

to_onnx_qdq(*, export_int32_bias=False, prequantize_constants=False, force_activation_as='unsigned')[source]

Return a copy of ModelProto with all QcQuantizeOp nodes replaced with QuantizeLinear and/or DequantizeLinear.

Example

>>> len([qc_op for qc_op in sim.model.nodes() if dq.op_type == "QcQuantizeOp"])
10
>>> onnx_qdq = sim.to_onnx_qdq()
>>> len([qc_op for qc_op in sim.model.nodes() if dq.op_type == "QcQuantizeOp"])
0
>>> len([dq for dq in onnx_qdq.graph.node if dq.op_type == "DequantizeLinear"])
10
Parameters:

  • export_int32_bias (bool) – If true, generate and export int32 bias encoding on the fly (default: True)

  • prequantize_constants (bool) – If True, weights will be represented as quantized weight followed by DequantizeLinear nodes. If False, weights will be represented as float tensors followed by QuantizeLinear and DequantizeLinear nodes.

  • force_activation_as (Union[Literal['unsigned'], Literal['signed'], None]) – Force representing quantized activations as signed or unsigned integers (default: “unsigned”)

Return type:

ModelProto

Note

FP8 (float8e4m3fn/float8e5m2) quantizers are exported as onnx::QuantizeLinear/DequantizeLinear with a float8 y_zero_point of 0, which requires opset >= 19. The opset is raised automatically when needed.

Such a model must be run with an onnxruntime graph optimization level of at most ORT_ENABLE_BASIC. From ORT_ENABLE_EXTENDED upward, onnxruntime fuses QDQ pairs into integer kernels such as QLinearConv, which do not accept float8 inputs and cause session creation to fail with “Type ‘tensor(float8e4m3fn)’ … is invalid”.

../../_images/conv_qdq.onnx.svg
aimet_onnx.compute_encodings(sim)[source]

Computes encodings for all quantizers in the model.

Under this context manager, QuantizationSimModel will observe all inputs that run through the model to calibrate the quantization encoding of each quantizer.

Example

>>> sim = QuantizationSimModel(...)
>>> with compute_encodings(sim):
...     for input in dataset:
...         _ = sim.session.run(None, {"input": input})
aimet_onnx.encodings_to_onnx_qdq(model, encodings, *, prequantize_constants=False, force_activation_as='unsigned')[source]

Convert an ONNX model and its AIMET encodings into an ONNX QDQ model.

This is a direct translation: the encodings determine exactly which tensors are quantized and how. No QuantizationSimModel is created, so no quantsim configuration (supergroups, op tying, encoding constraints) can add or remove quantizers relative to what the encodings specify.

[b]float16 is not a quantized data type and onnx QDQ has no representation for it, so plain-cast float encodings are omitted with a warning, whether they name a parameter or an activation. Scaled float formats such as FP8 do carry a scale, and are emitted as Q/DQ.

An encoding naming a tensor that is not in the model is an error.

Parameters:

  • model (ModelProto | ONNXModel | str) – ONNX model as a ModelProto or ONNXModel, or a path to one. The model is not modified in place.

  • encodings (dict | str) – AIMET encodings as a dict, or a path to an encodings file. Versions 0.6.1, 1.0.0, 2.0.0, and 2.1.0 are accepted.

  • prequantize_constants (bool) – If True, store quantized parameters as integer initializers rather than emitting QuantizeLinear for them. (default: False)

  • force_activation_as (Optional[Literal['unsigned', 'signed']]) – Cast all activation encodings to signed or unsigned, a workaround for QNN converter limitations. Pass None to leave them as-is. (default: “unsigned”)

Return type:

ModelProto

Returns:

The equivalent ONNX QDQ model.

Examples

>>> qdq_model = encodings_to_onnx_qdq("model.onnx", "model.encodings")
>>> onnx.save_model(qdq_model, "model_qdq.onnx")
aimet_onnx.quantsim.set_param_type(sim, param_type, **kwargs)[source]

Set parameter quantization data type for specified layers.

Overloads:

  • sim (QuantizationSimModel), param_type (QSpec | qtype | str), *, op_types (Union[Tuple[str], str, None]) = None, nodes_to_exclude (Optional[Set[str]]) = None, strict (bool) = True, shift_zero_point (bool) = False

  • sim (QuantizationSimModel), param_type (QSpec | qtype | str), *, nodes_to_include (Optional[Set[str]]) = None, strict (bool) = True, shift_zero_point (bool) = False

This function is overloaded with the following signatures:

aimet_onnx.quantsim.set_param_type(sim, param_type, *, nodes_to_include=None, strict=True, shift_zero_point=False)[source]
Parameters:

  • sim (QuantizationSimModel) – Quantsim to set param type for

  • param_type (QSpec | qtype | str) – Quantization data type to set for the parameters

  • nodes_to_include (Set[str]) – Set of onnx node names for which to set parameter quantization data type. If None, all nodes are included

  • strict (bool) – If False, skip quantizers incompatible with param_type.

  • shift_zero_point (bool) – (Deprecated) Whether to shift the quantizer’s zero point (only for int2 param type).

aimet_onnx.quantsim.set_param_type(sim, param_type, *, nodes_to_include=None, strict=True, shift_zero_point=False)[source]
Parameters:

  • sim (QuantizationSimModel) – Quantsim to set param type for

  • param_type (QSpec | qtype | str) – Quantization data type to set for the parameters

  • op_types (Set[str]) – Set of onnx op types for which to set parameter quantization data type. If None, all types are included

  • nodes_to_exclude (Set[str]) – Set of onnx node names to exclude for setting parameter quantization data type

  • strict (bool) – If False, skip quantizers incompatible with param_type.

  • shift_zero_point (bool) – (Deprecated) Whether to shift the quantizer’s zero point (only for int2 param type).

Examples

>>> sim = QuantizationSimModel(...)

>>> # Set all parameter quantizers to int8 data type
>>> set_param_type(sim, aimet_onnx.int8)

>>> # Set parameter quantizers of Conv, MatMul, and Gemm layers to int4 data type
>>> set_param_type(sim, aimet_onnx.int4, op_types={"Conv", "MatMul", "Gemm"})

>>> # Set parameter quantizers of "/lm_head/MatMul" to int2 with shifted zero point
>>> spec = aimet_onnx.QSpec.blockwise(int2, block_size=64, symmetric=True, shift_zero_point=True)
>>> set_param_type(sim, spec, nodes_to_include={"/lm_head/MatMul"})

>>> # Set parameter quantizers of "/lm_head/MatMul" to LPBQ
>>> spec = aimet_onnx.QSpec.lpbq(int4, block_size=64)
>>> set_param_type(sim, spec, nodes_to_include={"/lm_head/MatMul"})

Quantization Spec

class aimet_onnx.QSpec(dtype, granularity=None, symmetric=None, shift_zero_point=False)[source]

Specifies how a tensor should be quantized (precision and granularity).

Parameters:

  • dtype (qtype) – Quantized data type

  • granularity (Optional[Granularity]) – How quantization parameters are shared across the tensor

  • symmetric (Optional[bool]) – If specified, determines whether encodings will be symmetric

  • shift_zero_point (bool) – Whether to shift quantizer’s 0 point by a half-step. Only supported for int2 param quantizers

classmethod blockwise(dtype, block_size, *, symmetric=None, shift_zero_point=False)[source]

Constructs a blockwise QSpec

Return type:

QSpec

classmethod lpbq(dtype, block_size=64, scale_bits=4)[source]

Constructs a low-power blockwise QSpec

Return type:

QSpec

classmethod per_channel(dtype, *, symmetric=None, shift_zero_point=False)[source]

Constructs a per-channel QSpec

Return type:

QSpec

classmethod per_tensor(dtype, *, symmetric=None, shift_zero_point=False)[source]

Constructs a per-tensor QSpec

Return type:

QSpec

Quant Scheme Enum

class aimet_onnx.common.defs.QuantScheme(*values)[source]

Quantization schemes

classmethod from_str(alias)[source]

Returns QuantScheme object from string alias

Return type:

QuantScheme