📚 课程目录第3章 TensorRT量化工具链总览

3. TensorRT量化工具链总览

TensorRT INT8量化并不是一个单点操作，而是一条从训练框架到推理引擎的完整工具链。理解这条链路上每个环节的输入、输出和职责边界，是后续做校准、精度验证和性能调优的前提。本节把这条链路拆成五个关键阶段：PyTorch导出ONNX、ONNX解析、TensorRT Builder与Network定义、校准器接口、引擎序列化与运行时加载。它们各自解决不同的问题，又通过明确的数据契约串联在一起。

从PyTorch到ONNX：把计算图“冻结”成标准格式

PyTorch本身是动态图框架，训练时每一步都依赖Python解释器执行。TensorRT无法直接消费PyTorch的nn.Module，因此第一步是把模型导出为ONNX（Open Neural Network Exchange）。导出的本质是：用一组示例输入触发一次前向计算，把动态执行路径“追踪”成静态计算图，并记录每个算子的类型、属性、输入输出张量形状。

典型导出代码如下，注意opset_version和动态轴的处理：

import torch
import torch.onnx

model = MyModel().eval()
dummy_input = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    opset_version=13,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}}
)

这里有几个容易踩坑的点。第一，必须调用model.eval()，否则BatchNorm和Dropout会以训练模式导出，导致推理结果异常。第二，opset_version要与TensorRT版本匹配，较新的TensorRT对opset 13～17支持较好，过高的opset可能引入尚未实现的算子。第三，dynamic_axes用于声明哪些维度是动态的，INT8校准通常需要固定batch维度，因此如果后续要做校准，建议先以固定batch导出，或在校准阶段显式指定batch大小。

导出完成后，建议用onnxruntime做一次数值比对，确认ONNX输出与PyTorch输出在1e-4量级内一致，再进入下一步。这一步的验证成本很低，却能避免后续把“导出错误”误判为“量化误差”。

ONNX解析：TensorRT如何“读懂”计算图

TensorRT通过ONNX Parser把ONNX文件转换成内部的Network定义。解析过程不是简单的格式转换，而是一次语义重建：Parser会遍历ONNX的节点列表，把每个算子映射为TensorRT的层（Layer），同时推断张量形状、数据类型和内存布局。

解析阶段最常遇到的问题是“不支持的算子”。TensorRT并非支持ONNX全部算子集，遇到不支持的算子时，Parser会报错并指出节点名称。常见处理方式有三种：一是用ONNX Simplifier或自定义脚本把不支持的算子替换为等价组合；二是通过IPluginV2或IPluginV3编写自定义插件；三是回到PyTorch侧修改模型结构，避免使用冷门算子。

另一个关键点是形状推断。如果ONNX中某些维度是动态的，Parser需要显式设置优化profile（Optimization Profile），告诉TensorRT每个动态维度的最小、最优、最大取值。INT8校准通常要求形状确定，因此建议在校准前把batch和空间维度固定下来。

Builder与Network：定义优化空间

TensorRT Builder是引擎构建的入口，它负责根据Network定义和BuilderConfig生成优化后的引擎。Network是Parser输出的中间表示，包含了所有层、张量和连接关系。BuilderConfig则承载了精度、显存、 tactic选择等策略。

在INT8量化场景中，BuilderConfig需要显式设置setFlag(trt.BuilderFlag.INT8)，并指定校准器。如果同时希望某些层保持FP16或FP32，可以设置trt.BuilderFlag.FP16或trt.BuilderFlag.STRICT_TYPES，再通过层级精度覆盖实现混合精度。Builder在构建时会为每一层搜索可用的tactic，INT8模式下会优先选择INT8 kernel，但若某层没有INT8实现，会自动回退到FP16或FP32。

BuilderConfig中还有两个与INT8强相关的参数：setMaxWorkspaceSize决定构建时可用的显存上限，影响tactic搜索空间；setCalibrator绑定校准器对象。校准器必须在buildEngine之前设置，否则INT8 flag不会生效。

校准器接口：INT8量化的“标尺”

INT8量化的核心是确定每个张量的量化尺度（scale）。TensorRT采用对称线性量化，即q = round(x / scale)，其中scale由校准过程统计得到。校准器接口IInt8Calibrator定义了四个必须实现的方法：

getBatchSize()：返回校准批次大小，通常与校准数据加载器的batch一致。
getBatch()：按批次提供校准数据，返回输入张量指针和是否还有数据的标志。
readCalibrationCache()：尝试读取已缓存的校准表，避免重复校准。
writeCalibrationCache()：校准完成后写入缓存，供后续构建复用。

TensorRT提供了几种内置校准器：IInt8EntropyCalibrator2（基于熵最小化，推荐默认使用）、IInt8MinMaxCalibrator（基于最小最大值，适合激活分布较均匀的网络）、IInt8EntropyCalibrator（旧版熵校准）。选择哪种校准器取决于网络类型：CNN通常用EntropyCalibrator2，Transformer类模型可尝试MinMaxCalibrator。

校准数据的选择同样关键。一般建议从验证集中随机抽取500～1000张样本，覆盖主要类别和场景。样本太少会导致scale估计偏差，样本太多则校准时间线性增长。校准过程不需要标签，只需要输入张量。

引擎序列化与运行时加载

Builder构建完成后得到的是ICudaEngine对象，它包含了优化后的计算图和权重。但引擎不能跨进程直接使用，需要序列化为字节流（plan文件）保存到磁盘。序列化通过engine.serialize()完成，反序列化通过trt.Runtime(TRT_LOGGER).deserialize_cuda_engine(plan_bytes)完成。

序列化的意义在于：引擎构建（尤其是INT8校准）耗时较长，通常几分钟到几十分钟；而反序列化只需几百毫秒。因此生产部署中，标准做法是离线构建引擎并保存plan文件，运行时直接加载。需要注意的是，plan文件与TensorRT版本、GPU架构、CUDA版本强绑定，跨环境加载可能失败，必须保证构建和运行环境一致。

运行时加载后，通过engine.create_execution_context()创建执行上下文，再绑定输入输出显存地址，即可调用context.execute_async_v3()执行推理。INT8引擎在运行时不需要校准器，校准信息已经固化在引擎权重中。

工具链的边界与协作关系

把五个阶段放在一起看，可以总结出几条工程上的边界规则。第一，PyTorch导出ONNX是“格式转换”，不涉及量化；量化发生在Builder阶段，由校准器提供scale。第二，ONNX Parser是“语义翻译”，它决定了哪些算子能被TensorRT识别，不支持的算子必须在此阶段解决。第三，Builder是“优化决策”，它决定每层用什么精度、什么tactic，校准器只是它的输入之一。第四，序列化是“持久化”，它把构建结果固化为plan文件，使运行时加载成为可能。第五，Runtime是“执行”，它不参与任何量化决策，只负责按引擎定义调度kernel。