量化 ONNX 模型

目录

量化概述

ONNX 运行时中的量化是指对 ONNX 模型进行 8 位线性量化。

在量化过程中,浮点值被映射到 8 位量化空间,其形式为:val_fp32 = scale * (val_quantized - zero_point)

scale 是一个正实数,用于将浮点数映射到量化空间。其计算方法如下:

对于非对称量化

 scale = (data_range_max - data_range_min) / (quantization_range_max - quantization_range_min)

对于对称量化

 scale = max(abs(data_range_max), abs(data_range_min)) * 2 / (quantization_range_max - quantization_range_min)

zero_point 表示量化空间中的零点。重要的是,浮点零值必须在量化空间中精确可表示。这是因为许多 CNN 中使用了零填充。如果在量化后无法唯一地表示 0,将会导致精度误差。

ONNX 量化表示格式

有两种方式表示量化后的 ONNX 模型

  • 算子导向 (QOperator)
    所有量化后的算子都有自己的 ONNX 定义,例如 QLinearConv、MatMulInteger 等。
  • 张量导向 (QDQ;量化和反量化)
    这种格式在原始算子之间插入 DeQuantizeLinear(QuantizeLinear(张量)),以模拟量化和反量化过程。
    在静态量化中,QuantizeLinear 和 DeQuantizeLinear 算子也携带量化参数。
    在动态量化中,插入 ComputeQuantizationParameters 函数原型,以便实时计算量化参数。
  • 通过以下方式生成的模型采用 QDQ 格式
    1. 使用下文解释的 quantize_static 量化的模型,参数为 quant_format=QuantFormat.QDQ
    2. 从 Tensorflow 转换或从 PyTorch 导出的量化感知训练 (QAT) 模型。
    3. 从 TFLite 和其他框架转换的量化模型。

对于后两种情况,您无需使用量化工具对模型进行量化。ONNX 运行时可以直接将它们作为量化模型运行。

下图展示了 QOperator 和 QDQ 格式表示量化后的 Conv 的等效方式。这个端到端示例演示了这两种格式。

Changes to nodes from basic and extended optimizations

量化 ONNX 模型

ONNX 运行时提供了 Python API,用于将 32 位浮点模型转换为 8 位整数模型,即量化。这些 API 包括预处理、动态/静态量化和调试。

预处理

预处理是对 float32 模型进行转换,以便为量化做准备。它包含以下三个可选步骤

  1. 符号形状推理。这最适合用于 Transformer 模型。
  2. 模型优化:此步骤使用 ONNX 运行时原生库重写计算图,包括合并计算节点、消除冗余以提高运行时效率。
  3. ONNX 形状推理。

这些步骤的目标是提高量化质量。我们的量化工具在张量形状已知时效果最好。符号形状推理和 ONNX 形状推理都有助于确定张量形状。符号形状推理最适合基于 Transformer 的模型,而 ONNX 形状推理适用于其他模型。

模型优化会执行某些算子融合,这使得量化工具的工作更容易。例如,Convolution 算子后跟 BatchNormalization 可以在优化期间融合为一个,从而可以非常高效地进行量化。

不幸的是,ONNX 运行时中一个已知问题是模型优化无法输出大于 2GB 的模型。因此对于大型模型,必须跳过优化。

预处理 API 位于 Python 模块 onnxruntime.quantization.shape_inference 中的 quant_pre_process() 函数。参见 shape_inference.py。要了解更多关于预处理的选项和更精细的控制,请运行以下命令

python -m onnxruntime.quantization.preprocess --help

模型优化也可以在量化期间执行。然而,尽管由于历史原因这是默认行为,但这并*不*推荐。在量化期间进行模型优化会给调试量化引起的精度损失带来困难,这一点将在后续章节中讨论。因此,最好在预处理阶段进行模型优化,而不是在量化期间。

动态量化

量化模型有两种方式:动态量化和静态量化。动态量化会动态计算激活的量化参数(比例和零点)。这些计算增加了推理成本,但通常比静态量化能获得更高的精度。

动态量化的 Python API 位于模块 onnxruntime.quantization.quantize 中的 quantize_dynamic() 函数。

静态量化

静态量化方法首先使用一组输入(称为校准数据)运行模型。在这些运行过程中,我们计算每个激活的量化参数。这些量化参数作为常量写入量化模型,并用于所有输入。我们的量化工具支持三种校准方法:MinMax、Entropy 和 Percentile。详情请参考 calibrate.py

静态量化的 Python API 位于模块 onnxruntime.quantization.quantize 中的 quantize_static() 函数。详情请参考 quantize.py

量化调试

