PyTorch模型转TensorRT全链路实战：ONNX导出、引擎构建与精度调优-开发者社区

模型训练完之后，真正让人头疼的往往不是网络结构本身，而是怎么把它塞进推理引擎里跑出该有的速度。我最近刚把一个 PyTorch 模型从.pth一路折腾到 TensorRT 的.trt引擎，中间踩的坑比训练本身还多。这篇就把整条链路——PyTorch 导出 ONNX、ONNX 校验与简化、TensorRT 环境安装、trtexec 与 Python API 两条构建路径、精度对齐和性能调优——完整捋一遍。不管你是刚接触部署的新手，还是已经跑通过一次但被精度问题折磨过的老手，下面这些内容应该都能帮你少走几个晚上的弯路。

1. 先把整条链路想清楚：为什么是 PyTorch → ONNX → TensorRT

1.1 三个环节各自解决什么问题

很多人一上来就问“能不能直接从 PyTorch 转到 TensorRT”，答案是能，但通常不推荐作为第一步。原因在于这三者的定位完全不同。

PyTorch 是训练框架，它的计算图是动态的，带着大量 Python 层面的控制流、autograd 元信息、以及训练专用的算子（比如 Dropout 在训练和推理时行为不同）。TensorRT 是推理引擎，它只关心前向计算，而且要求计算图是静态的、算子集合是固定的。这两者之间需要一个“中间语言”来过渡，ONNX 就是这个中间语言。

ONNX 的价值在于它是一个开放的、框架无关的模型表示格式。PyTorch 官方提供了torch.onnx.export，把动态图 trace 或 script 成静态图；TensorRT 官方提供了 ONNX Parser，把 ONNX 图翻译成自己的网络定义。中间隔一层 ONNX，好处是：导出环节的问题和构建引擎环节的问题可以分开排查。如果直接走 PyTorch → TensorRT 的路径（比如通过torch-tensorrt），一旦出错，你很难判断是导出阶段的问题还是引擎构建阶段的问题。

所以这条链路的分工是：

PyTorch → ONNX：负责把训练图“冻结”成推理图，去掉训练专用逻辑，固定输入输出形状（或标记动态维度）。
ONNX → TensorRT：负责算子融合、精度校准（FP16/INT8）、kernel 自动调优，最终生成针对特定 GPU 架构优化的.trt引擎文件。

1.2 什么情况下值得走这条路

不是所有模型都值得上 TensorRT。我的经验判断标准是：

模型在 GPU 上推理，且对延迟敏感（比如实时视频处理、在线服务）；
模型结构相对稳定，不会一天改三次；
batch size 和输入分辨率基本固定，或者变化范围可控；
愿意花时间做精度对齐。

如果模型本身很小（比如几层全连接），或者推理频率极低，那用 ONNX Runtime 甚至原生 PyTorch 就够了，上 TensorRT 的收益可能还不如你调试花的时间。TensorRT 的收益主要来自三块：层融合（把 Conv+BN+ReLU 合成一个 kernel）、精度降低（FP32→FP16→INT8）、以及针对特定 GPU 的 kernel 自动调优。模型越大、层越多，收益越明显。

1.3 版本兼容性是这条链路最大的隐形杀手

在动手之前，我必须强调一件事：PyTorch、ONNX、onnxruntime、TensorRT、CUDA、cuDNN 这几个东西的版本是强耦合的。我见过太多人卡在“导出成功但 TensorRT 解析失败”或者“解析成功但输出全是 NaN”，最后发现是版本不匹配。

一个相对稳妥的组合思路是：先确定你的 GPU 驱动支持的 CUDA 版本上限，然后选一个 TensorRT 版本，再反推它支持的 ONNX opset 和 PyTorch 版本。TensorRT 每个版本对 ONNX opset 的支持是有上限的，比如某些版本最高支持到 opset 17，你导出时用了 opset 18 的新算子，Parser 直接报 unsupported。

