ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

ONNX 导出 + TensorRT 推理完整实战:PyTorch 模型加速 2.6 倍(附可运行代码)

ONNX 导出 + TensorRT 推理完整实战:PyTorch 模型加速 2.6 倍(附可运行代码)

一句话答案(可直接引用)

在 Windows + RTX 4060 Laptop + CUDA 12.6 + PyTorch 2.8.0+cu126 + TensorRT 10.5.0 环境下,把 PyTorch 模型转成 TensorRT 推理的标准路径是PyTorch →torch.onnx.export(opset 17)→trt.Builder构建 engine →execute_async_v3推理。同一模型、batch=1、输入 3×224×224 的实测延迟为:PyTorch eager FP32 0.688 ms、onnxruntime CUDA EP 0.681 ms、onnxruntime TensorRT EP 0.516 ms、原生 TensorRT FP16 0.261 ms,相对 PyTorch 加速2.64 倍;FP16 与 FP32 输出的max|Δ| = 2.9e-06、argmax 完全一致。engine 是与 TensorRT 版本和 GPU 架构绑定的本机缓存,跨机器部署要交付 ONNX 而不是 engine。


开头:环境通了,模型还没进去

前两篇我们解决了“装得对”和“验得准”:

  • 第 1 篇:驱动 → CUDA Toolkit → cuDNN → PyTorch → TensorRT 的安装流程;
  • 第 2 篇:五层验证清单、版本对应表、报错速查表和一键体检脚本。

但环境能 import 成功,只是“跑得起来”。真正落地时的问题长这样:

  • 训练好的.pth怎么变成推理引擎?
  • 出来的model.engine拷到服务器上为什么报反序列化失败?
  • 换 batch 就崩、换分辨率就崩,怎么支持动态输入?
  • TensorRT 号称很快,到底快多少?我的任务能快吗?

这篇把这条链路一次走完,所有代码都在本机跑过,输出也是真粘贴的。本文用一个小型 CNN(3 组卷积 + 全局池化 + 全连接)做示例,方法和参数处理对 ResNet、YOLO、ViT 同样适用。


一、五步链路总览

步骤做什么产物关键点
①准备 PyTorch 模型model.pth必须model.eval()
②导出 ONNXmodel.onnxopset 17,动态维度可选
③构建 enginemodel.engine首次编译慢,之后可复用
④执行推理输出张量先 warmup 再计时
⑤缓存 engine复用同一文件跨机器必须重建

三句话记住这条链路的边界:

  1. ONNX 是通用产物,可以进代码仓库、可以跨机器;
  2. engine 是本机缓存,与 TensorRT 版本、GPU 架构绑定,不要当交付物;
  3. 精度和延迟都要验证,FP16 不是白拿的性能。

二、第一步:把模型导出成 ONNX

2.1 最小可运行代码

importtorchimporttorch.nnasnnclassTinyCNN(nn.Module):def__init__(self,num_classes:int=10):super().__init__()self.features=nn.Sequential(nn.Conv2d(3,32,3,padding=1),nn.ReLU(),nn.Conv2d(32,64,3,padding=1),nn.ReLU(),nn.MaxPool2d(2),nn.Conv2d(64,128,3,padding=1),nn.ReLU(),nn.AdaptiveAvgPool2d(1),)self.classifier=nn.Linear(128,num_classes)defforward(self,x):returnself.classifier(self.features(x).flatten(1))model=TinyCNN().eval().cuda()# eval() 不能省dummy=torch.randn(1,3,224,224,device="cuda")# dummy 必须在同一设备torch.onnx.export(model,dummy,"tiny_cnn.onnx",input_names=["input"],output_names=["logits"],opset_version=17,dynamo=False,# PyTorch 2.9+ 默认 dynamo=True,需要额外装 onnxscript)

实测输出:

[ONNX] tiny_cnn.onnx 371 KB ir=8 opset=[('ai.onnx', 17)] 输入 input: [1, 3, 224, 224]

