小白用 TensorRT 给模型加速：ONNX 转 Engine 踩坑实录

• TensorRT 是什么？为什么需要它？
• 一个意外：TensorRT 11.x 的 API 大改
• 准备环境：Windows 上装 TensorRT
• 核心概念：三层架构
• 精度模式对比
• Step 1：ONNX 转 FP32 Engine（入门）
  • 核心代码
  • 推理时的坑：get_tensor_mode 参数类型
• Step 2：GPU 推理和 CUDA 内存管理
• Step 3：ModelOpt FP16 转换
  • ModelOpt 自动转换
  • 效果对比
• Step 4：多 batch 基准测试（GPU 到底快在哪？）
  • 关键发现
  • 这个结果说明了什么？
• 算子融合：TensorRT 的核心优化手段
• 文件说明与实践总结
• 踩坑记录
  • OptimizationProfile 忘记设置
  • get_tensor_mode 传入整数
  • FP16 Engine 比 FP32 还慢
  • cuMemcpyHtoD 的地址参数
  • nvidia-modelopt 的安装源
• 总结
• 系列文章

这是 AI 实操第 5 天。第 1 天训练了模型、第 2 天导出了 ONNX。今天轮到 TensorRT：把 ONNX 编译成 Engine，在 GPU 上跑。 本文首发地址 https://h89.cn/archives/671.html 项目地址 https://gitee.com/chenjim/cockpit-ai-from-zero

TensorRT 是什么？为什么需要它？

前几天的流程走下来，模型已经能在 ONNX Runtime 上跑了。但 ONNX Runtime 是个通用引擎，你可以在 CPU 上跑也可以在 GPU 上跑，但它不会对你的模型做特别激进的优化。

TensorRT（Tensor RunTime，张量运行时）是 NVIDIA 专门为 GPU 做的推理加速库。 你给它一个模型，它花时间编译成针对你这块 GPU 最优化的一份执行计划，这个过程叫 Engine（引擎）构建。

打个比方：ONNX Runtime 像一本通用菜谱，谁都能照着做。TensorRT 像一位名厨，先研究你的菜谱，重新安排工序，把能同时做的合并、没必要的步骤砍掉，最后给你一份只有他自己能执行的优化手册。这份手册就是 TensorRT Engine。

TensorRT 的优化从三个方向入手：

• 算子融合（Operator Fusion）：把多个连续的小算子合并成一个 kernel（GPU 上的最小计算单元），减少启动开销。
• 精度校准（Precision Calibration）：把模型从 FP32 降到 FP16 甚至 INT8，减少计算量和显存占用。
• 自动调优（Auto-tuning）：对同一个算子尝试多种实现，挑出针对你显卡最快的那个。

一个意外：TensorRT 11.x 的 API 大改

写这篇文章的时候，我装的是 TensorRT 11.1.0.106。如果你看网上搜到的教程，可能会发现它们写的是这样的代码：

config.set_flag(trt.BuilderFlag.FP16)  # 老版本写法

但这个写法在 TensorRT 11.x 会直接报错：BuilderFlag 里没有 FP16 这个属性了。因为 11.x 直接把精度相关的 builder flag 给移除了。

新版本的逻辑变成了：你要先把 ONNX 模型用 NVIDIA 的 ModelOpt（Model Optimization，模型优化工具）预处理成 FP16 或 INT8 的 ONNX，然后把这份已经标好精度的 ONNX 交给 TensorRT。TensorRT 不再负责决定精度，它只负责执行模型里已经写好的精度。

也就是说，TensorRT 11.x 的转换流程从原来的"一步走"变成了"两步走"：

老流程：ONNX → TensorRT（内部决定精度）→ Engine
新流程：ONNX → ModelOpt（标好精度）→ 精度ONNX → TensorRT（只负责编译）→ Engine

这个变化对新手来说增加了复杂度，但也更清晰了：精度转换和 Engine 编译是两件不同的事。

