此ONNX PTQ工具包提供了一套全面的工具，旨在通过量化优化ONNX（开放神经网络交换）模型。我们的工具包旨在帮助开发人员在使用TensorRT部署时，在不影响神经网络准确性的情况下提高性能、减小模型大小和加快推理时间。

量化是一种有效的模型优化技术，可以压缩模型。使用模型优化器进行量化可以将模型大小压缩2x-4x，在保持模型质量的同时加快推理速度。

Model Optimizer支持NVFP4、FP8、INT8、INT4等高性能量化格式，并通过易于使用的Python API支持AWQ和双量化等高级算法。

先决条件

Docker

请使用TensorRT docker镜像（例如，nvcr.io/nvidia/tensort:25.08-py3）或访问安装文档以获取更多信息。

在TensorRT docker中设置以下环境变量。

export CUDNN_LIB_DIR=/usr/lib/x86_64-linux-gnu/
export LD_LIBRARY_PATH="${CUDNN_LIB_DIR}:${LD_LIBRARY_PATH}"

还要按照以下安装步骤升级到最新版本的Model Optimizer，并安装特定于示例的依赖关系

本地安装

使用PyPI中的pip安装具有onnx依赖关系的模型优化器，并安装示例的要求：

pip install -U nvidia-modelopt[onnx]
pip install -r requirements.txt

对于TensorRT编译器框架工作负载：

从这里安装最新的TensorRT

现在开始

准备样例模型

本文档中的大部分示例均采用vit_base_patch16_224.onnx作为输入模型。可通过以下脚本下载该模型：

python download_example_onnx.py \
    --vit \
    --onnx_save_path=vit_base_patch16_224.onnx \
    --fp16 # <Optional, if the desired output ONNX precision is FP16>

准备校准数据

校准数据是训练或验证数据集中用于量化过程的一个代表性子集，用于确定将浮点数值转换为低精度格式（INT8、FP8、INT4）的最佳缩放因子。通过分析网络中激活值的分布，这些数据有助于在量化后保持模型的准确性。

首先，准备一些校准数据。TensorRT建议卷积神经网络（CNN）和视觉变换器（ViT）模型的校准数据量至少为500。以下命令将从tiny-imagenet数据集中选取500张图像，并将其转换为numpy格式的校准数组。在资源受限的环境中，可减少校准数据量。

python image_prep.py \
    --calibration_data_size=500 \
    --output_path=calib.npy \
    --fp16 # <Optional, if the input ONNX is in FP16 precision>

将ONNX模型量化为FP8、INT8或INT4

可以使用CLI或Python API将模型量化为FP8、INT8或INT4模型。对于FP8和INT8量化，您可以在最大值和熵校准算法之间进行选择。对于INT4量化，可以选择awq_clip或rtn_dq算法。

有关NVFP4和MXFP8 ONNX，请参阅PyTorch到ONNX部分。

最低操作系统要求：int8（13+），fp8（21+），int4（21+。ModelOpt将自动升级较低的opset版本以满足这些要求。

方法一:使用命令行

python -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch16_224.onnx \
    --quantize_mode=<fp8|int8|int4> \
    --calibration_data=calib.npy \
    --calibration_method=<max|entropy|awq_clip|rtn_dq> \
    --output_path=vit_base_patch16_224.quant.onnx

方法二：python api

from modelopt.onnx.quantization import quantize

quantize(
    onnx_path="vit_base_patch16_224.onnx",
    quantize_mode="int8",       # fp8, int8, int4 etc.
    calibration_data="calib.npy",
    calibration_method="max",   # max, entropy, awq_clip, rtn_dq etc.
    output_path="vit_base_patch16_224.quant.onnx",
)

评估量化的ONNX模型

以下评估需要ImageNet数据集的val目录。或者，您可以从这个hugging face数据集中准备它。一旦你有了它，量化的ONNX ViT模型就可以在ImageNet数据集上进行评估，如下所示

python evaluate.py \
    --onnx_path=<path to classification model> \
    --imagenet_path=<path to the ImageNet dataset> \
    --engine_precision=stronglyTyped \
    --model_name=vit_base_patch16_224

此脚本将量化的ONNX模型转换为TensorRT引擎，并使用该引擎进行评估。最后，评估结果将报告如下：

The top1 accuracy of the model is <accuracy score between 0-100%>
The top5 accuracy of the model is <accuracy score between 0-100%>
Inference latency of the model is <X> ms

torch量化到ONNX示例，用于MXFP8、INT4或NVFP4精度

此示例演示了如何使用MXFP8、INT4或NVFP4精度格式量化timm-vision模型，然后将其导出到ONNX。该脚本利用ModelOpt工具包进行量化和ONNX导出。

Opset 20 is used to export the torch models to ONNX.

它做了什么

• 加载预训练的timm-torch模型（默认值：ViT-Base）。

• 使用ModelOpt将焊炬模型量化为MXFP8、INT4或NVFP4。

• 将量化模型导出到ONNX。

• 对ONNX模型进行后处理，使其与TensorRT兼容。