提示：导出 ONNX 时不要盲目用最新 opset。先查清楚你目标 TensorRT 版本支持的 opset 范围，再在torch.onnx.export里用opset_version参数指定。

2. PyTorch 导出 ONNX：那些文档里不会写的细节

2.1 导出前的模型准备：eval 模式和设备选择

导出前第一件事，把模型切到 eval 模式：model.eval()。这一步看起来简单，但影响很大。BatchNorm 在 train 和 eval 模式下行为不同，Dropout 也是。如果你忘了切 eval，导出的 ONNX 里 BN 层会带着训练时的统计逻辑，推理结果直接偏掉。

第二件事是设备。我建议在 CPU 上导出，或者至少保证导出时用的 device 和后续推理一致。有些算子在不同 device 上的导出行为不一样，尤其是涉及自定义算子的时候。导出时给一个 dummy input，这个 input 的 shape 决定了 ONNX 图的输入 shape。

import torch model.eval() dummy_input = torch.randn(1, 3, 224, 224) torch.onnx.export( model, dummy_input, "model.onnx", opset_version=13, input_names=["input"], output_names=["output"], dynamic_axes={ "input": {0: "batch_size"}, "output": {0: "batch_size"} }, do_constant_folding=True )

这里几个参数值得展开说。do_constant_folding=True会在导出时把能算的常量提前算掉，减小图体积，一般建议开。dynamic_axes是动态维度的关键，如果你希望 batch size 可变，就必须在这里标记。但要注意，标记了动态维度之后，TensorRT 构建引擎时需要配置 optimization profile，否则默认按固定 shape 处理。

2.2 trace 和 script 的选择逻辑

torch.onnx.export默认走 trace 模式，也就是喂一个 dummy input 跑一遍，记录实际执行的操作。trace 的优点是简单，缺点是它记录的是“这一次执行路径”。如果你的模型里有if分支依赖输入值（比如根据输入内容走不同分支），trace 只会记录走的那一条，其他分支丢失。

遇到这种情况有两个选择：一是用torch.jit.script把模型转成 TorchScript 再导出，script 模式能保留控制流；二是改写模型，把数据依赖的控制流改成算子形式（比如用torch.where替代 if）。我的经验是，能改写就改写，因为 script 模式对某些 Python 语法支持有限，而且导出的图往往更复杂，TensorRT 解析时更容易出问题。

2.3 导出后的第一道校验：onnx.checker 和 onnxruntime 对齐

导出成功不代表导出正确。我养成的习惯是导出后立刻做两件事：

第一，用onnx.checker.check_model做结构校验，能抓出图结构层面的明显错误。

import onnx onnx_model = onnx.load("model.onnx") onnx.checker.check_model(onnx_model)

第二，用 onnxruntime 跑一遍，和 PyTorch 的输出做数值对比。这一步是精度对齐的第一道防线，如果这里就对不上，后面 TensorRT 更不可能对。

import onnxruntime as ort import numpy as np ort_session = ort.InferenceSession("model.onnx") ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()} ort_outputs = ort_session.run(None, ort_inputs) with torch.no_grad(): torch_outputs = model(dummy_input).numpy() diff = np.max(np.abs(ort_outputs[0] - torch_outputs)) print("max diff:", diff)

一般来说，FP32 下这个 diff 应该在 1e-5 到 1e-4 量级。如果 diff 到了 1e-2 甚至更大，说明导出过程有问题，常见原因是某些算子被替换成了近似实现，或者 eval 模式没切对。

2.4 onnx-simplifier：不是万能但常常有用

onnx-simplifier这个工具能做一些图优化，比如常量折叠、冗余算子消除、把一些复杂算子拆解成基础算子。它的价值在于：有些 PyTorch 导出的图里会带一些 TensorRT 不支持的算子组合，simplifier 有时能帮你化简掉。

pip install onnx-simplifier python -m onnxsim model.onnx model_sim.onnx

但要注意，simplifier 不是万能的，而且它本身也可能引入问题。我的做法是：先不简化，直接试 TensorRT 解析；如果报 unsupported operator，再试简化。简化后一定要重新做一次 onnxruntime 对齐，确认数值没变。

