3、NVIDIA TensorRT（上）：架构原理、层融合与精度校准（FP16/INT8）、动态形状支持

聊到模型部署，TensorRT 是个绕不开的名字。我最早接触它是在 2018 年，当时一个检测模型在 Pytorch 上跑 30ms，转成 TensorRT 后直接掉到 8ms。说实话，那个数字让我愣了好几秒。从那以后，我养成了一个习惯：任何要上线的模型，第一件事就是看看能不能走 TensorRT。

这一章，咱们先啃上半部分。我会把 TensorRT 的架构原理、层融合、精度校准和动态形状这几个核心概念掰开揉碎。你想想看，搞懂了这些，后面调优的时候心里就有底了。

3.1 TensorRT 整体架构：它到底干了什么？

TensorRT 本质上是一个深度学习推理优化器。它不负责训练，只负责把你的训练好的模型，变成能在 NVIDIA GPU 上跑得飞快的推理引擎。

它的工作流程分两步：

  • 构建期（Build Time）：读入模型，做各种优化，生成一个优化后的计划文件（Plan File）。
  • 推理期（Runtime）：加载计划文件，执行推理。

我个人习惯把构建期想象成一个「编译器」。你给它源代码（模型），它给你生成可执行文件（Plan）。这个 Plan 是高度优化的，甚至针对你指定的 GPU 型号做了特化。

核心要点：TensorRT 的优化是离线完成的。这意味着构建期花的时间长一点没关系，只要推理期够快就行。

下面这张图，是我自己总结的 TensorRT 核心优化链路，你可以看看它到底对模型做了什么：

TensorRT 核心优化链路 输入模型 ONNX / UFF / Caffe 解析与图优化 消除无用节点 / 常量折叠 层融合（Layer Fusion） CBR / 拼接 / 横向融合 精度校准 FP32 → FP16 / INT8 内核自动调优 选择最优 CUDA Kernel 推理引擎（Plan）

3.2 层融合：把多个小操作合并成一个

层融合是 TensorRT 最核心的优化手段之一。说白了，就是把多个连续的小算子，合并成一个大的 Kernel 来执行。

为什么要这么做？因为 GPU 最怕「启动开销」。你每调用一个 CUDA Kernel，GPU 都要花时间做上下文切换、数据搬运。如果能把 5 个 Kernel 合并成 1 个，启动开销直接降到 1/5。

最常见的融合模式是 CBR（Conv + Bias + ReLU）。我见过太多新手把 Conv、Bias、ReLU 分开写，结果 TensorRT 在构建期默默把它们揉在一起了。你想想看，本来要启动三次 GPU，现在一次搞定，速度能不快吗？

我的经验：在项目中，我经常用 trtexec 工具加 --verbose 参数来查看融合日志。如果发现某些层没有被融合，通常是因为算子不支持或者图结构太复杂。这时候我会手动调整一下网络结构，比如把一些零散的激活函数合并进去。

除了 CBR，还有几种常见的融合模式：

  • 拼接融合（Concat Fusion）：多个输入拼接成一个张量时，TensorRT 会尝试把前面的操作直接写到目标位置，省去中间缓存。
  • 横向融合：比如两个并行的 1x1 Conv，可以合并成一个更大的 Conv 操作。
  • 残差融合：把 Add 操作和前面的 Conv 融合，减少一次内存读写。

3.3 精度校准：FP16 和 INT8 到底怎么选？

精度校准是 TensorRT 的另一个杀手锏。它允许你用更低的精度（FP16 或 INT8）来推理，从而获得更高的吞吐量。

FP16 相对简单。你只需要在构建时设置精度模式为 FP16，TensorRT 会自动把能转的层转成半精度。我实测下来，大部分模型从 FP32 切到 FP16，精度损失几乎可以忽略不计，但速度能提升 1.5 到 2 倍。

INT8 就复杂多了。它需要你提供一个校准数据集，TensorRT 会统计每一层激活值的分布，然后找到最优的量化参数。

注意：INT8 校准数据集的选择非常关键。我曾经在一个项目中，用训练集的 500 张图片做校准，结果上线后精度掉了 3 个点。后来换成验证集的 1000 张图片，精度只掉了 0.5 个点。校准数据一定要覆盖真实场景的分布。

下面这个表格，是我根据多个项目经验总结的精度选择建议：

精度模式 速度提升（相对 FP32） 精度损失 适用场景
FP32 1x 无 精度要求极高的场景（如医疗影像）
FP16 1.5x - 2x 几乎无 大多数通用场景
INT8 3x - 4x 0.5% - 2% 对延迟敏感、精度容忍度高的场景

校准的流程，我一般这样写代码：

import tensorrt as trt

# 创建构建器
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
config = builder.create_builder_config()

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

# 设置校准器
calibrator = MyCalibrator(calibration_files, batch_size)
config.int8_calibrator = calibrator

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

这里要注意，MyCalibrator 需要你自己实现，继承 trt.IInt8Calibrator。我一般用 IInt8EntropyCalibrator2，它在大多数场景下效果最好。

3.4 动态形状支持：告别固定尺寸

早期的 TensorRT 只支持固定输入尺寸。你构建时设成 224x224，推理时就必须是 224x224。这在实际项目中非常蛋疼——比如目标检测，输入图片大小经常不一样。

从 TensorRT 7 开始，引入了动态形状支持。你可以在构建时指定一个范围，比如 [1x3x224x224, 1x3x448x448]，推理时只要在这个范围内，TensorRT 都能处理。

实现动态形状，需要在构建时做三件事：

  1. 定义优化范围：指定最小、最优、最大三个尺寸。
  2. 设置优化配置：告诉 TensorRT 你期望的优化目标。
  3. 推理时绑定形状：每次推理前，设置实际的输入尺寸。

避坑指南：我曾经在项目中把优化范围设得太大（比如从 1x3x64x64 到 1x3x2048x2048），结果构建时间暴涨到 40 分钟，而且生成的引擎文件有 2GB。后来我把范围缩小到实际业务需要的区间，构建时间降到了 5 分钟。

动态形状的代码示例：

import tensorrt as trt

# 定义动态形状范围
min_shape = (1, 3, 224, 224)
opt_shape = (1, 3, 384, 384)
max_shape = (1, 3, 448, 448)

# 创建优化配置
profile = builder.create_optimization_profile()
profile.set_shape("input", min_shape, opt_shape, max_shape)
config.add_optimization_profile(profile)

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

# 推理时设置形状
context.set_binding_shape(0, (1, 3, 320, 320))

嗯，这里要注意一点：动态形状会带来额外的性能开销。因为 TensorRT 需要在运行时根据实际尺寸重新选择 Kernel。所以如果你的业务场景输入尺寸变化不大，我建议用固定尺寸，性能会更好。

好了，这一章的内容就到这里。TensorRT 的上半部分，我们聊了架构、层融合、精度校准和动态形状。这些是理解 TensorRT 的基石。下一章，我们会深入插件开发、多流推理和实际部署的坑。

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

← 上一章 📖 返回目录 下一章 →