4、量化工具链实战：PyTorch量化工具包（torch.quantization）、TensorRT量化流程、ONNX Runtime量化

说实话，理论讲再多，最后都得落到工具上。我见过不少团队，量化原理说得头头是道，一上手就卡在工具链上——要么精度掉得莫名其妙，要么推理速度反而更慢。

这一章，咱们就把三个主流工具链挨个过一遍。PyTorch、TensorRT、ONNX Runtime，每个我都踩过坑，咱们直接聊实战。

核心观点： 没有最好的工具链，只有最适合你部署场景的工具链。PyTorch 适合快速验证，TensorRT 适合极致性能，ONNX Runtime 适合跨平台统一。

4.1 PyTorch 量化工具包（torch.quantization）

PyTorch 的量化工具包，我用了三年多。它最大的好处是跟训练流程无缝衔接。你想想看，模型本来就在 PyTorch 里训，量化时不用换框架，省心不少。

PyTorch 支持三种量化模式：

模式 校准方式 适用场景 精度损失
动态量化 运行时统计 LSTM、Transformer 等权重主导模型 较小
静态量化 校准数据集 CNN 等激活值范围固定的模型 中等
量化感知训练 训练中模拟量化 精度敏感场景 最小

动态量化 最简单。权重提前量化好，激活值在推理时实时计算范围。我一般用它快速验证量化效果，几分钟就能跑通。

import torch
from torch.quantization import quantize_dynamic

model = torch.load('model.pth')
quantized_model = quantize_dynamic(
    model,
    {torch.nn.Linear, torch.nn.LSTM},
    dtype=torch.qint8
)

静态量化 需要校准数据集。说白了，就是拿一批有代表性的数据跑一遍，统计每层激活值的 min/max 范围。我习惯用验证集的 100-200 张图片做校准，多了没必要，少了不够稳。

model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
model = torch.quantization.prepare(model, inplace=True)
# 校准循环
for data in calib_loader:
    model(data)
model = torch.quantization.convert(model, inplace=True)
注意： 静态量化后，模型里会出现 Quant/DeQuant 节点。我曾经遇到过一个模型，量化后推理反而变慢了——就是因为频繁的 Quant/DeQuant 操作抵消了计算加速。排查了半天才发现是某些层不支持 INT8 计算。

量化感知训练 是我最推荐的方式，尤其对精度要求高的场景。它在训练时插入 FakeQuantize 节点，模拟量化误差，让模型自己学会适应。

model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')
model = torch.quantization.prepare_qat(model, inplace=True)
# 正常训练
for epoch in range(epochs):
    for data, target in train_loader:
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
# 转换
model = torch.quantization.convert(model, inplace=True)

4.2 TensorRT 量化流程

TensorRT 是 NVIDIA 的推理优化引擎。我最早接触它是在一次自动驾驶项目上，模型在 PyTorch 里跑 30ms，转到 TensorRT 后直接降到 8ms。嗯，当时我就服了。

TensorRT 的量化流程分三步：

  1. 模型转换：把 PyTorch/ONNX 模型转成 TensorRT 的推理引擎
  2. 校准：用校准数据集计算 INT8 的缩放因子
  3. 构建引擎：生成优化后的推理引擎文件

我最常用的是 ONNX → TensorRT 这条路径。先把 PyTorch 模型导出为 ONNX，再用 TensorRT 的 ONNX Parser 加载。

import tensorrt as trt

TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)
parser = trt.OnnxParser(network, TRT_LOGGER)

with open('model.onnx', 'rb') as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

# INT8 校准
config.set_flag(trt.BuilderFlag.INT8)
calibrator = MyCalibrator(calib_data)
config.int8_calibrator = calibrator

engine = builder.build_serialized_network(network, config)
我的经验： TensorRT 的校准器有四种：IInt8EntropyCalibrator2（默认）、IInt8MinMaxCalibrator、IInt8EntropyCalibrator、IInt8LegacyCalibrator。我一般先用 EntropyCalibrator2，如果精度不够再换 MinMax。MinMax 对极端值更敏感，但有时反而精度更好。

