4、PTQ实战：使用PyTorch FX进行PTQ、使用TensorRT进行PTQ、使用ONNX Runtime进行PTQ

好，咱们直接进入正题。PTQ（训练后量化）是实际落地中最常用的手段。为什么？因为它快，不需要重新训练模型。你想想看，一个几百兆的大模型，跑一遍PTQ可能就几分钟，而QAT得折腾好几天。

我个人习惯是：先上PTQ，看看精度损失能不能接受。如果不行，再考虑QAT。今天咱们就聊聊三个主流框架的PTQ实战：PyTorch FX、TensorRT、ONNX Runtime。

4.1 用PyTorch FX做PTQ

PyTorch FX是PyTorch官方推出的量化工具。它通过符号跟踪（symbolic trace）把模型转换成一张计算图，然后在这张图上插入量化/反量化节点。

我记得第一次用FX时，最大的坑是模型必须符合FX的trace规则。如果你的模型里有动态控制流（比如if-else），FX会直接报错。嗯，这里要注意。

下面是一个完整的PTQ流程：

import torch
import torch.quantization as quant
from torch.quantization import get_default_qconfig
from torch.quantization.quantize_fx import prepare_fx, convert_fx

# 1. 准备模型
model = MyModel().eval()

# 2. 设置量化配置
qconfig = get_default_qconfig('fbgemm')  # x86用fbgemm，arm用qnnpack
model.qconfig = qconfig

# 3. 准备量化（插入观察器）
example_inputs = (torch.randn(1, 3, 224, 224),)
model_prepared = prepare_fx(model, qconfig, example_inputs)

# 4. 校准（喂入校准数据）
for batch in calib_dataloader:
    model_prepared(batch)

# 5. 转换为量化模型
model_quantized = convert_fx(model_prepared)

# 6. 推理
with torch.no_grad():
    output = model_quantized(test_input)

FX的优点是跟PyTorch生态无缝集成。缺点也很明显——对模型结构有要求。如果你的模型里有自定义算子，FX可能处理不了。

4.2 用TensorRT做PTQ

TensorRT是NVIDIA的推理优化库。它的PTQ走的是另一条路——用INT8校准器（calibrator）来统计激活值分布。

说白了，TensorRT的PTQ分两步：第一步，跑一遍校准数据，收集每层的激活值直方图；第二步，根据直方图选择最优的量化阈值。

我项目中遇到过最头疼的问题：TensorRT的INT8量化对某些层特别敏感，稍微量化一下精度就崩了。后来发现，可以手动指定哪些层跳过量化。

import tensorrt as trt

# 创建builder和config
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network()
config = builder.create_builder_config()

# 设置INT8模式
config.set_flag(trt.BuilderFlag.INT8)

# 自定义校准器
class MyCalibrator(trt.IInt8EntropyCalibrator2):
    def __init__(self, calib_loader):
        trt.IInt8EntropyCalibrator2.__init__(self)
        self.calib_loader = calib_loader
        self.buffer = None
        
    def get_batch(self, names):
        try:
            data = next(self.calib_loader)
            self.buffer = np.ascontiguousarray(data)
            return [self.buffer]
        except StopIteration:
            return None

# 设置校准器
calibrator = MyCalibrator(calib_loader)
config.int8_calibrator = calibrator

# 构建引擎
engine = builder.build_serialized_network(network, config)

TensorRT的INT8量化效果通常很好，尤其是对卷积层。但要注意，TensorRT只支持部分算子。如果你的模型里有不支持的算子，它会回退到FP32，影响加速效果。

4.3 用ONNX Runtime做PTQ

ONNX Runtime（ORT）是微软开源的推理引擎。它的PTQ支持两种模式：静态量化和动态量化。

静态量化需要校准数据，跟FX类似。动态量化不需要校准数据，但精度通常差一些。我个人习惯用静态量化，除非模型太大、校准数据不好找。

ORT的PTQ接口很简洁：

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

# 加载ONNX模型
model_fp32 = 'model.onnx'
model_int8 = 'model_int8.onnx'

# 定义校准数据读取器
class MyDataReader:
    def __init__(self, dataloader):
        self.dataloader = dataloader
        self.iter = iter(dataloader)
    
    def get_next(self):
        try:
            data = next(self.iter)
            return {'input': data.numpy()}
        except StopIteration:
            return None

# 执行静态量化
quantize_static(
    model_input=model_fp32,
    model_output=model_int8,
    calibration_data_reader=MyDataReader(calib_loader),
    quant_format=QuantType.QInt8,
    per_channel=True,
    calibration_method=CalibrationMethod.MinMax
)

ORT的优点是跨平台支持好，Windows、Linux、ARM都能跑。缺点是对某些复杂模型的支持不如TensorRT成熟。

4.4 三个框架的对比

你可能会问：到底该选哪个？我个人的建议是：

框架	适用场景	精度	速度	易用性
PyTorch FX	PyTorch生态、快速验证	中等	中等	高（需模型兼容）
TensorRT	NVIDIA GPU、生产部署	高	最高	中（需转ONNX）
ONNX Runtime	跨平台、多后端	中等	高	高

说白了，如果你在NVIDIA GPU上部署，TensorRT是首选。如果你需要跨平台，ORT更合适。如果你只是快速验证PTQ效果，FX最方便。

4.5 PTQ精度调优的通用技巧

不管用哪个框架，PTQ精度调优的思路是相通的。我总结了几条经验：

• 校准数据要够多、够多样——500-1000张，覆盖各种场景
• 逐层分析精度损失——找出哪些层量化后损失最大，单独处理
• 尝试不同的校准方法——MinMax、Entropy、Percentile，效果差别很大
• 考虑per_channel量化——对权重用per_channel，精度提升明显
• 敏感层跳过量化——如果某层量化后精度崩了，就让它跑FP32

好了，PTQ实战就聊到这里。三个框架各有千秋，关键是根据你的场景选对工具。记住：PTQ是量化的第一步，也是性价比最高的一步。