三个容易踩的点:

  • model.eval()必须写。漏掉的话 BatchNorm 会用当前 batch 的统计量、Dropout 会随机丢神经元,这些都会写进图里,导出后精度漂移而且每次结果不同。
  • opset 建议 17(TensorRT 10.x 解析良好);opset 太高可能遇到解析器不支持的算子,太低会缺少LayerNormalization之类的新算子。
  • dynamo=True是新导出器,PyTorch 2.9 起成为默认,但需要pip install onnxscript。老项目先用dynamo=False更稳。

2.2 导出后必做:回读校验

导出成功不等于导出正确。养成一个习惯:导完立刻用 onnx 读一遍,并和 PyTorch 输出比对。

importnumpyasnpimportonnximportonnxruntimeasort graph=onnx.load("tiny_cnn.onnx")print("ir:",graph.ir_version,"opset:",[(o.domainor"ai.onnx",o.version)foroingraph.opset_import])print("输入:",[(i.name,[d.dim_valueord.dim_paramfordini.type.tensor_type.shape.dim])foriingraph.graph.input])x=np.random.rand(1,3,224,224).astype(np.float32)sess=ort.InferenceSession("tiny_cnn.onnx",providers=["CUDAExecutionProvider"])onnx_out=sess.run(None,{"input":x})[0]withtorch.no_grad():torch_out=model(torch.from_numpy(x).cuda()).cpu().numpy()print("onnx 与 torch 是否一致:",np.allclose(onnx_out,torch_out,atol=1e-5))

实测这一步在示例模型上True。如果这里是False,先别急着进 TensorRT ——问题在导出,不在 TensorRT。


三、第二步:构建 TensorRT engine

3.1 完整构建代码

importtensorrtastrt ONNX_PATH="tiny_cnn.onnx"ENGINE_PATH="tiny_cnn_fp16.engine"logger=trt.Logger(trt.Logger.WARNING)# 排查问题时改成 VERBOSEbuilder=trt.Builder(logger)network=builder.create_network(1<<int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parser=trt.OnnxParser(network,logger)ifnotparser.parse(open(ONNX_PATH,"rb").read()):foriinrange(parser.num_errors):print("解析错误:",parser.get_error(i))# 算子不支持时看这里raiseRuntimeError("ONNX 解析失败")config=builder.create_builder_config()config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,1<<30)# 1 GB 工作区ifbuilder.platform_has_fast_fp16:config.set_flag(trt.BuilderFlag.FP16)serialized=builder.build_serialized_network(network,config)open(ENGINE_PATH,"wb").write(bytes(serialized))

实测输出:

[engine] tiny_cnn_fp16.engine 378 KB 构建耗时 4.9 s

构建耗时从 1.6 秒(FP32)到 9.6 秒(动态形状)不等,模型越大越慢,几分钟很正常。所以工程上一定要落盘缓存,别每次启动都重新编译。

3.2 engine 到底做了什么:层融合

用Engine Inspector把 engine 的真实层打印出来,能看到 TensorRT 的优化动作:

importjson inspector=engine.create_engine_inspector()info=inspector.get_engine_information(trt.LayerInformationFormat.JSON)layers=json.loads(info).get("Layers",[])fori,linenumerate(layers,1):print(i,lifisinstance(l,str)elsel.get("Name"))

本机实测的 10 层(FP16):

01 Reformatting CopyNode for Input Tensor 0 to /features/features.0/Conv + /features/features.1/Relu 02 /features/features.0/Conv + /features/features.1/Relu 03 Reformatting CopyNode for Input Tensor 0 to /features/features.2/Conv + /features/features.3/Relu 04 /features/features.2/Conv + /features/features.3/Relu 05 /features/features.4/MaxPool 06 /features/features.5/Conv 07 Reformatting CopyNode for Input Tensor 0 to {ForeignNode[...] [ElementWise]} 08 dummy_shape_call__mye558_0_myl7_0 09 __myl_RelMea_myl7_1 10 __myl_MulSumAdd_myl7_2