TensorRT 11.x 的其他变化：

• 强类型网络（Strongly Typed Network）：精度来自 ONNX 模型本身，不再由 Builder flag 控制。
• EXPLICIT_BATCH 默认启用：create_network() 不用传参了。
• 动态 batch 必须设 OptimizationProfile：遇到动态形状必须告诉 TensorRT 最小/最优/最大尺寸。

准备环境：Windows 上装 TensorRT

我们这个项目原本是在 Linux 上跑的，但这次换到有 NVIDIA RTX 显卡的 Windows 上，所以选择了直接在 Python 环境里安装 TensorRT。

为什么换平台？TensorRT 是 NVIDIA GPU 专用加速库，前 4 天用的开发机没有独立显卡，而手边这台 Windows 笔记本带 RTX 4060，正好用来体验 GPU 推理和 FP16 加速。Linux 和 Windows 的 TensorRT Python API 几乎一样，你完全可以把下面的代码搬回 Linux 跑。

# 在项目的 .venv 虚拟环境里安装
pip install tensorrt
pip install cuda-python       # CUDA 驱动 API 的 Python 封装

这里有个坑：如果你需要做 FP16/INT8 精度转换，还需要安装 nvidia-modelopt，但它在默认的 PyPI（Python 包索引）上找不到，要从 NVIDIA 的源安装：

pip install nvidia-modelopt polygraphy onnx-graphsurgeon onnxslim lief \
  --extra-index-url https://pypi.nvidia.com

装完后验证一下：

python -c "import tensorrt as trt; print(trt.__version__)"
# 输出 11.1.0.106

核心概念：三层架构

写代码前先看 TensorRT 的三层架构：

• IBuilder（构建器）：负责配置构建参数，比如显存上限、优化配置、精度选项。相当于总工头。
• INetworkDefinition（网络定义）：存放神经网络的结构，你可以手搭每一层，也可以通过 ONNX Parser 自动导入。相当于蓝图。
• ICudaEngine（引擎）：最终产物，包含优化后的可执行计划。可以序列化保存到磁盘，也可以反序列化加载。相当于已经排好工序的施工方案。

转换路径也很直接：

ONNX → ONNXParser → INetworkDefinition → IBuilder → ICudaEngine

精度模式对比

TensorRT 支持三种精度模式，简单对比一下：

FP16 的精度损失在大多数推理场景下可以忽略不计，INT8 则需要额外的校准步骤（用一小批真实数据跑一遍，确定量化范围）。今天我们会做 FP32 和 FP16，INT8 留作进阶。

Step 1：ONNX 转 FP32 Engine（入门）

先从最基础的做起——不降精度，直接把 ONNX 转成 FP32 的 Engine。

核心代码

import tensorrt as trt

# 1. Logger：控制 TensorRT 打印什么级别的日志
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)

# 2. Builder：构建引擎的总控制台
builder = trt.Builder(TRT_LOGGER)

# 3. Network：定义网络结构
network = builder.create_network()

# 4. ONNX Parser：把 ONNX 文件翻译成 Network 能理解的结构
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("mnist_cnn.onnx", "rb") as f:
    if not parser.parse(f.read()):
        for i in range(parser.num_errors):
            print(f"解析错误 {i}: {parser.get_error(i)}")
        raise RuntimeError("ONNX 解析失败")

# 5. Config：配置构建参数
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 1GB

# 6. OptimizationProfile：处理动态形状
profile = builder.create_optimization_profile()
profile.set_shape("input", (1, 1, 28, 28), (32, 1, 28, 28), (256, 1, 28, 28))
config.add_optimization_profile(profile)

# 7. 构建 Engine
serialized = builder.build_serialized_network(network, config)
engine_data = bytes(serialized)

# 8. 保存到文件
with open("mnist_cnn_fp32.trt", "wb") as f:
    f.write(engine_data)

