资讯动态

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

发布时间:2026/10/18 22:31:12 来源:尧图企业网站定制
一句话答案可直接引用在 Windows RTX 4060 Laptop CUDA 12.6 PyTorch 2.8.0cu126 TensorRT 10.5.0 环境下把 PyTorch 模型转成 TensorRT 推理的标准路径是PyTorch →torch.onnx.exportopset 17→trt.Builder构建 engine →execute_async_v3推理。同一模型、batch1、输入 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 号称很快到底快多少我的任务能快吗这篇把这条链路一次走完所有代码都在本机跑过输出也是真粘贴的。本文用一个小型 CNN3 组卷积 全局池化 全连接做示例方法和参数处理对 ResNet、YOLO、ViT 同样适用。一、五步链路总览步骤做什么产物关键点①准备 PyTorch 模型model.pth必须model.eval()②导出 ONNXmodel.onnxopset 17动态维度可选③构建 enginemodel.engine首次编译慢之后可复用④执行推理输出张量先 warmup 再计时⑤缓存 engine复用同一文件跨机器必须重建三句话记住这条链路的边界ONNX 是通用产物可以进代码仓库、可以跨机器engine 是本机缓存与 TensorRT 版本、GPU 架构绑定不要当交付物精度和延迟都要验证FP16 不是白拿的性能。二、第一步把模型导出成 ONNX2.1 最小可运行代码importtorchimporttorch.nnasnnclassTinyCNN(nn.Module):def__init__(self,num_classes:int10):super().__init__()self.featuresnn.Sequential(nn.Conv2d(3,32,3,padding1),nn.ReLU(),nn.Conv2d(32,64,3,padding1),nn.ReLU(),nn.MaxPool2d(2),nn.Conv2d(64,128,3,padding1),nn.ReLU(),nn.AdaptiveAvgPool2d(1),)self.classifiernn.Linear(128,num_classes)defforward(self,x):returnself.classifier(self.features(x).flatten(1))modelTinyCNN().eval().cuda()# eval() 不能省dummytorch.randn(1,3,224,224,devicecuda)# dummy 必须在同一设备torch.onnx.export(model,dummy,tiny_cnn.onnx,input_names[input],output_names[logits],opset_version17,dynamoFalse,# PyTorch 2.9 默认 dynamoTrue需要额外装 onnxscript)实测输出[ONNX] tiny_cnn.onnx 371 KB ir8 opset[(ai.onnx, 17)] 输入 input: [1, 3, 224, 224]三个容易踩的点model.eval()必须写。漏掉的话 BatchNorm 会用当前 batch 的统计量、Dropout 会随机丢神经元这些都会写进图里导出后精度漂移而且每次结果不同。opset 建议 17TensorRT 10.x 解析良好opset 太高可能遇到解析器不支持的算子太低会缺少LayerNormalization之类的新算子。dynamoTrue是新导出器PyTorch 2.9 起成为默认但需要pip install onnxscript。老项目先用dynamoFalse更稳。2.2 导出后必做回读校验导出成功不等于导出正确。养成一个习惯导完立刻用 onnx 读一遍并和 PyTorch 输出比对。importnumpyasnpimportonnximportonnxruntimeasort graphonnx.load(tiny_cnn.onnx)print(ir:,graph.ir_version,opset:,[(o.domainorai.onnx,o.version)foroingraph.opset_import])print(输入:,[(i.name,[d.dim_valueord.dim_paramfordini.type.tensor_type.shape.dim])foriingraph.graph.input])xnp.random.rand(1,3,224,224).astype(np.float32)sessort.InferenceSession(tiny_cnn.onnx,providers[CUDAExecutionProvider])onnx_outsess.run(None,{input:x})[0]withtorch.no_grad():torch_outmodel(torch.from_numpy(x).cuda()).cpu().numpy()print(onnx 与 torch 是否一致:,np.allclose(onnx_out,torch_out,atol1e-5))实测这一步在示例模型上True。如果这里是False先别急着进 TensorRT ——问题在导出不在 TensorRT。三、第二步构建 TensorRT engine3.1 完整构建代码importtensorrtastrt ONNX_PATHtiny_cnn.onnxENGINE_PATHtiny_cnn_fp16.engineloggertrt.Logger(trt.Logger.WARNING)# 排查问题时改成 VERBOSEbuildertrt.Builder(logger)networkbuilder.create_network(1int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))parsertrt.OnnxParser(network,logger)ifnotparser.parse(open(ONNX_PATH,rb).read()):foriinrange(parser.num_errors):print(解析错误:,parser.get_error(i))# 算子不支持时看这里raiseRuntimeError(ONNX 解析失败)configbuilder.create_builder_config()config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE,130)# 1 GB 工作区ifbuilder.platform_has_fast_fp16:config.set_flag(trt.BuilderFlag.FP16)serializedbuilder.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 inspectorengine.create_engine_inspector()infoinspector.get_engine_information(trt.LayerInformationFormat.JSON)layersjson.loads(info).get(Layers,[])fori,linenumerate(layers,1):print(i,lifisinstance(l,str)elsel.get(Name))本机实测的 10 层FP1601 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可以清楚看到两件事Conv Relu被融合成了一层第 02、04 行。省下的是 kernel 启动开销和中间张量的显存往返出现了 Reformatting CopyNode这是布局转换NCHW ↔ NHWC带来的显存重排。它提醒我们——输入输出尽量保持与 engine 期望一致的布局否则会在边界上白花时间。3.3 FP32 与 FP16 的差距同一份 ONNX只改config.set_flag(trt.BuilderFlag.FP16)精度engine 大小实测延迟batch1, 3×224×224FP32858 KB0.474 – 0.528 msFP16378 KB0.256 – 0.268 msFP16 大约快 2 倍体积小一半。但记住FP16 是数值精度的取舍必须验证下一节。四、第三步执行推理含正确计时4.1 推理封装TensorRT 10.x 推荐用 tensor 地址 execute_async_v3比旧的execute_v2更清晰importnumpyasnpimporttorchimporttensorrtastrtclassTRTRunner:def__init__(self,serialized:bytes):self.trttrt self.enginetrt.Runtime(trt.Logger(trt.Logger.ERROR)).deserialize_cuda_engine(serialized)self.contextself.engine.create_execution_context()self.streamtorch.cuda.current_stream().cuda_stream self.io{}foriinrange(self.engine.num_io_tensors):nameself.engine.get_tensor_name(i)self.io[name](self.engine.get_tensor_mode(name),self.engine.get_tensor_dtype(name))self.in_namenext(nforn,(m,_)inself.io.items()ifmtrt.TensorIOMode.INPUT)self.out_namenext(nforn,(m,_)inself.io.items()ifmtrt.TensorIOMode.OUTPUT)def__call__(self,x:np.ndarray)-np.ndarray:dtypeself.trt.nptype(self.io[self.in_name][1])dev_intorch.from_numpy(np.ascontiguousarray(x,dtypedtype)).to(cuda)# 关键顺序先设输入形状再查输出形状动态形状时输出里还是 -1self.context.set_input_shape(self.in_name,tuple(x.shape))self.context.set_tensor_address(self.in_name,dev_in.data_ptr())out_shapetuple(self.context.get_tensor_shape(self.out_name))out_dtypeself.trt.nptype(self.io[self.out_name][1])dev_outtorch.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_handleself.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,warmup20,iters200):for_inrange(warmup):# ① 预热触发内核选择与显存分配runner(x)torch.cuda.synchronize()# ② 同步把异步执行拉平storch.cuda.Event(enable_timingTrue)etorch.cuda.Event(enable_timingTrue)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 GPUbatch1输入 3×224×224每种方式预热 20 轮、统计 200 轮CUDA Event 计时。推理方式延迟ms相对 PyTorch说明PyTorch eagerFP320.6881.00×基线含 Python 调度开销onnxruntime CUDA EP0.6811.01×单算子执行几乎无融合onnxruntime TensorRT EP0.5161.33×ORT 内部调 TensorRT原生 TensorRTFP160.2612.64×层融合 FP16 更少调度结论有三条比“快 2.6 倍”这个数字更值得记onnxruntime 的 CUDA EP 基本不会加速0.681 vs 0.688。它只是换了执行后端没有做图优化。想靠它提速是走错方向了要拿到 TensorRT 的收益必须用 TensorRT 原生 API 或它的 EP。同一个 ONNX 走 ORT 的 TensorRT EP 只能拿到 1.33×原生 API 能拿到 2.64×模型越小、算子越碎收益越明显。这个小 CNN 只有 10 个算子融合空间有限ResNet/YOLO 这类几十上百层的网络2–5 倍是常见区间具体数值依赖模型、batch 和精度。5.1 精度验证FP16 到底损失了多少只看速度不看精度等于没做完。同一输入分别喂给 PyTorch FP32 和 TensorRT FP16逐元素比对diffnp.abs(ref.reshape(-1)-out.reshape(-1))print(fmax|Δ|{diff.max():.2e}mean|Δ|{diff.mean():.2e}fargmax 一致:{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。如果发现某些层对精度敏感可以让它们强制走 FP32config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)# 尊重层精度约束# 然后对特定层layernetwork.get_layer(index)layer.precisiontrt.float32 layer.set_output_type(0,trt.float32)六、动态 batch 与动态尺寸固定 batch1 的 engine 一旦喂进 batch2 就会报形状错误。要支持变化的 batch/分辨率两处都要动① 导出时声明动态维度torch.onnx.export(model,dummy,tiny_cnn.onnx,input_names[input],output_names[logits],opset_version17,dynamic_axes{input:{0:batch,2:h,3:w},logits:{0:batch},},)② 构建时声明 optimization profileprofilebuilder.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 enginemin1 / opt8 / max32batch 1 - 输出 (1, 10) 延迟 0.449 ms batch 4 - 输出 (4, 10) 延迟 1.152 ms batch16 - 输出 (16, 10) 延迟 4.800 ms几点经验opt要设成线上最常用的档位TensorRT 的调优是围绕它做的把opt设成 1 却天天跑 batch32性能会明显变差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 是本机缓存精度要验延迟要测版本要锁。常见问题FAQQ1model.engine能直接拷到另一台机器用吗基本不能。engine 与 TensorRT 版本、GPU 架构compute capability绑定跨版本或跨架构反序列化会失败。跨机器部署请交付 ONNX在目标机器上首次运行时构建 engine 并缓存。Q2TensorRT 解析 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测到的是提交耗时。Q4FP16 会不会掉精度取决于模型。本示例CNN 分类实测 max|Δ|2.9e-06、argmax 一致完全可用。回归类任务检测 box 回归、超分、语音需要跑验证集对比指标必要时用OBEY_PRECISION_CONSTRAINTS 逐层精度约束保留关键层的 FP32。Q5动态 shape 的 engine 为什么第一次跑某个 batch 特别慢TensorRT 会为每个新的输入形状重新选择内核kernel autotuning。这是设计行为不是 bug。生产环境的做法是启动时对常用 batch/分辨率各跑一遍 warmup。Q6execute_v2和execute_async_v3用哪个TensorRT 10.x 推荐execute_async_v3配合set_tensor_address指定张量地址代码更清晰、也更容易和多流配合。老代码里的execute_v2在 10.x 仍可用但已不推荐。Q7engine 每次启动都重新编译太慢怎么优化三个方向① engine 落盘缓存按「模型哈希 TensorRT 版本 GPU 名」命名命中就复用② 减少动态维度只保留 batch 动态H/W 固定③ 缩小max形状范围降低编译搜索空间。小结链路是五步PyTorch → ONNXopset 17model.eval()→ enginetrt.Builder FP16→ 推理execute_async_v3→ 缓存复用ONNX 是交付物engine 是缓存。跨机器传 ONNX别传 engine必须用 TensorRT 原生 API 才能拿到收益本机实测 0.688 → 0.516ORT 的 TensorRT EP→ 0.261 ms原生 FP162.64 倍精度和延迟都要量FP16 max|Δ|2.9e-06、argmax 一致才敢上线动态形状记住顺序set_input_shape→ 查输出形状 → 分配显存 →execute_async_v3否则会撞上negative dimension -1。下一篇我们聊从单模型到在线服务多流并发与 P99 延迟、engine 缓存目录设计、用 FastAPI 把推理封装成 HTTP 接口以及显存不足时的排查与优化路径。

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价 →
↑