量化不是无损转换,它可能对模型的精度产生负面影响。解决这个问题的一个方法是比较原始计算图与量化计算图的权重和激活张量,找出它们差异最大的地方,然后避免对这些张量进行量化,或者选择另一种量化/校准方法。这被称为量化调试。为了方便这个过程,我们提供了 Python API,用于匹配 float32 模型与其量化模型对应的权重和激活张量。

用于调试的 API 位于模块 onnxruntime.quantization.qdq_loss_debug 中,包含以下函数

  • 函数 create_weight_matching()。它接收一个 float32 模型及其量化模型,并输出一个字典,用于匹配这两个模型对应的权重。
  • 函数 modify_model_output_intermediate_tensors()。它接收一个 float32 或量化模型,并对其进行增强以保存其所有激活。
  • 函数 collect_activations()。它接收一个由 modify_model_output_intermediate_tensors() 增强过的模型和一个输入数据读取器,运行增强后的模型以收集所有激活。
  • 函数 create_activation_matching()。您可以想象,您在 float32 模型及其量化模型上运行 collect_activations(modify_model_output_intermediate_tensors()),以收集两组激活。此函数接收这两组激活,并匹配对应的激活,以便用户可以轻松比较它们。

总之,ONNX 运行时提供了 Python API,用于匹配 float32 模型与其量化模型对应的权重和激活张量。这使得用户可以轻松比较它们,从而找出差异最大的地方。

然而,在量化期间进行模型优化会给此调试过程带来困难,因为它可能显著改变计算图,导致量化模型与原始模型差异巨大。这使得很难匹配两个模型中对应的张量。因此,我们建议在预处理阶段进行模型优化,而不是在量化过程中进行。

示例

  • 动态量化
import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

model_fp32 = 'path/to/the/model.onnx'
model_quant = 'path/to/the/model.quant.onnx'
quantized_model = quantize_dynamic(model_fp32, model_quant)

方法选择

动态量化和静态量化的主要区别在于激活的比例和零点是如何计算的。对于静态量化,它们是使用校准数据集预先(离线)计算的。因此,激活在每次前向传递过程中具有相同的比例和零点。对于动态量化,它们是实时(在线)计算的,并且特定于每次前向传递。因此它们更准确,但也引入了额外的计算开销。

一般来说,建议对 RNN 和基于 Transformer 的模型使用动态量化,对 CNN 模型使用静态量化。

如果两种后训练量化方法都无法满足您的精度目标,您可以尝试使用量化感知训练 (QAT) 对模型进行再训练。ONNX 运行时目前不提供再训练功能,但您可以使用原始框架对模型进行再训练,然后再将其转换回 ONNX 格式。

数据类型选择

量化值是 8 位宽的,可以是带符号的 (int8) 或无符号的 (uint8)。我们可以分别选择激活和权重的符号性,因此数据格式可以是(激活:uint8,权重:uint8)、(激活:uint8,权重:int8)等。我们使用 U8U8 作为(激活:uint8,权重:uint8)的简称,U8S8 作为(激活:uint8,权重:int8)的简称,类似地,S8U8 和 S8S8 代表其余两种格式。

ONNX 运行时在 CPU 上的量化可以运行 U8U8、U8S8 和 S8S8。使用 QDQ 的 S8S8 是默认设置,它平衡了性能和精度,应作为首选。只有当精度下降很多时,您可以尝试 U8U8。请注意,在 x86-64 CPU 上,使用 QOperator 的 S8S8 会很慢,通常应避免使用。ONNX 运行时在 GPU 上的量化仅支持 S8S8。

我何时以及为何需要尝试 U8U8?

在带有 AVX2 和 AVX512 扩展的 x86-64 机器上,ONNX 运行时使用 VPMADDUBSW 指令进行 U8S8 量化以提升性能。该指令可能存在饱和问题:输出可能无法完全放入 16 位整数,必须进行截断(饱和)以适应。一般来说,这对最终结果影响不大。但是,如果您确实遇到较大的精度下降,可能是由饱和引起的。在这种情况下,您可以尝试 reduce_range 或者尝试没有饱和问题的 U8U8 格式。

在其他 CPU 架构(带有 VNNI 的 x64 和 Arm®)上没有此问题。

支持的量化算子列表

支持的算子列表请参考注册表

量化和模型 opset 版本

模型必须是 opset10 或更高版本才能被量化。opset 小于 10 的模型必须使用更新的 opset 从其原始框架重新转换为 ONNX 格式。

基于 Transformer 的模型

针对基于 Transformer 的模型有一些特定的优化,例如用于注意力层量化的 QAttention。为了利用这些优化,您需要在量化模型之前使用Transformer 模型优化工具对模型进行优化。

此笔记本演示了该过程。

在 GPU 上量化

要在 GPU 上通过量化获得更好的性能,需要硬件支持。您需要一个支持 Tensor Core int8 计算的设备,如 T4 或 A100。较旧的硬件无法从量化中受益。