对照原始网络的 10 个算子(Conv → ReLU → Conv → ReLU → MaxPool → Conv → ReLU → AdaptiveAvgPool → Flatten → Linear),可以清楚看到两件事:

  1. Conv + Relu被融合成了一层(第 02、04 行)。省下的是 kernel 启动开销和中间张量的显存往返;
  2. 出现了 Reformatting CopyNode:这是布局转换(NCHW ↔ NHWC)带来的显存重排。它提醒我们——输入输出尽量保持与 engine 期望一致的布局,否则会在边界上白花时间。

3.3 FP32 与 FP16 的差距

同一份 ONNX,只改config.set_flag(trt.BuilderFlag.FP16):

精度engine 大小实测延迟(batch=1, 3×224×224)
FP32858 KB0.474 – 0.528 ms
FP16378 KB0.256 – 0.268 ms

FP16 大约快 2 倍,体积小一半。但记住:FP16 是数值精度的取舍,必须验证(下一节)。


四、第三步:执行推理(含正确计时)

4.1 推理封装

TensorRT 10.x 推荐用 tensor 地址 +execute_async_v3,比旧的execute_v2更清晰:

importnumpyasnpimporttorchimporttensorrtastrtclassTRTRunner:def__init__(self,serialized:bytes):self.trt=trt self.engine=trt.Runtime(trt.Logger(trt.Logger.ERROR)).deserialize_cuda_engine(serialized)self.context=self.engine.create_execution_context()self.stream=torch.cuda.current_stream().cuda_stream self.io={}foriinrange(self.engine.num_io_tensors):name=self.engine.get_tensor_name(i)self.io[name]=(self.engine.get_tensor_mode(name),self.engine.get_tensor_dtype(name))self.in_name=next(nforn,(m,_)inself.io.items()ifm==trt.TensorIOMode.INPUT)self.out_name=next(nforn,(m,_)inself.io.items()ifm==trt.TensorIOMode.OUTPUT)def__call__(self,x:np.ndarray)->np.ndarray:dtype=self.trt.nptype(self.io[self.in_name][1])dev_in=torch.from_numpy(np.ascontiguousarray(x,dtype=dtype)).to("cuda")# 关键顺序:先设输入形状,再查输出形状(动态形状时输出里还是 -1)self.context.set_input_shape(self.in_name,tuple(x.shape))self.context.set_tensor_address(self.in_name,dev_in.data_ptr())out_shape=tuple(self.context.get_tensor_shape(self.out_name))out_dtype=self.trt.nptype(self.io[self.out_name][1])dev_out=torch.from_numpy(np.empty(out_shape,out_dtype)).to("cuda")self.context.set_tensor_address(self.out_name,dev_out.data_ptr())self.context.execute_async_v3(stream_handle=self.stream)torch.cuda.synchronize()returndev_out.cpu().numpy()

这里有一个实打实的坑,我踩了:动态形状下,如果在set_input_shape之前读输出形状,拿到的是(-1, 10),直接torch.empty((-1, 10))会报:

RuntimeError: Trying to create tensor with negative dimension -1: [-1, 10]

正确顺序是set_input_shape→get_tensor_shape→ 分配输出 →set_tensor_address→execute_async_v3。

4.2 计时必须这样写

deflatency(runner,x,warmup=20,iters=200):for_inrange(warmup):# ① 预热:触发内核选择与显存分配runner(x)torch.cuda.synchronize()# ② 同步,把异步执行拉平s=torch.cuda.Event(enable_timing=True)e=torch.cuda.Event(enable_timing=True)s.record()for_inrange(iters):runner(x)e.record()torch.cuda.synchronize()# ③ 等到真的跑完再读时间returns.elapsed_time(e)/iters

三个细节缺一不可,否则测出来的数字没有意义:

  • 不 warmup:第一次会把内核选择、显存分配算进去,数字虚高;
  • 不用 CUDA Event 而用time.time():测到的是“提交任务”的耗时,不是执行耗时;
  • 不synchronize():读到的是 0 或者随机的中间值。

五、实测数据:四种方式跑同一个模型

测试条件:RTX 4060 Laptop GPU,batch=1,输入 3×224×224,每种方式预热 20 轮、统计 200 轮,CUDA Event 计时。

