量化 ONNX 模型

目录

量化概述

ONNX Runtime 中的量化指的是 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(tensor)) 以模拟量化和反量化过程。
    在静态量化中,QuantizeLinear 和 DeQuantizeLinear 运算符也携带量化参数。
    在动态量化中,插入 ComputeQuantizationParameters 函数原型以动态计算量化参数。
  • 以下列方式生成的模型采用 QDQ 格式
    1. 通过 quantize_static 量化的模型(如下所述),使用 quant_format=QuantFormat.QDQ
    2. 从 Tensorflow 转换或从 PyTorch 导出的量化感知训练 (QAT) 模型。
    3. 从 TFLite 和其他框架转换的量化模型。

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

下图显示了量化 Conv 的 QOperator 和 QDQ 格式的等效表示形式。此端到端示例演示了这两种格式。

Changes to nodes from basic and extended optimizations

量化 ONNX 模型

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

预处理

预处理是将 float32 模型转换为准备进行量化的过程。它包括以下三个可选步骤

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

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

模型优化执行某些运算符融合,使量化工具的工作更容易。例如,卷积运算符后跟 BatchNormalization 可以在优化期间融合为一个,这样可以非常有效地进行量化。

遗憾的是,ONNX Runtime 中一个已知的问题是模型优化无法输出大于 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 Runtime 提供了 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 Runtime 目前不提供重新训练,但您可以使用原始框架重新训练模型,然后将其转换回 ONNX。

数据类型选择

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

CPU 上的 ONNX Runtime 量化可以运行 U8U8、U8S8 和 S8S8。带有 QDQ 的 S8S8 是默认设置,可在性能和精度之间取得平衡。它应该是首选。只有在精度大幅下降的情况下,您才可以尝试 U8U8。请注意,带有 QOperator 的 S8S8 在 x86-64 CPU 上会很慢,一般应避免使用。GPU 上的 ONNX Runtime 量化仅支持 S8S8。

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

在具有 AVX2 和 AVX512 扩展的 x86-64 机器上,ONNX Runtime 使用 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 Runtime 现在利用 TensorRT 执行提供程序进行 GPU 上的量化。与 CPU 执行提供程序不同,TensorRT 接受全精度模型和输入的校准结果。它使用自己的逻辑决定如何量化。利用 TensorRT EP 量化的总体过程是

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

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

量化为 Int4/UInt4

ONNX Runtime 可以将模型中的某些运算符量化为 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 域版本 < 21,则会强制升级到 opset 21。请确保模型中的运算符与 onnx opset 21 兼容。

要运行具有 GatherBlockQuantized 节点的模型,需要 ONNX Runtime 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 这样的运算符没有被量化?

某些运算符(例如 MaxPool)的 8 位类型支持已添加到 ONNX opset 12 中。请检查您的模型版本并将其升级到 opset 12 及以上。