3. TensorRT核心概念与工作流程

TensorRT 并不是一个“训练框架”，而是一个面向 NVIDIA GPU 的推理编译器与运行时。它接收来自 PyTorch、TensorFlow、ONNX 等来源的已训练模型，把网络重新组织成一张经过优化的执行图，再编译成与具体 GPU 架构、精度模式、批大小强绑定的引擎（Engine）。理解 TensorRT，本质上就是理解这条“模型 → 优化图 → 序列化引擎 → 反序列化 → 运行时执行”的链路，以及每一环上真正发生了什么。

一、计算图优化：从原始网络到可执行图

导入模型后，TensorRT 首先构建一张内部计算图。此时它做的是“图级”工作，而不是逐层照搬：

  • 常量折叠：把权重与常量参与的计算提前算掉，例如 BN 的均值方差在推理阶段是常量，可直接并入卷积权重。
  • 死代码消除：训练用的 Dropout、辅助分支、未参与输出的节点被移除。
  • 冗余转置与重塑消除：连续的 Transpose、Reshape、Squeeze 若可合并或抵消，会被折叠。
  • 布局与格式选择：TensorRT 会为张量选择 NCHW、NHWC、NC/32HW32 等格式，以匹配底层 Tensor Core 的访问模式。

这些优化的目标是减少 kernel 启动次数与显存往返。判断优化是否生效，最直接的方式是查看构建日志中的层数与显存占用，或使用 trtexec --verbose 观察每层被融合、被替换的情况。

二、层融合：减少访存才是关键

GPU 推理的瓶颈往往不是算力，而是显存带宽。层融合把多个算子合并成一个 kernel，使中间结果留在寄存器或共享内存中，不再写回显存。常见融合模式包括：

  • Conv + Bias + ReLU（或 SiLU、GELU）→ 单个卷积 kernel；
  • Conv + BN → 权重与偏置重算后合并；
  • 多个逐元素算子（Add、Mul、Scale）→ 一个 element-wise kernel；
  • 矩阵乘 + 转置 + 激活 → 一个 GEMM 变体。

融合不是越多越好。若某层输出需要被多个下游节点共享，或需要作为网络输出暴露，融合会被打断。因此同一模型在不同输出配置下，融合结果可能不同。

TensorRT 推理链路：图优化 → 融合 → 校准 → 序列化 → 运行时 原始计算图（ONNX） Conv BN ReLU Add Mul 图优化 融合后执行图 Conv+BN+ReLU 单 kernel Add+Mul element-wise 精度校准 精度模式 FP32 / FP16 直接构建 INT8 校准 校准集 → 缩放因子 序列化 引擎与运行时 .plan 引擎文件 与 GPU/精度绑定 反序列化 加载 → 执行上下文 运行时执行（IExecutionContext） 输入张量 H2D 拷贝 enqueueV3 异步提交 GPU 执行 CUDA Stream 输出张量 D2H 拷贝 绑定张量地址 → 提交推理 → 同步流 → 读取结果 同一引擎可创建多个执行上下文，共享权重、独立状态

三、精度校准：INT8 的缩放因子从哪来

FP16 通常只需在构建时设置标志即可，几乎不损失精度。INT8 则不同：它需要为每个张量确定一个缩放因子 scale，把浮点值映射到 [-127, 127]。TensorRT 提供两种方式：

  • 训练后校准（PTQ）：准备 200～1000 张有代表性的校准图，用 IInt8EntropyCalibrator2 跑一遍前向，统计激活分布，生成校准缓存。
  • 量化感知训练（QAT）：在训练阶段插入伪量化节点，精度通常更好，但需要重新训练。

校准集必须覆盖真实部署场景的输入分布。若校准集全是白天图像，夜间输入可能出现明显精度下降。校准完成后应保存 CalibrationCache，避免每次构建重复校准。

四、序列化与反序列化

构建引擎是耗时操作，可能几十秒到几分钟。因此 TensorRT 允许把引擎序列化为 .plan 文件：

IBuilderConfig* config = builder->createBuilderConfig();
config->setMemoryPoolLimit(MemoryPoolType::kWORKSPACE, 1ULL << 30);
config->setFlag(BuilderFlag::kFP16);

IHostMemory* serialized = builder->buildSerializedNetwork(*network, *config);
std::ofstream out("model.plan", std::ios::binary);
out.write((char*)serialized->data(), serialized->size());

部署时只需反序列化，不再需要 ONNX 解析器和构建器：

IRuntime* runtime = createInferRuntime(logger);
std::ifstream in("model.plan", std::ios::binary);
std::string blob((std::istreambuf_iterator<char>(in)), {});
ICudaEngine* engine = runtime->deserializeCudaEngine(blob.data(), blob.size());

关键约束：.plan 与构建时的 GPU 架构、TensorRT 版本、精度模式、最大批大小绑定。换卡或升级 TensorRT 后通常需要重新构建。

五、运行时执行

引擎本身不保存执行状态。真正执行推理的是 IExecutionContext，它持有中间激活的显存。典型流程是：

  1. 创建执行上下文，查询输入输出张量名与形状；
  2. 为输入输出分配 GPU 显存，调用 setTensorAddress 绑定地址；
  3. 把主机数据异步拷贝到输入显存；
  4. 调用 enqueueV3(stream) 提交推理；
  5. 在同一个流上异步拷贝输出回主机，最后 cudaStreamSynchronize。

若需要并发处理多路请求，可以为一个引擎创建多个执行上下文，它们共享权重，但各自拥有独立的中间张量。此时应关注显存占用：每个上下文都会额外占用一份激活显存。

构建阶段关注“图优化与精度”，运行阶段关注“显存、流与并发”。把这两件事分开排查，定位问题会快很多。

六、一个可复用的构建与执行检查清单

阶段关键动作判断依据
解析ONNX → network，检查不支持的算子解析日志无 warning，层数合理
优化设置 workspace、精度标志、动态形状 profile构建成功，显存未超限
校准INT8 时提供校准集并缓存校准缓存生成，精度回归通过
序列化保存 .plan，记录 GPU 与版本文件可被反序列化
执行绑定地址、异步 enqueue、同步流输出与参考实现误差在阈值内

把这条链路走通一次，后续所有模型部署都只是替换解析器和输入输出配置。真正需要反复打磨的，是校准集质量、动态形状 profile 的取值，以及执行上下文的并发与显存平衡。

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