TensorRT 的 INT8 量化，核心是 校准算法。它通过最小化 KL 散度来找到最优的阈值，把激活值映射到 INT8 范围。我遇到过一个问题：校准数据集跟实际推理数据分布不一致，导致精度崩了。后来我强制要求校准数据必须覆盖所有典型场景。

4.3 ONNX Runtime 量化

ONNX Runtime 是微软开源的推理引擎。它的优势在于跨平台——同一个 ONNX 模型，在 Windows、Linux、ARM 上都能跑。我有个项目需要在 Jetson 和 x86 服务器上同时部署，ONNX Runtime 就是最佳选择。

ONNX Runtime 的量化方式有两种：

方式 工具 特点
离线量化 onnxruntime.quantization 静态量化，需要校准数据
动态量化 onnxruntime.quantization 运行时统计，无需校准
QAT 模型导入 PyTorch/其他框架导出 精度最高，需要训练

离线量化 是我用得最多的。它跟 PyTorch 的静态量化类似，但 ONNX Runtime 的量化工具会自动处理算子融合和优化。

from onnxruntime.quantization import quantize_static, QuantType
from onnxruntime.quantization import CalibrationMethod

quantize_static(
    model_input='model.onnx',
    model_output='model_quant.onnx',
    calibration_data_reader=calib_reader,
    quant_format=QuantType.QInt8,
    calibrate_method=CalibrationMethod.MinMax,
    per_channel=True
)
避坑指南： 我曾经在 ONNX Runtime 量化时遇到一个坑——某些算子（比如 Resize、Pad）不支持 INT8 计算。量化工具会把这些算子留在 FP32，导致推理时频繁切换精度。解决办法是手动检查量化后的模型，用 Netron 可视化看看有没有多余的 Cast 节点。

动态量化 适合那些输入尺寸变化大的模型。比如 NLP 模型，句子长度不固定，静态量化很难选一个合适的校准集。动态量化虽然速度提升不如静态量化，但胜在通用。

下面这张图展示了三个工具链的量化流程对比：

量化工具链流程对比 PyTorch 模型 prepare / prepare_qat 校准 / QAT 训练 convert 转换 量化模型 (torch.jit) ONNX 模型 ONNX Parser 解析 INT8 校准器校准 Builder 构建引擎 TensorRT 引擎 (.engine) ONNX 模型 quantize_static/dynamic 校准数据读取 算子 INT8 替换 量化 ONNX 模型 选择建议 PyTorch：快速验证、研究原型 TensorRT：NVIDIA GPU 极致性能 ONNX Runtime：跨平台、多后端统一部署

4.4 工具链选型对比

我整理了一张对比表，方便你根据项目需求快速选型：

维度 PyTorch TensorRT ONNX Runtime
硬件支持 CPU/GPU 通用 NVIDIA GPU 专属 CPU/GPU/ARM 通用
量化精度 中等（QAT 较高） 高（校准算法优化） 中等
推理速度 中等 最快（算子融合+内核调优） 较快
易用性 高（与训练框架集成） 低（需要学习 TRT API） 中（ONNX 生态成熟）
调试难度 低（Python 原生） 高（黑盒引擎） 中（可视化工具支持）
社区生态 最活跃 NVIDIA 官方维护 微软+社区
我的建议： 如果你刚开始接触量化，先从 PyTorch 的动态量化入手。跑通了再尝试静态量化。等对量化原理有感觉了，再挑战 TensorRT 的 INT8 校准。别一上来就搞 TensorRT，容易劝退。

最后说一句，工具链只是手段。真正决定量化效果的，是你对模型的理解和对数据的把控。我见过有人用最简单的 PyTorch 动态量化，配合精心设计的校准集，效果比某些人用 TensorRT 瞎调还好。

← 上一章 📖 返回目录 下一章 →