4、TensorRT模型构建：Builder配置、网络定义、优化策略

说实话，模型训练出来只是第一步。真正要把模型用起来，还得靠TensorRT这把"手术刀"来精雕细琢。今天咱们就聊聊怎么用Builder这把刀，把模型切得又快又准。

我个人习惯把TensorRT的构建过程分成三块：Builder配置、网络定义、优化策略。这三块搞定了，模型部署基本就稳了。

4.1 Builder配置：你得告诉引擎要干啥

Builder就像是一个工厂的厂长。你得告诉它：你要造什么产品？用什么工艺？预算多少？

先看一段最基础的配置代码：

import tensorrt as trt

# 创建Builder
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(
    1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
)

# 创建配置
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB
config.set_flag(trt.BuilderFlag.FP16)  # 开启FP16

这里有个坑，我踩过好几次。EXPLICIT_BATCH这个标志，在TensorRT 8.x之后是必须加的。不加的话，动态batch就用不了。你想想看，线上流量忽高忽低，没有动态batch怎么扛得住？

4.2 网络定义：把模型"翻译"给TensorRT听

网络定义说白了，就是把PyTorch或者ONNX的模型，转成TensorRT能理解的结构。我一般推荐用ONNX作为中间格式，兼容性最好。

但这里有个问题——ONNX导出时经常丢东西。比如一些自定义算子，ONNX根本不认识。怎么办？

嗯，这时候就得手动补了。TensorRT提供了两种方式：

1. Plugin插件：自己写CUDA代码实现算子
2. IFunc：用TensorRT内置的算子组合拼出来

我个人更推荐第二种。为什么呢？因为写Plugin太容易出bug了。我曾经为了一个LayerNorm插件，调了整整两天，最后发现是shared memory没对齐。你说冤不冤？

4.3 优化策略：精度、内存、延迟的三角博弈

做部署优化，本质上就是在精度、内存、延迟这三个维度上找平衡。你不可能三个都要，得学会取舍。

4.3.1 精度优化：FP16和INT8怎么选？

FP16基本是无损的。我测过几十个模型，FP16相比FP32，精度损失基本在0.1%以内。但INT8就不一定了，有些模型掉点很厉害。

怎么判断能不能用INT8？我的经验是：先跑一遍校准集，看看激活值的分布。如果分布比较集中，INT8问题不大。如果分布很散，那就要小心了。

# 开启INT8校准
config.set_flag(trt.BuilderFlag.INT8)
config.int8_calibrator = MyCalibrator()  # 自定义校准器

4.3.2 内存优化：别让显存爆了

内存优化这块，我一般关注三个点：

• 工作空间大小：设太大浪费，设太小性能差。我习惯先设1GB，然后根据实际使用情况调整。
• 张量复用：TensorRT会自动做内存复用，但有些场景需要手动干预。
• 序列化：把构建好的引擎存成文件，下次直接加载，省去重新构建的内存开销。

说到序列化，有个细节要注意：引擎文件是跟CUDA版本、TensorRT版本绑定的。换版本就得重新构建。别问我怎么知道的，都是泪。

4.3.3 延迟优化：把每一微秒都榨出来

延迟优化是重头戏。我常用的手段有这些：

优化手段	效果	适用场景
算子融合	减少kernel启动开销	所有模型
动态batch	提高吞吐	在线服务
CUDA图	减少CPU调度开销	固定输入尺寸
多流并行	提高GPU利用率	多batch场景

这里我想重点说说算子融合。TensorRT会自动把Conv+BN+ReLU这种常见组合融合成一个kernel。但有些特殊组合它不会自动做，比如Conv+LayerNorm。这时候就得手动写Plugin来实现了。

4.4 构建流程全景图

说了这么多，咱们用一张图把整个流程串起来：

这张图把整个流程串起来了。从Builder创建，到网络定义，再到配置优化，最后构建引擎。每一步都有它的讲究。

最后说一句，构建引擎是个体力活，也是个经验活。多跑几次，多看看profiling结果，慢慢就能找到感觉。别指望一次就能配出最优解，那不太现实。

---

公众号:蓝海资料掘金营，微信deep3321