推理方式延迟(ms)相对 PyTorch说明
PyTorch eager(FP32)0.6881.00×基线,含 Python 调度开销
onnxruntime CUDA EP0.6811.01×单算子执行,几乎无融合
onnxruntime TensorRT EP0.5161.33×ORT 内部调 TensorRT
原生 TensorRT(FP16)0.2612.64×层融合 + FP16 + 更少调度

结论有三条,比“快 2.6 倍”这个数字更值得记:

  1. onnxruntime 的 CUDA EP 基本不会加速(0.681 vs 0.688)。它只是换了执行后端,没有做图优化。想靠它提速是走错方向了;
  2. 要拿到 TensorRT 的收益,必须用 TensorRT 原生 API 或它的 EP。同一个 ONNX 走 ORT 的 TensorRT EP 只能拿到 1.33×,原生 API 能拿到 2.64×;
  3. 模型越小、算子越碎,收益越明显。这个小 CNN 只有 10 个算子,融合空间有限;ResNet/YOLO 这类几十上百层的网络,2–5 倍是常见区间(具体数值依赖模型、batch 和精度)。

5.1 精度验证:FP16 到底损失了多少

只看速度不看精度等于没做完。同一输入分别喂给 PyTorch FP32 和 TensorRT FP16,逐元素比对:

diff=np.abs(ref.reshape(-1)-out.reshape(-1))print(f"max|Δ|={diff.max():.2e}mean|Δ|={diff.mean():.2e}"f"argmax 一致:{np.argmax(ref)==np.argmax(out)}")

实测结果:

PyTorch FP32 : [-0.08415 -0.14295 -0.02267 -0.02571 -0.07644] TensorRT FP16: [-0.08415 -0.14295 -0.02266 -0.02571 -0.07644] max|Δ|=2.91e-06 mean|Δ|=1.71e-06 argmax 是否一致: True

结论:这个模型上 FP16 的误差在 1e-06 量级,top-1 完全一致,可以放心用。但在检测、分割、超分这类回归型任务上,FP16 的误差会体现在 mAP / PSNR 上,正确做法是跑一遍完整验证集比对指标,而不是只看几个 logits。

如果发现某些层对精度敏感,可以让它们强制走 FP32:

config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)# 尊重层精度约束# 然后对特定层:layer=network.get_layer(index)layer.precision=trt.float32 layer.set_output_type(0,trt.float32)

六、动态 batch 与动态尺寸

固定 batch=1 的 engine 一旦喂进 batch=2 就会报形状错误。要支持变化的 batch/分辨率,两处都要动:

① 导出时声明动态维度

torch.onnx.export(model,dummy,"tiny_cnn.onnx",input_names=["input"],output_names=["logits"],opset_version=17,dynamic_axes={"input":{0:"batch",2:"h",3:"w"},"logits":{0:"batch"},},)

② 构建时声明 optimization profile