ONNX 运行时现在利用 TensorRT 执行提供程序在 GPU 上进行量化。与 CPU 执行提供程序不同,TensorRT 接受一个全精度模型和输入校准结果。它会根据自己的逻辑决定如何量化。利用 TensorRT EP 量化的总体步骤是

  • 实现一个 CalibrationDataReader
  • 使用校准数据集计算量化参数。注意:为了包含模型中的所有张量以获得更好的校准效果,请先运行 symbolic_shape_infer.py。详情请参考此处
  • 将量化参数保存到 flatbuffer 文件中
  • 加载模型和量化参数文件,并使用 TensorRT EP 运行。

我们提供了两个端到端示例:Yolo V3resnet50

量化到 Int4/UInt4

ONNX 运行时可以将模型中的某些算子量化为 4 位整数类型。对这些算子应用分块仅权重量化。支持的算子类型有

  • MatMul:
    • 仅当输入 B 是常量时才对该节点进行量化
    • 支持 QOperator 或 QDQ 格式。
    • 如果选择 QOperator 格式,该节点将被转换为 MatMulNBits 节点。权重 B 会被分块量化并保存在新节点中。支持 HQQGPTQ 和 RTN(默认)算法。
    • 如果选择 QDQ 格式,MatMul 节点将被替换为 DequantizeLinear -> MatMul 对。权重 B 会被分块量化并作为初始化器保存在 DequantizeLinear 节点中。
  • Gather:
    • 仅当输入 data 是常量时才对该节点进行量化。
    • 支持 QOperator 格式
    • Gather 被量化为 GatherBlockQuantized 节点。输入 data 会被分块量化并保存在新节点中。仅支持 RTN 算法。

由于 Int4/UInt4 类型是在onnx opset 21 中引入的,如果模型的 onnx domain 版本小于 21,它将被强制升级到 opset 21。请确保模型中的算子与 onnx opset 21 兼容。

要运行包含 GatherBlockQuantized 节点的模型,需要 ONNX 运行时 1.20 版本。

代码示例

from onnxruntime.quantization import (
    matmul_4bits_quantizer,
    quant_utils,
    quantize
)
from pathlib import Path

model_fp32_path="path/to/orignal/model.onnx"
model_int4_path="path/to/save/quantized/model.onnx"

quant_config = matmul_4bits_quantizer.DefaultWeightOnlyQuantConfig(
  block_size=128, # 2's exponential and >= 16
  is_symmetric=True, # if true, quantize to Int4. otherwsie, quantize to uint4.
  accuracy_level=4, # used by MatMulNbits, see https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#attributes-35
  quant_format=quant_utils.QuantFormat.QOperator, 
  op_types_to_quantize=("MatMul","Gather"), # specify which op types to quantize
  quant_axes=(("MatMul", 0), ("Gather", 1),) # specify which axis to quantize for an op type.

model = quant_utils.load_model_with_shape_infer(Path(model_fp32_path))
quant = matmul_4bits_quantizer.MatMul4BitsQuantizer(
  model, 
  nodes_to_exclude=None, # specify a list of nodes to exclude from quantizaiton
  nodes_to_include=None, # specify a list of nodes to force include from quantization
  algo_config=quant_config,)
quant.process()
quant.model.save_model_to_file(
  model_int4_path,
  True) # save data to external file

有关 AWQ 和 GTPQ 量化的用法,请参考Gen-AI 模型构建器

常见问题

为什么我没有看到性能提升?

性能提升取决于您的模型和硬件。量化带来的性能提升体现在计算和内存两个方面。较旧的硬件没有或只有少量执行高效 int8 推理所需的指令。而且量化本身存在开销(量化和反量化),因此在旧设备上性能变差并非罕见。

通常情况下,带有 VNNI 的 x86-64、支持 Tensor Core int8 的 GPU 以及带有点积指令的 Arm® 处理器可以获得更好的性能。

我应该选择哪种量化方法,动态还是静态?

请参考方法选择章节。

何时使用 reduce-range 和 per-channel 量化?

Reduce-range 将权重量化为 7 位。它设计用于 AVX2 和 AVX512(非 VNNI)机器上的 U8S8 格式,以缓解饱和问题。在支持 VNNI 的机器上则不需要。

对于权重范围较大的模型,per-channel 量化可以提高精度。如果精度损失较大,请尝试使用。在 AVX2 和 AVX512 机器上,如果启用 per-channel,通常还需要同时启用 reduce-range。

为什么 MaxPool 等算子没有被量化?

ONNX opset 12 中添加了对 MaxPool 等某些算子的 8 位类型支持。请检查您的模型版本,并将其升级到 opset 12 或更高版本。