3. TensorRT 环境安装：比想象中更需要耐心

3.1 安装方式的取舍：pip 还是 tar 包

TensorRT 的安装方式主要有几种：pip 安装、tar 包解压、以及通过系统包管理器。我个人的偏好是 tar 包，原因是可控性强，能清楚知道每个库文件在哪，出问题好排查。pip 安装最省事，但有时候会和系统里已有的 CUDA 库产生冲突。

pip 安装大致是这样：

pip install tensorrt

但要注意，pip 上的 TensorRT 包通常只包含 Python 绑定和部分库，完整的推理库（比如libnvinfer.so）可能需要额外安装。而且 pip 包的版本和你的 CUDA 版本必须匹配，装错了会在 import 时报找不到符号。

tar 包方式则是下载对应 CUDA 版本的 TensorRT 压缩包，解压后把lib目录加到LD_LIBRARY_PATH，把python目录下的 whl 装进环境。这种方式的好处是版本对应关系一目了然。

3.2 环境变量配置：LD_LIBRARY_PATH 是重灾区

TensorRT 运行时依赖一堆动态库，如果LD_LIBRARY_PATH没配好，会出现ImportError: libnvinfer.so.X: cannot open shared object file。我的做法是在激活环境后显式设置：

export LD_LIBRARY_PATH=/path/to/tensorrt/lib:$LD_LIBRARY_PATH export PATH=/path/to/tensorrt/bin:$PATH

把bin也加进去是为了能用trtexec这个命令行工具，后面构建引擎会频繁用到它。

注意：LD_LIBRARY_PATH的顺序很重要。如果系统里同时存在多个版本的 CUDA 或 TensorRT 库，排在前面的会优先被加载。建议把 TensorRT 的路径放在最前面，避免加载到旧版本。

3.3 验证安装是否成功

装完之后别急着跑模型，先做最小验证：

import tensorrt as trt print(trt.__version__) logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) print("builder created:", builder is not None)

如果这几行能跑通，说明 Python 绑定和底层库都加载正常。如果import tensorrt就报错，八成是LD_LIBRARY_PATH或者 CUDA 版本的问题。这时候可以用ldd检查具体缺哪个库：

ldd $(python -c "import tensorrt; print(tensorrt.__file__)")

它会列出所有依赖库的解析情况，缺哪个一目了然。

3.4 一个容易被忽略的点：CUDA 版本与 TensorRT 的对应

TensorRT 每个版本都明确声明支持的 CUDA 版本范围。比如某个 TensorRT 版本可能只支持 CUDA 11.x，你系统里装的是 CUDA 12.x，那即使装上了，运行时也可能出各种诡异问题。查对应关系最靠谱的方式是看官方 release notes 里的 compatibility matrix，而不是凭感觉。

我踩过一次坑：系统 CUDA 是 12.1，装了个只支持到 11.8 的 TensorRT，import能过，但一构建引擎就 segfault。后来换成匹配版本，问题消失。所以这一步千万别省。

4. 用 trtexec 构建引擎：最快跑通第一条链路

4.1 trtexec 是什么，为什么先用它

trtexec是 TensorRT 自带的命令行工具，能把 ONNX 直接转成.trt引擎，还能顺便跑性能测试。我强烈建议第一次转换先用它，原因是：它把构建引擎的所有步骤都封装好了，你不需要写一行 Python 代码就能验证“这个 ONNX 到底能不能被 TensorRT 解析”。如果 trtexec 都过不了，那写 Python API 也是白搭。

基本用法：

trtexec --onnx=model.onnx --saveEngine=model.trt --fp16

这一行做了几件事：加载 ONNX、构建网络、应用 FP16 精度、序列化引擎到文件。如果 ONNX 里有 TensorRT 不支持的算子，这一步会直接报错并告诉你哪个算子不支持。

4.2 动态 shape 的 optimization profile 配置

