ONNX 量化 - Linux (测试版)

ModelOpt 提供了与 TensorRT 显式量化 (EQ) 协同工作的 ONNX 量化。ModelOpt 的 ONNX 量化的主要优势包括：

易于非专家用户使用。
白盒设计允许专家用户自定义量化过程。
更好地支持视觉变换器。

目前ONNX量化支持FP8、INT4和INT8量化。

注意

ModelOpt ONNX量化根据TensorRT规则生成带有QDQ节点的新ONNX模型。为了获得实际的速度提升，生成的ONNX应编译为TensorRT引擎。

需求

TensorRT >= 8.6 (推荐 >= 10.0)。请参考 TensorRT 10.0 下载链接。

应用训练后量化（PTQ）

PTQ 应该使用校准数据集来完成。如果没有提供校准数据集，ModelOpt 将使用随机比例来设置 QDQ 节点。

准备校准数据集

ModelOpt 支持 npz/npy 文件作为校准数据格式，并且该 numpy 文件应该是一个字典，其中键为模型输入名称，值为 numpy 数组。

# Example numpy file for single-input ONNX
calib_data = np.random.randn(batch_size, channels, h, w)
np.save("calib_data.npy", calib_data)

# Example numpy file for single/multi-input ONNX
# Dict key should match the input names of ONNX
calib_data = {
    "input_name": np.random.randn(*shape),
    "input_name2": np.random.randn(*shape2),
}
np.savez("calib_data.npz", calib_data)

调用PTQ函数

import modelopt.onnx.quantization as moq

calibration_data = np.load(calibration_data_path)

moq.quantize(
    onnx_path=onnx_path,
    calibration_data=calibration_data,
    output_path="quant.onnx",
    quantize_mode="int8",
)

或者，您可以在命令行中调用 PTQ 函数：

python -m modelopt.onnx.quantization \
    --calibration_data_path /calibration/data/in/npz/npy/format \
    --output_path /path/to/the/quantized/onnx/output \
    --quantize_mode int8

默认情况下，运行校准后，量化工具将按照TensorRT友好的QDQ插入算法插入QDQ节点。用户可以通过调整API参数（如op_types_to_quantize、op_types_to_exclude等）来更改默认的量化行为。详情请参见modelopt.onnx.quantization.quantize()。

部署量化ONNX模型

trtexec 是 TensorRT 提供的一个命令行工具。通常，它位于 /usr/src/tensorrt/bin/ 目录下。以下是一个简单的命令，用于将上一步生成的量化 onnx 模型编译为 TensorRT 引擎文件。

trtexec --onnx=quant.onnx --saveEngine=quant.engine --best

比较性能

以下命令将使用fp16精度构建引擎。构建完成后，检查报告的“延迟”和“吞吐量”字段并进行比较。

trtexec --onnx=original.onnx --saveEngine=fp16.engine --fp16

注意

如果将--fp16标志替换为--best标志，此命令将使用TensorRT的隐式量化创建一个int8引擎。