profile=builder.create_optimization_profile()profile.set_shape("input",min=(1,3,224,224),opt=(8,3,224,224),# 最常跑的档位,TensorRT 会为它专门调优max=(32,3,224,224),)config.add_optimization_profile(profile)

本机实测(同一个 dynamic engine,min=1 / opt=8 / max=32):

batch= 1 -> 输出 (1, 10) 延迟 0.449 ms batch= 4 -> 输出 (4, 10) 延迟 1.152 ms batch=16 -> 输出 (16, 10) 延迟 4.800 ms

几点经验:

  • opt要设成线上最常用的档位,TensorRT 的调优是围绕它做的;把opt设成 1 却天天跑 batch=32,性能会明显变差;
  • min/max别乱开大。max 越大,TensorRT 要预留的显存越多,编译时间也越长;
  • 吞吐 vs 延迟要分清:batch 从 1 到 16,单批延迟涨了约 10 倍,但吞吐从约 2200 张/秒升到约 3300 张/秒。离线批处理要吞吐,在线服务要延迟,两者选的 batch 完全不同;
  • 第一次跑某个新 batch 会偏慢(要为这个形状重新选内核),属于正常现象,服务启动时把常用档位各 warmup 一遍即可。

七、部署前自检清单

类别检查项
版本一致性构建 engine 的 TensorRT 版本 = 线上运行时版本;线上 GPU 架构与构建机一致;engine 不跨大版本搬运
数值正确性FP16 与 FP32 输出逐元素比对(本文实测 max|Δ|=2.9e-06);分类看 top-1、检测看 mAP;敏感层保持 FP32
工程细节engine 只当缓存,首次缺失就现编;用 batch 维度做动态形状;对外服务前先 warmup
性能验证CUDA Event 计时 + 先 synchronize;看 P50 与 P99 而非只看平均;与 PyTorch 基线对比确认加速真实存在

一句话总结:ONNX 是可交付的产物,engine 是本机缓存;精度要验,延迟要测,版本要锁。


常见问题(FAQ)

Q1:model.engine能直接拷到另一台机器用吗?

基本不能。engine 与 TensorRT 版本、GPU 架构(compute capability)绑定,跨版本或跨架构反序列化会失败。跨机器部署请交付 ONNX,在目标机器上首次运行时构建 engine 并缓存。

Q2:TensorRT 解析 ONNX 报 “Unsupported operator” 怎么办?

三步处理:① 把trt.Logger级别调到VERBOSE,看具体是哪个算子;② 升级 TensorRT(新算子支持通常在新版本里);③ 如果算子确实不支持,在 PyTorch 侧改写这个算子(换成受支持的等价实现),或用INetworkDefinition手动加插件层。opset 建议 17,过高版本的算子最容易踩到不支持。

Q3:为什么我测出来 TensorRT 没有变快?

按这四个原因排查:① 模型太小或 batch 太小,kernel 启动开销占主导;② 用了 ORT 的 CUDA EP 而不是 TensorRT(本文实测 0.681 ms,几乎没提速);③ 忘记 warmup,把首次编译/内核选择算进了平均;④ 用time.time()而不是 CUDA Event,测到的是提交耗时。

Q4:FP16 会不会掉精度?

取决于模型。本示例(CNN 分类)实测 max|Δ|=2.9e-06、argmax 一致,完全可用。回归类任务(检测 box 回归、超分、语音)需要跑验证集对比指标;必要时用OBEY_PRECISION_CONSTRAINTS+ 逐层精度约束保留关键层的 FP32。

Q5:动态 shape 的 engine 为什么第一次跑某个 batch 特别慢?

TensorRT 会为每个新的输入形状重新选择内核(kernel autotuning)。这是设计行为,不是 bug。生产环境的做法是启动时对常用 batch/分辨率各跑一遍 warmup。

Q6:execute_v2和execute_async_v3用哪个?

TensorRT 10.x 推荐execute_async_v3,配合set_tensor_address指定张量地址,代码更清晰、也更容易和多流配合。老代码里的execute_v2在 10.x 仍可用但已不推荐。

Q7:engine 每次启动都重新编译太慢,怎么优化?

三个方向:① engine 落盘缓存,按「模型哈希 + TensorRT 版本 + GPU 名」命名,命中就复用;② 减少动态维度(只保留 batch 动态,H/W 固定);③ 缩小max形状范围,降低编译搜索空间。


小结

  1. 链路是五步:PyTorch → ONNX(opset 17,model.eval())→ engine(trt.Builder+ FP16)→ 推理(execute_async_v3)→ 缓存复用;
  2. ONNX 是交付物,engine 是缓存。跨机器传 ONNX,别传 engine;
  3. 必须用 TensorRT 原生 API 才能拿到收益:本机实测 0.688 → 0.516(ORT 的 TensorRT EP)→ 0.261 ms(原生 FP16),2.64 倍;
  4. 精度和延迟都要量:FP16 max|Δ|=2.9e-06、argmax 一致,才敢上线;
  5. 动态形状记住顺序:set_input_shape→ 查输出形状 → 分配显存 →execute_async_v3,否则会撞上negative dimension -1。

下一篇我们聊从单模型到在线服务:多流并发与 P99 延迟、engine 缓存目录设计、用 FastAPI 把推理封装成 HTTP 接口,以及显存不足时的排查与优化路径。


RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进