如果你的 ONNX 标记了动态 batch，trtexec 需要你指定 shape 范围，否则它默认按固定 shape 处理。配置方式是通过--minShapes、--optShapes、--maxShapes：

trtexec --onnx=model.onnx \ --saveEngine=model.trt \ --fp16 \ --minShapes=input:1x3x224x224 \ --optShapes=input:8x3x224x224 \ --maxShapes=input:16x3x224x224

这里三个 shape 的含义要理解清楚：min和max定义了引擎支持的输入范围，opt是“最优”shape，TensorRT 会针对这个 shape 做 kernel 调优。所以opt应该设成你实际推理时最常用的 batch size。如果你实际跑 batch=8 居多，就把 opt 设成 8，这样性能最好。

提示：min、opt、max三个 shape 的维度数必须一致，而且opt必须落在min和max之间。设错了 trtexec 会直接报错。

4.3 精度选项：FP32、FP16、INT8 怎么选

--fp16是最常用的选项，绝大多数模型在 FP16 下精度损失可以忽略，速度提升明显。--int8则需要校准（calibration），因为 INT8 的量化范围需要根据实际数据分布确定。校准需要一个校准数据集，配置起来更复杂。

我的建议是：先用 FP32 跑通，确认精度没问题；再切 FP16，对比精度和速度；如果 FP16 还不够快，再考虑 INT8。不要一上来就 INT8，否则精度出问题你都不知道是量化的问题还是导出本身的问题。

4.4 解读 trtexec 的输出

trtexec 跑完会打印一堆信息，其中最有价值的是这几块：

Layer 信息：告诉你哪些层被融合了，哪些层用了什么精度。
Throughput：吞吐量，单位是 qps（queries per second）。
Latency：延迟，分 min/median/max/percentile。
GPU Compute Time：纯计算时间，不含数据传输。

我一般重点看 median latency 和 throughput。如果发现某个层没被融合（比如出现了独立的 BN 层），可能是导出时 BN 没被折叠，可以考虑在导出前手动 fold BN。

5. Python API 构建引擎：需要精细控制时的选择

5.1 什么时候必须用 Python API

trtexec 够用，但有些场景必须用 Python API：

需要在构建时动态决定精度策略（比如某些层 FP16，某些层 FP32）；
需要自定义插件（custom plugin）；
需要把引擎构建集成到服务里，而不是手动跑命令行；
需要更细粒度地控制 optimization profile。

Python API 的流程是：创建 logger → 创建 builder → 创建 network → 用 parser 解析 ONNX → 配置 builder config → 构建并序列化。

5.2 完整构建代码与逐行解释

import tensorrt as trt logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network( 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) ) parser = trt.OnnxParser(network, logger) with open("model.onnx", "rb") as f: if not parser.parse(f.read()): for i in range(parser.num_errors): print(parser.get_error(i)) raise RuntimeError("ONNX parse failed") config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) config.set_flag(trt.BuilderFlag.FP16) profile = builder.create_optimization_profile() profile.set_shape( "input", min=(1, 3, 224, 224), opt=(8, 3, 224, 224), max=(16, 3, 224, 224) ) config.add_optimization_profile(profile) serialized_engine = builder.build_serialized_network(network, config) with open("model.trt", "wb") as f: f.write(serialized_engine)

几个关键点展开说。EXPLICIT_BATCH这个 flag 现在基本是必须的，它表示网络使用显式 batch 维度。set_memory_pool_limit设置 workspace 大小，workspace 是 TensorRT 构建时用来做 kernel 调优的临时内存，太小会导致某些 tactic 无法使用，太大浪费显存。1GB 是个常见的起点，如果构建时报 workspace 不足，就往上加。

parser.parse返回 False 时，一定要遍历parser.num_errors把错误打出来，否则你只知道失败不知道原因。这是排查 ONNX 兼容性问题最重要的信息来源。

5.3 构建失败的常见原因与排查顺序

构建失败时，我一般按这个顺序排查：