Logger 控制日志输出级别，WARNING 只显示警告和错误。Builder 是引擎构建的入口，Network 是网络结构的容器，Parser 负责解析 ONNX。

set_memory_pool_limit 设置 TensorRT 构建 Engine 时最多用多少显存。1 << 30 是 1GB，对于 MNIST 这种小模型完全够用。

OptimizationProfile（优化配置） 是我踩的第一个坑。第 2 天导出 ONNX 时用了 dynamic_axes，所以输入形状是 (-1, 1, 28, 28) —— batch 维度是动态的。TensorRT 遇到动态输入必须告诉它三个值：min（最小的输入大小）、opt（最常见的输入大小）、max（最大的输入大小）。TensorRT 会针对 opt 形状做最大优化，同时保证 min 到 max 范围内的形状都能正常跑。如果没设 OptimizationProfile，TensorRT 直接报错：

Network has dynamic or shape inputs, but no optimization profile has been defined.

build_serialized_network 返回的是 IHostMemory 对象，需要用 bytes() 转成普通字节才能写文件。

推理时的坑：get_tensor_mode 参数类型

写完 Engine 构建，我写推理部分时又踩了一脚。获取输入输出信息时，get_tensor_mode 接受的是 tensor 名称（字符串），不是索引：

# 错误写法：
engine.get_tensor_mode(0)  # TypeError

# 正确写法：
for i in range(engine.num_io_tensors):
    name = engine.get_tensor_name(i)
    if engine.get_tensor_mode(name) == trt.TensorIOMode.OUTPUT:
        output_name = name
        break

TensorRT 11.x 的 API 统一把名称当主键，迁移代码时容易在这翻车。

Step 2：GPU 推理和 CUDA 内存管理

Engine 构建好了，下一步就是加载它做推理。但 TensorRT 的推理和 PyTorch 不一样——它不管理内存，你需要自己用 CUDA API 操作 GPU 数据。

from cuda.bindings import driver, runtime

# 初始化 CUDA
driver.cuInit(0)
runtime.cudaSetDevice(0)

