4、TensorRT核心机制：层融合、精度校准(INT8/FP16)、动态形状、插件编写
聊到TensorRT，我脑子里蹦出来的第一个词就是「加速」。但加速不是魔法，它背后有一套非常扎实的工程机制。今天咱们就把这四个核心机制掰开揉碎，一个个讲清楚。
4.1 层融合：把碎片拼成利刃
你想想看，一个神经网络模型，动辄几百层。每一层都有读写操作、内存分配、kernel启动。这些开销累加起来，比实际计算还耗时。层融合要做的，就是把多个小操作合并成一个大的kernel。
我个人习惯把层融合分成三类：
- 垂直融合：比如Conv + BN + ReLU，三个小算子合成一个。我在项目中遇到过，一个ResNet-50不做融合，推理延迟是融合后的2.3倍。
- 水平融合：比如多个相同大小的卷积，可以合并成一个batch处理。说白了就是让GPU一次干更多活。
- 跨层融合：比如把残差结构中的add操作和前面的卷积融合。嗯，这里要注意，不是所有层都能随便融合，得看数据依赖关系。
核心收益：减少kernel启动次数、减少全局内存读写、提高缓存命中率。我实测过，层融合能带来30%-50%的加速。
我曾经踩过一个坑：自己手写了一个自定义层，结果TensorRT死活不给我做融合。后来发现是算子注册时没设置正确的tensor格式。从那以后，我写插件都会先检查一下格式兼容性。
4.2 精度校准：INT8/FP16的取舍之道
FP16加速是白送的，大部分GPU都支持。但INT8就不一样了——它需要校准。为什么会这样？因为FP16的精度损失很小，模型几乎无感。但INT8只有256个量化级别，一不小心精度就崩了。
我建议的INT8校准流程是这样的：
- 准备校准数据集：500-1000张有代表性的图片就够了。别用训练集，用验证集。
- 选择校准算法：TensorRT提供了三种——Entropy、MinMax、Percentile。我个人习惯用Entropy，它在大多数场景下表现最稳。
- 运行校准：TensorRT会统计每层的激活值分布，然后找到最优的量化阈值。
- 验证精度：跑一遍验证集，看精度下降是否在可接受范围内。
| 精度模式 | 加速比（相对FP32） | 精度损失 | 适用场景 |
|---|---|---|---|
| FP32 | 1x | 无 | 精度敏感场景 |
| FP16 | 1.5x - 2x | 几乎无 | 大部分场景 |
| INT8 | 2x - 4x | 0.5%-2% | 对延迟要求高的场景 |
小技巧：如果INT8精度下降太多，试试用Percentile校准，把阈值设到99.9%。我有个项目就是这么救回来的。
4.3 动态形状：别让输入尺寸卡住你
很多模型需要处理不同尺寸的输入，比如目标检测、OCR。TensorRT 8.0之后，动态形状的支持已经很成熟了。说白了就是允许你在运行时指定输入尺寸，而不是编译时写死。
配置动态形状需要三步：
- 定义优化范围：设置最小、最优、最大三个尺寸。比如输入是[1,3,224,224]到[1,3,1024,1024]，最优尺寸设为[1,3,512,512]。
- 构建engine：TensorRT会根据最优尺寸做层融合和内存分配，同时保证其他尺寸也能运行。
- 运行时绑定：每次推理前，调用
setBindingDimensions设置当前输入尺寸。
// 动态形状配置示例
IBuilderConfig* config = builder->createBuilderConfig();
IOptimizationProfile* profile = builder->createOptimizationProfile();
profile->setDimensions(inputName, OptProfileSelector::kMIN, Dims4(1, 3, 224, 224));
profile->setDimensions(inputName, OptProfileSelector::kOPT, Dims4(1, 3, 512, 512));
profile->setDimensions(inputName, OptProfileSelector::kMAX, Dims4(1, 3, 1024, 1024));
config->addOptimizationProfile(profile);
注意：动态形状会增加显存占用，因为TensorRT要为所有可能的尺寸预留空间。我建议最优尺寸选你最常用的那个，这样大部分场景都能跑在最佳性能上。
4.4 插件编写：当内置算子不够用时
TensorRT内置了上百个算子，但总有一些奇奇怪怪的操作它不支持。这时候就得自己写插件了。我写过不下20个插件，从简单的激活函数到复杂的ROI Align都有。
写插件其实不复杂，核心就三步：
- 继承IPluginV2DynamicExt：这是TensorRT推荐的插件基类，支持动态形状。
- 实现核心方法：
getOutputDimensions、enqueue、serialize、deserialize。其中enqueue是真正干活的地方，你在这里写CUDA kernel。 - 注册插件：用
REGISTER_TENSORRT_PLUGIN宏注册，这样TensorRT就能识别你的插件了。
// 插件注册示例
class MyCustomPlugin : public IPluginV2DynamicExt {
public:
// ... 实现所有虚函数
int32_t enqueue(const PluginTensorDesc* inputDesc,
const PluginTensorDesc* outputDesc,
const void* const* inputs,
void* const* outputs,
void* workspace,
cudaStream_t stream) noexcept override {
// 在这里调用你的CUDA kernel
myKernel<<<grid, block, 0, stream>>>(inputs[0], outputs[0], ...);
return 0;
}
};
REGISTER_TENSORRT_PLUGIN(MyCustomPluginCreator);
我曾经写过一个NMS插件，一开始性能很差。后来发现是没做warp-level的并行优化。改完之后，速度提升了4倍。嗯，写插件这事，CUDA功底很重要。
避坑指南：插件写完后一定要做序列化/反序列化测试。我遇到过插件在内存里跑得好好的，一保存再加载就崩了。原因是serialize方法没把内部状态写全。
4.5 知识体系总览
这四个机制不是孤立的，它们共同构成了TensorRT的加速引擎。我画了一张图，帮你理清它们之间的关系：
你看，层融合和精度校准是「编译期优化」，动态形状和插件编写是「运行时能力」。四者缺一不可。我个人觉得，理解透这四点，TensorRT就算入门了。剩下的就是多踩坑、多积累经验。
想要更多资料，联系可可