看 parser 错误：如果是 unsupported operator，说明 ONNX 里有 TensorRT 不支持的算子，需要换实现或写插件。
看 shape 问题：如果报 shape 相关错误，检查 dynamic_axes 和 optimization profile 是否匹配。
看 workspace：如果报内存不足，加大 workspace。
看精度：如果 FP16 构建失败但 FP32 成功，可能是某些层不支持 FP16，可以设置config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)或者逐层指定精度。

5.4 引擎序列化与反序列化

构建好的引擎可以序列化到磁盘，下次直接反序列化加载，省去构建时间。构建一个复杂引擎可能要几分钟甚至更久，所以生产环境一定是构建一次、多次加载。

runtime = trt.Runtime(logger) with open("model.trt", "rb") as f: engine = runtime.deserialize_cuda_engine(f.read())

要注意，引擎是和 GPU 架构绑定的。在 A 卡上构建的引擎，拿到 B 卡上可能无法加载。所以部署时要确保构建环境和运行环境的 GPU 架构一致，或者干脆在目标机器上构建。

6. 精度对齐：从 FP32 到 FP16 的数值验证

6.1 为什么精度对齐是必须的

TensorRT 会做算子融合和精度转换，这些操作会引入数值差异。FP16 的精度损失是必然的，关键是要控制在可接受范围内。如果你的模型对数值敏感（比如涉及 softmax 后的概率、或者小目标的检测），精度对齐没做好，线上表现可能和离线差很多。

6.2 三层对齐法：PyTorch vs ONNX vs TensorRT

我的做法是做三层对比：

第一层：PyTorch 输出 vs ONNX Runtime 输出（验证导出正确性）；
第二层：ONNX Runtime 输出 vs TensorRT FP32 输出（验证 TensorRT 解析正确性）；
第三层：TensorRT FP32 输出 vs TensorRT FP16 输出（验证精度损失可接受）。

每一层都算 max abs diff 和相对误差。经验阈值是：第一层 diff < 1e-4，第二层 diff < 1e-3，第三层 diff < 1e-2。超过这些阈值就要警惕。

6.3 用 polygraphy 做自动对比

手动写对比代码比较繁琐，NVIDIA 提供了polygraphy这个工具，能自动做多层对比：

polygraphy run model.onnx \ --trt \ --fp16 \ --onnxrt \ --val-range input:0:1 \ --atol 1e-2 --rtol 1e-2

它会同时跑 ONNX Runtime 和 TensorRT，对比输出并报告差异。--atol和--rtol是绝对和相对容差。这个工具在排查精度问题时非常高效，强烈推荐。

6.4 精度不达标时的调整手段

如果 FP16 精度不达标，可以尝试：

部分层保持 FP32：对数值敏感的层（比如最后的分类头）单独设成 FP32。
关闭某些融合：有些融合会引入误差，可以尝试禁用。
调整校准：如果是 INT8，重新做校准，用更有代表性的校准数据。

我遇到过一次，某个模型的 FP16 输出在边界样本上偏差很大，最后发现是某一层的 LayerNorm 在 FP16 下精度不够，把那一层单独设成 FP32 后问题解决。

7. 性能调优与部署时的实际考量

7.1 batch size 与吞吐的权衡

batch size 不是越大越好。大 batch 能提高 GPU 利用率，但会增加延迟。在线服务通常追求低延迟，所以 batch 要小；离线批处理追求吞吐，batch 可以大。我的做法是测几个不同 batch size 下的延迟和吞吐，画一条曲线，找拐点。

用 trtexec 可以方便地测：

trtexec --loadEngine=model.trt --batch=8 --iterations=100

7.2 workspace 大小对构建的影响

workspace 太小会导致 TensorRT 无法使用某些高性能 tactic，构建出的引擎性能可能不是最优。但 workspace 太大也没必要，因为它只在构建时占用。我的经验是给 1-2GB，如果构建日志里出现“some tactics were skipped due to workspace limit”，就往上加。

7.3 引擎缓存与版本管理

