3、TensorRT基础：模型转换流程（ONNX→TensorRT）、Engine构建、序列化与反序列化、上下文（Context）与执行模式

TensorRT 的核心价值在于将训练框架（PyTorch、TensorFlow 等）产出的模型，通过一系列优化手段转换为可在 NVIDIA GPU 上高效执行的推理引擎。理解从 ONNX 到 TensorRT Engine 的完整转换链路、Engine 的构建与生命周期管理、序列化与反序列化的适用场景，以及 Context 与执行模式的底层机制，是进行 CUDA Stream 并行推理优化的前提。本章将系统拆解这一基础链路，并重点剖析每个环节的工程陷阱与性能关键点。

ONNX 模型 计算图 + 权重 动态/静态 Shape Builder Network + Config 精度校准 / 层融合 Engine 优化后推理计划 含 Kernel 选择 Context 执行状态 可多实例 序列化 Engine → .plan 文件 反序列化 .plan → Engine 执行模式 同步 / 异步 / Stream 关键路径： ONNX 解析 → Builder 构建（含层融合/精度校准）→ 生成 Engine → 序列化保存 .plan 运行时：加载 .plan → 反序列化 → 创建 Context → 绑定输入输出 → enqueueV3 执行

3.1 模型转换流程：从 ONNX 到 TensorRT 的完整链路

核心原理

ONNX（Open Neural Network Exchange）作为中间表示格式，将 PyTorch、TensorFlow 等框架的模型统一为计算图描述。TensorRT 的转换流程本质上是将 ONNX 计算图映射到 TensorRT 的优化层体系（Layer）中，并经过图优化、内核自动调优、精度校准等步骤生成推理引擎。转换过程并非简单的格式翻译，而是涉及算子映射、图融合（如 Conv+BN+ReLU 融合为 CBR）、张量布局调整（NHWC 与 NCHW 的自动选择）、以及 Kernel 级别的自动调优。

操作流程

  1. 导出 ONNX：在训练框架中导出模型时，需固定输入维度（或设置动态轴），并确保算子集合在 TensorRT 支持范围内。PyTorch 使用 torch.onnx.export，需设置 opset_version（建议 ≥ 11）和 dynamic_axes。
  2. 创建 Builder 与 Network：通过 nvinfer1::createInferBuilder(logger) 创建 Builder，再调用 builder->createNetworkV2(0) 创建 Network 定义。
  3. 解析 ONNX：使用 nvonnxparser::createParser 将 ONNX 文件解析并填充到 Network 中。解析器会报告不支持的算子或属性。
  4. 配置 BuilderConfig：设置工作空间大小（setMemoryPoolLimit）、精度模式（FP32/FP16/INT8）、动态 Shape 优化配置等。
  5. 构建 Engine：调用 builder->buildSerializedNetwork(network, config) 生成序列化引擎。
关键参数说明：工作空间（Workspace）大小直接影响层融合策略。较大的工作空间允许 TensorRT 尝试更多融合方案，但会增加显存占用。建议初始设置为 1GB，根据实际显存余量调整。

案例拆解：ResNet50 转换中的算子映射问题

某项目需要将 ResNet50 从 PyTorch 转换到 TensorRT，初始转换时遇到 aten::max_pool2d 算子不支持的错误。分析发现 PyTorch 导出的 ONNX 中 MaxPool 的 ceil_mode 属性为 True，而 TensorRT 的 MaxPool 层默认不支持该模式。解决方案是在导出时设置 ceil_mode=False，或使用 onnx-simplifier 对计算图进行简化。转换后对比精度，FP32 模式下 Top-1 精度损失小于 0.1%，满足要求。

常见误区与复盘

  • 误区一：认为 ONNX 转换是“万能”的。实际上，TensorRT 对某些动态控制流（如 torch.where 的复杂条件）支持有限，需提前规避。
  • 误区二：忽略 opset_version 的影响。低版本 opset 可能导致部分算子无法映射。
  • 复盘：转换前应使用 polygraphy 工具进行模型结构预检，提前发现不支持的算子，避免在构建阶段才报错。

3.2 Engine 构建：生命周期与优化策略

核心原理

Engine 是 TensorRT 优化后的推理计划，包含具体的 Kernel 选择、内存分配策略、层融合结果等。Engine 的构建过程是离线完成的，构建时间通常较长（从几秒到几分钟不等），但构建完成后可序列化保存，运行时直接加载。Engine 的生命周期管理包括创建、配置、构建、序列化、反序列化和销毁。

构建配置对比表

配置项 FP32 模式 FP16 模式 INT8 模式
精度设置 默认 setFlag(kFP16) setFlag(kINT8) + 校准
性能提升 基线 约 1.5-2x 约 3-4x
精度风险 无 低（需验证） 中高（需校准）
适用场景 精度敏感 通用推理 大规模部署

操作流程