# 反序列化 Engine
runtime = trt.Runtime(TRT_LOGGER)
with open("mnist_cnn_fp32.trt", "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())

context = engine.create_execution_context()
context.set_input_shape("input", (1, 1, 28, 28))

# 申请 GPU 显存
_, input_dev = driver.cuMemAlloc(input_np.nbytes)
_, output_dev = driver.cuMemAlloc(output_bytes)

# 把输入数据拷贝到 GPU
driver.cuMemcpyHtoD(input_dev, input_np.ctypes.data, input_np.nbytes)

# 告诉 Engine 输入输出在显存的位置
context.set_tensor_address("input", int(input_dev))
context.set_tensor_address(output_name, int(output_dev))

# 执行推理
context.execute_async_v3(0)    # 0 表示默认 CUDA 流
driver.cuCtxSynchronize()      # 等 GPU 算完

# 把结果拷回 CPU
driver.cuMemcpyDtoH(output_np.ctypes.data, output_dev, output_np.nbytes)

# 清理显存
driver.cuMemFree(input_dev)
driver.cuMemFree(output_dev)

每一步都是手动的：申请显存 → 拷入数据 → 设置地址 → 执行 → 同步 → 拷回 → 释放。PyTorch 里 tensor.cuda() 干的事，这里全得自己写。

Step 3：ModelOpt FP16 转换

当 FP32 转换跑通后，可以尝试 FP16 了。如前所述，TensorRT 11.x 不支持在 Builder 里直接设 FP16 标志，需要用 ModelOpt 预处理 ONNX。

ModelOpt 自动转换

import modelopt.onnx.autocast as autocast
import onnx

converted = autocast.convert_to_mixed_precision(
    "mnist_cnn.onnx",
    low_precision_type="fp16",
    keep_io_types=True,  # 输入输出保持 FP32，只在中间层用 FP16
)
onnx.save(converted, "mnist_cnn_fp16_autocast.onnx")

参数 keep_io_types=True 的意思是模型的和输入输出用 FP32，中间的 Conv、ReLU 等层用 FP16。这样上层调用方不用改代码，底层的计算已经加速了。

ModelOpt 内部做了三件事：

1. 把 ONNX 的 opset version 从 11 升级到 13（onnxscript 自动处理）
2. 用 ONNX Runtime 跑一次推理，拿每层的输出当"参考值"
3. 逐层分析哪些算子可以安全转 FP16，哪些要保留 FP32

日志里会看到 Converted 17/17 nodes (100.00%) to fp16，说明 17 个算子全部成功转换。

转换完成后，这个 FP16 ONNX 直接喂给 TensorRT，走和 FP32 完全一样的构建流程。不需要在 BuilderConfig 里做任何特殊配置，因为精度信息已经写在 ONNX 文件里了。

效果对比

FP16 的 Engine 比 FP32 小了约 42%。FP16 的权重用 2 字节存一个浮点数，FP32 需要 4 字节，实际效果符合预期。

我用 10 张测试图分别用两个 Engine 跑了一遍，最大数值差异在 0.001 以内，预测结果完全一致。

Step 4：多 batch 基准测试（GPU 到底快在哪？）

下面是基准测试结果。对于 MNIST 这种小模型，batch 不够大时 GPU 反而更慢。

不同 batch size 下的完整对比，包括 ONNX Runtime（CPU）、TensorRT FP32、TensorRT FP16：

关键发现

• batch=1 时 GPU 比 CPU 慢 5 倍。 GPU 每次执行需要把 kernel 加载到计算单元，这个启动开销（kernel launch overhead）几十到几百微秒。当模型只有两层卷积时，GPU 大部分时间花在启动上。
• GPU 的延迟基本不随 batch 增长。 CPU 从 batch=1 到 256 延迟涨了 10 倍（0.07ms→0.68ms），GPU 只从 0.42ms 涨到 0.47ms。batch 越大 GPU 优势越明显。
• batch=128 时 GPU 开始反超。 同时处理 128 张图足够分摊启动开销。
• batch=256 时 GPU 快 1.5 倍。 和网上宣传的 5-10 倍差距很大——MNIST 的 CNN 太小了，两层卷积不到 10 万参数，GPU 利用率可能不到 5%。
• FP16 和 FP32 表现接近，FP16 甚至略慢。 不是 bug，是模型太小，格式转换的开销超过了计算节省。

这个结果说明了什么？

GPU 加速不是免费的。模型够大、batch 够大的时候才能发挥实力。跑 ResNet-50（2500 万参数）或 YOLOv8（300 万参数）时，GPU 利用率能到 80% 以上，加速比会好看很多。MNIST 这个 10 万参数的小模型，还没喂饱 GPU 就被 CPU 追上了。作为学习 TensorRT API 的入门案例，理解流程和 API 的目标是达到了。

算子融合：TensorRT 的核心优化手段

构建 Engine 时，TensorRT 还会做一件事——算子融合（Operator Fusion）。就是把多个连续的算子合并成一个 GPU kernel。比如 Conv2d → ReLU → MaxPool2d 这种模式，TensorRT 会把三者合并在一个 kernel 里执行。好处是少启动几次 kernel，省掉开销；数据直接在寄存器或 L1 缓存上传，不用反复读写显存。

垂直融合把前后串联的算子合并。Conv→ReLU→MaxPool 就是典型例子。更复杂的模型里还有 Conv→BatchNorm→ReLU，BN 的乘加运算可以折叠进 Conv 权重里，推理时 BN 层完全消失。

水平融合把并行且同类型的算子合并，比如多个 Add 合并成一个更大的 Add kernel。

TensorRT 常见的融合模式：

想亲眼确认哪些算子被融合了，可以用 trtexec --profilingVerbosity=detailed 或 config.set_profiling_verbosity(trt.ProfilingVerbosity.DETAILED)。不过 MNIST 这种小模型的日志可能很简略。

文件说明与实践总结

这一天的实操产生了几个关键脚本：

运行方式：

# 基础转换（FP32）
python convert_onnx_to_trt.py

# FP16 精度转换（需要提前安装 nvidia-modelopt）
python precision_autocast.py --precision fp16

# 批量基准测试
python benchmark.py

踩坑记录

这一路走下来踩了不少坑，总结成清单，方便你自查：

OptimizationProfile 忘记设置

导出 ONNX 时用了 dynamic_axes，TensorRT 认为输入是动态形状，必须通过 OptimizationProfile 告诉它 min/opt/max。否则报错：

Network has dynamic or shape inputs, but no optimization profile has been defined.

解决方法就是构建 Config 之后、调用 build_serialized_network 之前加上那段 profile 代码。

get_tensor_mode 传入整数

TensorRT 11.x 中，get_tensor_mode 接受的是 tensor 名称（字符串），不是索引。set_tensor_address 同理。这和旧版本 API 不同，迁移时要注意。

FP16 Engine 比 FP32 还慢

这不是 bug，是 MNIST 模型太小了。FP16 在 ResNet、BERT 那种规模才能发挥优势。MNIST CNN 两层卷积不到 10 万参数，FP16 的格式转换开销超过了计算节省。

cuMemcpyHtoD 的地址参数

cuMemcpyHtoD 的第二个参数是 host 内存的整数地址，需要通过 numpy 数组的 .ctypes.data 获取。这个地址必须指向连续内存，numpy 数组默认是连续的，但如果你做了切片或转置操作，记得先 .copy() 确保连续。

nvidia-modelopt 的安装源

nvidia-modelopt 不在默认的 PyPI 上，需要加 --extra-index-url https://pypi.nvidia.com，而且它有很多依赖（polygraphy、onnx-graphsurgeon、onnxslim、lief），最好一起装。

总结

今天做的事：

1. FP32 转换：用 TensorRT Python API 把 ONNX 转成 FP32 Engine，手写了 GPU 推理流程，理解了 Builder/Network/Engine 三层架构。
2. FP16 转换：用 ModelOpt 预处理 ONNX，Engine 从 0.47MB 降到 0.27MB，精度差异 0.001 以内。
3. 基准测试：batch=1 时 GPU 比 CPU 慢 5 倍，batch=256 时 GPU 快 1.5 倍。确认了 GPU 加速在小模型上的边界条件。

核心收获：API 用法和优化思路是通用的，学会这套流程，下次换大模型直接上手。

如果你搜到的教程代码跑不通，先检查 TensorRT 版本。trt.BuilderFlag.FP16 这种经典写法在 11.x 已经不存在了，精度转换改由 ModelOpt 负责。

第 6 天会回到高通平台，用 SNPE 把 ONNX 转成骁龙车机芯片能跑的 DLC 格式。

系列文章

本系列「车载端侧 AI 工程化从零上手」共 10 篇，建议按序阅读：

1. 零基础用 PyTorch 识别手写数字：MNIST 实战入门
2. 把 PyTorch 模型变成 ONNX：导出、验证和可视化一次学会
3. MLflow 实验追踪：从入门到上手
4. 用 MNN 在 CPU 上跑 AI 推理：阿里端侧推理框架上手记
5. 小白用 TensorRT 给模型加速：ONNX 转 Engine 踩坑实录
6. 零基础上手：高通 SNPE 模型转换实战（从 ONNX 到 DLC）
7. 零基础搞懂 SNPE 模型量化：INT8 精度损失 = 0% 的秘密
8. Python 工程化重构：从 print 脚本到 pytest 项目
9. 车载语音指令识别：从 44% 到 96% 的调优之路
10. 车载语音指令识别：SNPE 转换与真机部署实测记录