生产环境里，引擎文件应该和模型版本、TensorRT 版本、GPU 型号绑定管理。我见过因为引擎文件和 TensorRT 版本不匹配导致加载失败的案例。建议在引擎文件名里带上这些信息，比如model_v1.2_trt8.6_sm86.trt，一目了然。

7.4 一个实际部署中的小技巧

如果服务需要频繁重启，每次重启都反序列化引擎会浪费时间。可以在服务启动时一次性加载引擎并创建 execution context，然后复用。execution context 是线程不安全的，多线程场景下每个线程需要独立的 context，但可以共享同一个 engine。

context = engine.create_execution_context() context.set_input_shape("input", (8, 3, 224, 224))

设置 input shape 这一步在动态 shape 场景下必须做，否则 context 不知道按哪个 shape 分配内存。

8. 我踩过的几个典型坑与排查思路

8.1 导出成功但 TensorRT 报 unsupported operator

这是最常见的坑。原因通常是 PyTorch 某个算子导出成了 ONNX 里的某个 op，而 TensorRT 不支持这个 op。排查方法是看 parser 的具体报错，它会告诉你哪个节点、哪个 op 不支持。解决办法有几种：换一个等价的算子实现、用 onnx-simplifier 化简、或者写自定义插件。

我有一次遇到Resize算子不支持某个 mode，最后把模型的 resize 逻辑改成了Interpolate的另一种模式才通过。

8.2 动态 shape 下输出 shape 不对

动态 shape 场景下，如果 optimization profile 配置不对，或者 ONNX 里的 shape 推导有问题，输出 shape 可能和预期不符。排查方法是先用固定 shape 跑通，确认逻辑正确，再逐步引入动态维度。

8.3 FP16 下出现 NaN

FP16 的数值范围比 FP32 小很多，某些中间结果可能溢出成 NaN。排查方法是逐层对比 FP32 和 FP16 的输出，找到第一个出现 NaN 的层。解决办法是把这个层设成 FP32，或者对输入做归一化。

8.4 引擎加载失败：版本或架构不匹配

前面提过，引擎和 GPU 架构绑定。如果加载时报错，先确认构建引擎的 GPU 和运行引擎的 GPU 是不是同一架构。另外，TensorRT 版本也要一致，不同版本的引擎格式可能不兼容。

排查这类问题的顺序是：先看错误信息里的关键词（是 version 还是 architecture），再对照构建和运行环境。我一般会在构建脚本里把 TensorRT 版本、CUDA 版本、GPU 型号都打印出来，方便对比。

8.5 一个关于 workspace 的隐蔽坑

有一次构建引擎时没报错，但性能比预期差很多。后来看构建日志，发现因为 workspace 设得太小，TensorRT 跳过了一些高性能 tactic。把 workspace 从 256MB 加到 2GB 后，性能提升了将近 40%。这个坑很隐蔽，因为构建是成功的，只是性能没到最优。所以构建日志一定要看，尤其是关于 tactic 的部分。

9. 收尾：一些个人经验

整条链路走下来，我最大的体会是：版本管理和精度对齐是这条路上两个最耗时的环节，但也是最值得提前投入的。与其出了问题再回头查版本，不如一开始就把版本对应关系理清楚，把每一层的数值对比脚本写好。

另外，trtexec 和 polygraphy 这两个工具真的能省很多时间。前者帮你快速验证 ONNX 能不能转，后者帮你快速定位精度问题。我现在的习惯是：任何新模型，先用 trtexec 跑一遍看能不能过，再用 polygraphy 做精度对比，两个都过了再写正式的部署代码。

最后分享一个小技巧：把整个转换流程写成一个脚本，从导出 ONNX 到构建引擎到精度验证一条龙跑完，每次模型更新直接重跑。这样既避免了手动操作的遗漏，也方便做版本对比。我现在的脚本里会记录每次转换的 PyTorch 版本、ONNX opset、TensorRT 版本、GPU 型号和精度对比结果，出问题时回溯起来非常方便。