3. 量化工具与框架：TensorRT、ONNX Runtime、PyTorch

做模型量化，光懂理论远远不够。

你得有趁手的工具。市场上量化工具不少，但我个人最常用的，就是这三板斧：TensorRT、ONNX Runtime 和 PyTorch 自带工具。它们各有各的脾气，也各有各的绝活。

3.1 先看一张总览图

在聊细节之前，我先把它们的定位理清楚。下图展示了三者的核心关系。

3.2 TensorRT 量化：GPU 上的极致性能

如果你追求推理速度，那 TensorRT 就是绕不开的选择。我最早接触 TensorRT 是在一个视频分析项目里，模型死活跑不到 30fps，量化做完，直接飙到 90fps。差距就是这么大。

3.2.1 核心思路

TensorRT 量化走的是 校准（Calibration） 路线。你给它几百张代表性图片，它自己统计每层激活值的分布，然后算出最优的 INT8 缩放系数和零点。这比粗暴的 min-max 要好得多。

我记得有一次，我用 500 张校准图就搞定了，精度损失不到 0.5%。关键是校准集要覆盖真实场景。你拿高清写真图去校准，结果部署时全是监控视频，那精度大概率会崩。

3.2.2 基本流程

TensorRT 做 INT8 量化，步骤不复杂。我给出一个典型的代码骨架：

// 1. 构建推理引擎，指定精度为 INT8
IBuilder* builder = createInferBuilder(gLogger);
IBuilderConfig* config = builder->createBuilderConfig();
config->setFlag(BuilderFlag::kINT8);

// 2. 设置校准器（IInt8Calibrator）
//    这里需要自己实现，核心是提供校准数据和缓存
// 来自定义校正器
class MyCalibrator : public IInt8EntropyCalibrator2 {
    bool getBatch(void* bindings[], const char* names[], int nbBindings) override {
        // 填充校准数据到 bindings，返回 false 表示结束
        return true;
    }
    const void* readCalibrationCache(size_t& length) override { return nullptr; }
    void writeCalibrationCache(const void* cache, size_t length) override {}
};

// 3. 创建网络定义，导出量化引擎
INetworkDefinition* network = builder->createNetworkV2(0);
// 添加网络层...
auto calibrator = new MyCalibrator();
auto engine = builder->buildSerializedNetwork(*network, *config);

3.2.3 避坑指南

我曾经在一个 NLP 模型上做 INT8 量化，结果精度直接掉了 3 个点。反复排查后发现是校准集大小写错了，只用了 10 张图。校准集必须能代表真实数据的分布，不是随便拿几张就算完。

还有一点：TensorRT 对动态形状支持有限。如果模型输入尺寸频繁变化，量化后的引擎可能报错。要么固定输入尺寸，要么用多 Profile 方案。

3.3 ONNX Runtime 量化：跨平台的稳健选择

如果你要在 CPU 上跑，或者需要同时支持多平台，那 ONNX Runtime 更合适。它的量化工具集成在 onnxruntime-tools 和 onnxruntime-quantization 里。

3.3.1 两种模式

我一般先用动态量化跑一遍，看看效果。如果精度掉得厉害，再试试静态量化。动态量化不需要校准数据，改几行代码就能用，性价比很高。

3.3.2 实操：动态量化

动态量化的代码极其简洁。说白了，就是对权重做 INT8，激活值还保持 FP32 或 BF16。

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

model_fp32 = "model.onnx"
model_int8 = "model_int8.onnx"

quantize_dynamic(
    model_input=model_fp32,
    model_output=model_int8,
    weight_type=QuantType.QInt8
)
print("动态量化完成！")

这就完了。你想想看，就这么几行代码，模型大小直接减到 1/4，推理速度快 2-3 倍。不过要注意，动态量化本质上是 weight-only，算力瓶颈在计算上，所以实际加速可能不如静态量化明显。

3.3.3 实操：静态量化

静态量化需要校准数据。它会在推理前跑一遍校准，收集激活值的范围，然后生成量化表。

import onnx
from onnxruntime.quantization import quantize_static, QuantType, CalibrationDataReader

# 自己写一个数据读取器
class MyDataReader(CalibrationDataReader):
    def get_next(self):
        # 返回校准数据，格式为 dict: {input_name: np_array}
        return {"input": your_calib_data}
    def __init__(self):
        self.iterator = iter(your_dataset)

# 触发静态量化
quantize_static(
    model_input="model.onnx",
    model_output="model_static_int8.onnx",
    calibration_data_reader=MyDataReader(),
    quant_format=QuantType.QInt8,
    per_channel=True
)

3.4 PyTorch 量化：研究者的好工具

如果你在做模型压缩研究，或者需要精细控制每个量化参数，PyTorch 自带的量化工具是最灵活的。它支持两种模式：Eager 模式 和 FX 图模式。

3.4.1 FX 图模式

我个人强烈推荐 FX 图模式。它做了一次图捕获，把整个计算图转换成统一表示，然后自动插入量化/反量化节点。你不需要手动改动每一层。

import torch
import torch.quantization as quant

model_fp32 = MyModel().eval()

# 准备：配置量化后端
model_prepared = quant.prepare(
    model_fp32,
    inplace=False,
    backend='fbgemm'  # x86 用 fbgemm，arm 用 qnnpack
)

# 校准：跑几个 batch，收集统计量
for images in calib_loader:
    model_prepared(images)

# 转换：执行量化
model_int8 = quant.convert(model_prepared, inplace=False)
print("FX 图模式量化完成！")

为什么好用？你仔细看，这里没有一句 for 循环去手动改层。FX 会自动分析模型结构，识别哪些层可以融合，哪些层需要跳过。它聪明着呢。

3.4.2 Eager 模式：老派但可控

Eager 模式需要你手动在模型前向中定义量化节点。适合你对每个模块有精确控制。

import torch
import torch.nn as nn
import torch.quantization as quant

class QuantModel(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.quant = quant.QuantStub()
        self.dequant = quant.DeQuantStub()
        self.model = model

    def forward(self, x):
        x = self.quant(x)        # 输入量化
        x = self.model(x)        # 推理在 INT8 中
        x = self.dequant(x)      # 输出反量化
        return x

嗯，这样写起来确实有点原始，但如果你需要对特定层手动设置量化参数，Eager 模式更顺手。

3.4.3 QAT（量化感知训练）

如果 Post-Training Quantization 掉精度，你可以试试 QAT。PyTorch 的 QAT 集成得很好。

# 准备 QAT 模型
model_qat = quant.prepare_qat(model_fp32, inplace=False)

# 微调几个 epoch
for images, labels in train_loader:
    output = model_qat(images)
    loss = criterion(output, labels)
    loss.backward()
    optimizer.step()

# 转换
model_int8 = quant.convert(model_qat, inplace=False)

我记得在一个人脸识别模型上，PTQ 掉了 2 个点，QAT 只掉了 0.2 个点。多花两天时间微调，值了。

3.5 如何选？一张决策表