• 保存最终ONNX模型。

使用

python torch_quant_to_onnx.py \
    --timm_model_name=vit_base_patch16_224 \
    --quantize_mode=<mxfp8|nvfp4|int4_awq> \
    --onnx_save_path=<path to save the exported ONNX model>

评价

如果输入模型属于图像分类类型，请使用以下脚本对其进行评估。

注意：评估MXFP8或NVFP4 ONNX模型需要TensorRT 10.11或更高版本。

python evaluate.py \
    --onnx_path=<path to the exported ONNX model> \
    --imagenet_path=<path to the ImageNet dataset> \
    --engine_precision=stronglyTyped \
    --model_name=vit_base_patch16_224

ONNX导出支持的LLM模型

高级功能

ONNX格式的每节点校准

每节点校准是一种内存优化功能，旨在减少大型ONNX模型量化过程中的内存消耗。该功能不是一次在整个网络上运行推理，而是逐个节点处理模型，这可以显著降低峰值内存使用率并防止内存不足（OOM）错误

运作原理

启用每节点校准时，量化过程：

• 分解模型：将原始ONNX模型拆分为多个单节点子模型

• 管理依赖关系：跟踪节点之间的输入/输出依赖关系，以确保正确的执行顺序

• 按顺序处理：使用拓扑处理顺序在每个节点上单独运行校准

• 管理内存：自动清理中间结果并管理引用计数，以最大限度地减少内存使用

• 聚合结果：组合来自所有节点的校准数据，以生成最终的量化模型

何时使用每节点校准

每节点校准对以下方面特别有益：

• 在标准校准过程中导致OOM误差的大型模型

• GPU内存有限的内存受限环境

• 具有高中间内存要求的复杂架构的模型

用法

要启用每节点校准，请在量化命令中添加--calibrate_per_node标志：

python -m modelopt.onnx.quantization \
    --onnx_path=vit_base_patch16_224.onnx \
    --quantize_mode=<int8/fp8> \
    --calibration_data=calib.npy \
    --calibrate_per_node \
    --output_path=vit_base_patch16_224.quant.onnx

注意：INT4量化方法（awq_clip、rtn_dq）不支持每节点校准

使用自定义操作量化ONNX模型

此功能需要TensorRT 10+和ORT>=1.20。为了正确使用，请确保libcudn*.so和TensorRT-lib/的路径在LD_LIBRARY_PATH环境变量中，并且安装了TensorRT-python包。

请参阅下面的示例

步骤1：从TensorRT自定义插件示例中获取示例ONNX模型和TensorRT插件。

1.1.将目录更改为TensorRT自定义插件示例：

cd /path/to/TensorRT-Custom-Plugin-Example

1.2.编译TensorRT插件：

cmake -B build \
    -DNVINFER_LIB=$TRT_LIBPATH/libnvinfer.so.10 \
    -DNVINFER_PLUGIN_LIB=$TRT_LIBPATH/libnvinfer_plugin.so.10 \
    -DNVONNXPARSER_LIB=$TRT_LIBPATH/libnvonnxparser.so.10 \
    -DCMAKE_CXX_STANDARD_INCLUDE_DIRECTORIES=/usr/include/x86_64-linux-gnu

cmake --build build --config Release --parallel

这将在TensorRT自定义插件示例/build/src/plugins/IdentityConvIPluginV2IOExt/libidentity_conv_iplgin_v2_io_ext.so中生成一个插件

1.3.创建ONNX文件

python scripts/create_identity_neural_network.py

这将在TensorRT自定义插件示例/data/identity_neural_network.onnx中生成identity_neural_network模型

步骤2：量化ONNX模型。在这个例子中，我们将使用libidentity_conv_iplogin_v2_io_ext.so插件。

python -m modelopt.onnx.quantization \
    --onnx_path=/path/to/identity_neural_network.onnx \
    --trt_plugins=/path/to/libidentity_conv_iplugin_v2_io_ext.so

步骤3：使用TensorRT部署量化模型

trtexec --onnx=/path/to/identity_neural_network.quant.onnx \
    --staticPlugins=/path/to/libidentity_conv_iplugin_v2_io_ext.so

量化好的checkpoint

已准备好部署可导出为ONNX格式的检查点（如果在支持列表支持）[hugging face-Nvidia TensorRT模型优化器集合]

技术资源

示例脚本中支持许多量化方案

• FP8格式可在Hopper和Ada GPU上使用，CUDA计算能力大于或等于8.9。

• INT4 AWQ是一种仅考虑INT4权重的量化和校准方法。INT4 AWQ对于低批量推理特别有效，其中推理延迟主要由权重加载时间而不是计算时间本身决定。对于低批量推理，INT4 AWQ可以比FP8/INT8提供更低的延迟，比INT8提供更少的精度下降。

• NVFP4是NVIDIA Blackwell GPU支持的新FP4格式之一，与其他4位替代品相比具有良好的准确性。NVFP4可以应用于模型权重和激活，与Blackwell上的FP8数据格式相比，它有可能显著提高数学吞吐量，减少内存占用和内存带宽使用。