// 创建 Builder
auto builder = nvinfer1::createInferBuilder(logger);
// 创建 Network
auto network = builder->createNetworkV2(0);
// 创建配置
auto config = builder->createBuilderConfig();
// 设置工作空间
config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 30);
// 启用 FP16
config->setFlag(nvinfer1::BuilderFlag::kFP16);
// 构建序列化引擎
auto serializedModel = builder->buildSerializedNetwork(*network, *config);

案例拆解：动态 Shape 的 Engine 构建

在 NLP 场景中，输入序列长度不固定。构建 Engine 时需设置优化 Profile（Optimization Profile），指定最小、最优、最大维度。例如，设置 min=(1,32)、opt=(1,128)、max=(1,512)。构建时 TensorRT 会针对最优维度进行 Kernel 调优，运行时根据实际输入维度选择最接近的 Kernel。实测中，动态 Shape 相比静态 Shape 性能损失约 5%-10%，但灵活性大幅提升。

风险提示：动态 Shape 的 Engine 构建时间显著增加，且显存占用更高。若业务场景输入维度相对固定，建议使用静态 Shape 以获得最佳性能。

3.3 序列化与反序列化：Engine 的持久化与加载

核心原理

序列化是将 Engine 转换为字节流（.plan 文件）的过程，反序列化则是将字节流恢复为 Engine 对象。序列化的意义在于避免每次启动都重新构建 Engine，节省启动时间。但需注意，序列化后的 Engine 与 GPU 架构、TensorRT 版本、CUDA 版本强相关，跨环境使用可能导致加载失败。

操作流程

  1. 序列化：buildSerializedNetwork 返回 IHostMemory，将其写入文件即可。
  2. 反序列化：读取 .plan 文件，调用 runtime->deserializeCudaEngine(data, size) 创建 Engine。
  3. 版本兼容性检查：加载时需确保 TensorRT 主版本一致，次版本向后兼容。

案例拆解：跨环境部署的序列化陷阱

某团队在开发机（RTX 3090，TensorRT 8.5）上构建 Engine，部署到生产服务器（A10，TensorRT 8.4）时反序列化失败。原因是 TensorRT 8.5 生成的 Engine 包含 8.5 特有的 Kernel 特性，8.4 无法解析。解决方案是在生产环境重新构建 Engine，或统一 TensorRT 版本。复盘经验：序列化 Engine 应视为“环境绑定”产物，部署时需在目标环境重新构建或严格匹配版本。

常见误区

  • 误区一：将 .plan 文件视为跨平台通用文件。实际上，不同 GPU 架构（如 Ampere 与 Turing）的 Engine 不兼容。
  • 误区二：忽略序列化文件的大小。大型模型的 .plan 文件可能超过 500MB，需考虑存储和加载时间。

3.4 上下文（Context）与执行模式

核心原理

Context（执行上下文）是 Engine 的运行时实例，持有推理所需的中间状态、内存绑定和 CUDA Stream 信息。一个 Engine 可以创建多个 Context，每个 Context 独立执行推理任务，共享 Engine 的权重和 Kernel 计划。执行模式分为同步执行（enqueueV3 阻塞式）和异步执行（配合 CUDA Stream 的非阻塞式）。

执行模式对比表

执行模式 API 阻塞行为 适用场景
同步执行 context->executeV3(stream) 阻塞 CPU 直到推理完成 简单推理、低并发
异步执行 context->enqueueV3(stream) 立即返回，GPU 异步执行 高并发、流水线

操作流程

// 创建 Context
auto context = engine->createExecutionContext();
// 设置输入输出地址
context->setTensorAddress("input", inputBuffer);
context->setTensorAddress("output", outputBuffer);
// 异步执行
cudaStream_t stream;
cudaStreamCreate(&stream);
context->enqueueV3(stream);
// 同步等待
cudaStreamSynchronize(stream);

案例拆解：多 Context 并发推理

在视频分析场景中，需要同时处理 4 路视频流。创建 4 个 Context，每个 Context 绑定独立的 CUDA Stream 和输入输出缓冲区。通过 enqueueV3 异步提交到各自的 Stream，GPU 可并行调度多个 Kernel。实测中，4 路并发相比串行推理吞吐量提升 3.2 倍，但显存占用增加约 30%（每个 Context 需要独立的中间缓冲区）。

经验提醒：Context 的数量并非越多越好。当 Context 数量超过 GPU 的 SM 数量时，性能提升趋于饱和，且显存开销线性增长。建议根据 GPU 规格和显存余量动态调整并发数。

常见误区与总结

  • 误区一：认为一个 Engine 只能创建一个 Context。实际上，多 Context 是提升并发吞吐的关键手段。
  • 误区二：忽略 Stream 的同步。异步执行后必须调用 cudaStreamSynchronize 或使用事件机制确保数据就绪。
  • 总结：Context 与 Stream 的配合是 CUDA Stream 并行推理的核心。理解 Engine 的“只读”特性与 Context 的“独立状态”特性，是设计高效推理流水线的基础。
← 上一章 📖 返回目录 下一章 →