★ FEATURED ARTICLE 1. 为什么模型部署绕不开ONNX这个中间站做过深度学习模型落地的朋友大概率都经历过这样的场景在PyTorch里训练好的模型推理速度怎么调都上不去GPU利用率忽高忽低延迟波动大得让人心里没底。这时候TensorRT就会进入视野——它能把推理性能压榨到极致但问题是PyTorch的.pth或.pb文件不能直接喂给TensorRT 7.0中间必须经过一次格式转换。这个转换的桥梁就是ONNX。ONNX的全称是Open Neural Network Exchange你可以把它理解成模型格式的“通用插座”——PyTorch、TensorFlow、MXNet等框架都能把自己的模型导出成ONNX格式而TensorRT则可以从ONNX文件解析出网络结构再针对NVIDIA GPU做深度优化。整个链路是这样的PyTorch训练 → 导出ONNX → TensorRT解析ONNX → 生成优化后的推理引擎.plan或.trt文件。为什么不能跳过ONNX直接转因为TensorRT 7.0自带的PyTorch解析器功能非常有限只支持极少数算子而且对动态图的支持几乎为零。ONNX作为中间表示既能保留计算图的完整结构又能被TensorRT的ONNX Parser稳定解析。所以不管你的原始模型是.pth还是.pb只要想上TensorRTONNX这一步基本绕不开。这里需要先澄清一个容易混淆的点标题里说的“pb文件”在不同语境下含义不同。TensorFlow的.pb是冻结计算图而PyTorch生态里有时也会把序列化后的模型叫.pb但更常见的是.pth或.pt。不管原始格式是什么核心思路是一致的——先转成ONNX再从ONNX转到TensorRT。下面我会以PyTorch模型为例把整条链路拆开讲透。注意TensorRT 7.0是一个相对早期的版本它对ONNX的算子集版本有明确要求太新的opset可能导致解析失败。这一点在后文会详细展开。2. 环境搭建版本匹配比什么都重要2.1 TensorRT 7.0的安装与CUDA/cuDNN的版本锁定TensorRT 7.0不是随便装一个版本就能用的它对CUDA和cuDNN的版本有硬性要求。根据我的实际经验TensorRT 7.0通常搭配CUDA 10.2和cuDNN 7.6.5使用。如果你用的是CUDA 11.x那TensorRT 7.0基本跑不起来必须升级到TensorRT 8.x。所以第一步是先确认自己的CUDA版本nvcc --version如果输出显示的是CUDA 10.2那就可以继续。安装TensorRT 7.0的方式有几种直接下载tar包解压、用deb包安装、或者通过pip安装tensorrt库。我推荐用tar包方式因为可控性最强# 解压TensorRT 7.0 tar包 tar -xzvf TensorRT-7.0.0.11.Ubuntu-18.04.x86_64-gnu.cuda-10.2.cudnn7.6.tar.gz # 设置环境变量 export LD_LIBRARY_PATH$LD_LIBRARY_PATH:/path/to/TensorRT-7.0.0.11/lib export PATH$PATH:/path/to/TensorRT-7.0.0.11/bin然后安装Python绑定cd TensorRT-7.0.0.11/python pip install tensorrt-7.0.0.11-cp37-none-linux_x86_64.whl这里有个坑Python版本必须和whl文件匹配。TensorRT 7.0官方提供的whl通常只覆盖Python 3.6和3.7如果你用的是Python 3.8要么自己编译要么换环境。我建议直接用conda创建一个Python 3.7的虚拟环境省去很多麻烦。2.2 PyTorch与ONNX的版本选择PyTorch这边建议用1.4到1.6之间的版本。太新的PyTorch比如1.10以上导出的ONNX可能包含TensorRT 7.0不支持的算子属性。ONNX的版本也要控制onnx库建议用1.6或1.7onnxruntime可以用来验证导出的模型是否正确。pip install torch1.5.0 torchvision0.6.0 pip install onnx1.7.0 onnxruntime1.4.0安装完之后用一段简单的代码验证环境是否正常import torch import onnx import tensorrt as trt print(torch.__version__) print(onnx.__version__) print(trt.__version__)如果三行都能正常打印出版本号说明基础环境没问题。接下来就可以进入模型转换的核心环节了。提示如果你在服务器上有多张GPU卡建议先用CUDA_VISIBLE_DEVICES0锁定一张卡做转换和测试避免多卡环境下的设备分配问题。3. 从PyTorch到ONNX导出过程中的关键决策3.1 动态轴设置与输入输出命名PyTorch导出ONNX的核心API是torch.onnx.export。这个函数看起来简单但里面的参数每一个都会影响后续TensorRT的解析结果。先看一个典型的导出代码import torch import torch.onnx # 假设model是你的PyTorch模型已经加载了权重 model.eval() dummy_input torch.randn(1, 3, 224, 224).cuda() torch.onnx.export( model, dummy_input, model.onnx, export_paramsTrue, opset_version11, do_constant_foldingTrue, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size, 2: height, 3: width}, output: {0: batch_size} } )这里有几个关键点需要展开说。opset_version的选择很讲究TensorRT 7.0对opset 11的支持最好opset 12及以上可能会遇到解析问题。dynamic_axes决定了哪些维度是动态的——如果你希望推理时batch size可以变化就必须把第0维标记为动态。但要注意动态维度越多TensorRT优化时的不确定性越大性能可能反而下降。我的建议是如果实际部署中batch size固定就不要设动态轴让TensorRT做全静态优化性能会更好。input_names和output_names必须显式指定因为TensorRT解析ONNX时会用这些名字来绑定输入输出。如果你不指定PyTorch会生成默认名字后续在TensorRT里做推理时容易搞混。3.2 导出后的ONNX模型验证导出完成不代表万事大吉必须验证ONNX模型的正确性。我见过太多次导出后精度对不上的情况问题往往出在某个算子的导出方式上。验证分两步第一步用onnx.checker检查模型结构是否合法import onnx onnx_model onnx.load(model.onnx) onnx.checker.check_model(onnx_model) print(ONNX model check passed.)第二步用onnxruntime跑一遍推理和PyTorch的输出做对比import onnxruntime as ort import numpy as np # PyTorch输出 with torch.no_grad(): torch_output model(dummy_input).cpu().numpy() # ONNX Runtime输出 ort_session ort.InferenceSession(model.onnx) ort_inputs {ort_session.get_inputs()[0].name: dummy_input.cpu().numpy()} ort_output ort_session.run(None, ort_inputs)[0] # 对比 np.testing.assert_allclose(torch_output, ort_output, rtol1e-3, atol1e-5) print(Outputs match.)如果这一步报错说明导出过程有问题需要回头检查模型里是否有TensorRT不支持的算子或者某些算子的属性在导出时丢失了。常见的坑包括自定义算子没有注册ONNX导出规则、某些归一化层的epsilon值被忽略、reshape操作在动态轴下行为不一致等。注意rtol和atol的阈值不要设得太松否则可能掩盖真实的精度问题。一般rtol1e-3、atol1e-5是比较合理的范围。4. ONNX转TensorRT解析、优化与序列化4.1 用ONNX Parser构建TensorRT网络TensorRT 7.0提供了trt.OnnxParser来解析ONNX模型。整个流程分为四步创建builder、创建network、解析ONNX、构建engine。代码如下import tensorrt as trt TRT_LOGGER trt.Logger(trt.Logger.WARNING) EXPLICIT_BATCH 1 (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) def build_engine(onnx_file_path, engine_file_path, max_batch_size1): with trt.Builder(TRT_LOGGER) as builder, \ builder.create_network(EXPLICIT_BATCH) as network, \ trt.OnnxParser(network, TRT_LOGGER) as parser: builder.max_batch_size max_batch_size builder.max_workspace_size 1 30 # 1GB builder.fp16_mode True # 开启FP16精度 # 解析ONNX with open(onnx_file_path, rb) as model: if not parser.parse(model.read()): print(Failed to parse ONNX file.) for error in range(parser.num_errors): print(parser.get_error(error)) return None # 构建engine engine builder.build_cuda_engine(network) if engine is None: print(Failed to build engine.) return None # 序列化保存 with open(engine_file_path, wb) as f: f.write(engine.serialize()) return engine这段代码里有几个参数直接决定最终性能。max_workspace_size是TensorRT在优化过程中可以使用的显存上限设得太小会导致某些优化策略无法启用设得太大则浪费显存。1GB是一个比较稳妥的起点如果你的模型很大或者层融合需求复杂可以调到2GB甚至4GB。fp16_mode开启后TensorRT会尝试用FP16精度做推理速度通常能提升30%到50%但精度可能会有轻微下降。如果你的任务对精度极其敏感比如某些医学图像分割建议先用FP32跑通再对比FP16的精度损失是否可接受。EXPLICIT_BATCH这个标志在TensorRT 7.0里非常重要。如果不加这个标志网络会以隐式batch模式构建动态shape的支持会受限。加上之后batch维度被当作显式维度处理配合ONNX里的dynamic_axes设置才能实现真正的动态batch推理。4.2 解析失败的常见原因与排查思路ONNX Parser报错是转换过程中最让人头疼的环节。错误信息往往很模糊比如“Unsupported operation”或者“Assertion failed”。根据我的经验解析失败通常有以下几个原因算子不支持。TensorRT 7.0支持的ONNX算子列表是有限的一些较新的算子如NonMaxSuppression的某些变体、Einsum等可能不被支持。排查方法是把ONNX模型用Netron可视化逐个检查算子类型然后对照TensorRT官方文档的算子支持列表。opset版本过高。前面提到过TensorRT 7.0对opset 11支持最好。如果你的PyTorch版本较新默认导出的opset可能是12或13这时候需要在torch.onnx.export里显式指定opset_version11。动态轴设置不当。如果ONNX里某些维度是动态的但TensorRT构建engine时没有对应的optimization profile解析会失败。解决办法是在构建engine时设置profileprofile builder.create_optimization_profile() profile.set_shape(input, min(1, 3, 224, 224), opt(4, 3, 224, 224), max(8, 3, 224, 224)) config builder.create_builder_config() config.add_optimization_profile(profile)权重形状不匹配。有些模型在导出时会把某些层的权重固化如果TensorRT解析时发现权重形状和预期不符也会报错。这种情况通常需要回到PyTorch侧检查模型定义。提示解析失败时把TRT_LOGGER的级别调到trt.Logger.VERBOSE可以看到更详细的日志有助于定位问题。5. 推理阶段CUDA上下文与内存绑定5.1 创建ExecutionContext并分配显存Engine构建好之后真正的推理需要通过ExecutionContext来执行。每一步都涉及显存的分配和绑定稍有不慎就会报“invalid argument”或者“out of memory”。下面是一个完整的推理示例import pycuda.driver as cuda import pycuda.autoinit import numpy as np def infer(engine, input_data): context engine.create_execution_context() # 分配显存 d_input cuda.mem_alloc(input_data.nbytes) output_shape (input_data.shape[0], 1000) # 假设输出是1000类 d_output cuda.mem_alloc(np.prod(output_shape) * 4) # float32 # 拷贝输入到显存 cuda.memcpy_htod(d_input, input_data) # 执行推理 bindings [int(d_input), int(d_output)] context.execute(batch_sizeinput_data.shape[0], bindingsbindings) # 拷贝输出回主机 output np.empty(output_shape, dtypenp.float32) cuda.memcpy_dtoh(output, d_output) return output这里有几个细节值得注意。bindings的顺序必须和engine的输入输出顺序一致可以通过engine.get_binding_index(name)来确认。batch_size参数在context.execute里传入如果engine是用EXPLICIT_BATCH模式构建的这个参数必须和实际输入数据的batch维度匹配。显存分配建议用cuda.mem_alloc而不是cuda.malloc前者是设备端分配后者是主机端锁页内存用途不同。5.2 性能调优的几个实操方向Engine构建完成后性能调优还有空间。第一个方向是调整workspace size。有些层融合策略需要较大的临时显存如果workspace不够TensorRT会放弃这些优化。你可以逐步增大workspace观察engine构建时间和推理延迟的变化找到一个平衡点。第二个方向是尝试不同的精度模式。TensorRT 7.0支持FP32、FP16和INT8三种精度。FP16通常是最佳选择速度提升明显且精度损失可控。INT8需要校准数据集流程更复杂但如果你的任务对延迟极其敏感INT8能带来数倍的提升。第三个方向是层融合与kernel自动调优。TensorRT会自动做层融合比如ConvBNReLU合并成一个kernel但融合效果受网络结构影响。你可以通过builder.config里的set_tactic_sources来限制或扩展调优策略不过大多数情况下默认设置已经足够好。注意每次修改builder参数后都需要重新构建engine这个过程可能耗时几分钟到几十分钟不等取决于模型大小和workspace设置。建议把最优参数记录下来避免重复试错。6. 那些文档里不会写的踩坑记录6.1 动态shape导致的engine构建失败我第一次做动态batch转换时ONNX里设置了dynamic_axes但构建engine时没有配置optimization profile结果build_cuda_engine直接返回None。日志里只写了一句“Profile has no shapes”排查了半天才意识到问题。后来总结出一个规律只要ONNX里有动态维度就必须在builder里设置对应的profile且profile的min/opt/max必须覆盖实际推理时可能出现的所有shape。另一个相关的坑是profile里的opt shape会影响TensorRT的kernel选择。如果你把opt设成(1,3,224,224)但实际推理时batch size经常是8那TensorRT选择的kernel可能不是最优的。建议把opt设成实际业务中最常见的shape。6.2 精度对不上的排查链路精度问题是转换过程中最隐蔽的坑。PyTorch输出和TensorRT输出对不上可能的原因有很多FP16精度损失、某个算子在转换时语义变了、动态轴下的reshape行为不一致等。我的排查链路是这样的第一步先用FP32模式构建engine排除精度损失的影响。如果FP32下输出一致那问题就出在FP16上可以尝试对某些敏感层保持FP32。第二步逐层对比输出。TensorRT支持engine.get_binding_shape和逐层调试但更实用的方法是在ONNX里插入中间输出节点然后用ONNX Runtime和TensorRT分别跑对比中间层的输出差异。第三步检查算子语义。有些算子在PyTorch和ONNX里的默认行为不同比如AdaptiveAvgPool2d在ONNX里可能被转换成GlobalAveragePool如果输入尺寸不匹配就会出错。6.3 显存泄漏与资源释放TensorRT的Python API在资源管理上不算友好如果不显式释放很容易造成显存泄漏。ExecutionContext、Engine、Builder这些对象在使用完毕后应该手动删除或者用上下文管理器。另外PyCUDA的mem_alloc分配的内存在Python GC时不一定能及时释放建议在推理循环外部分配好显存循环内复用避免频繁分配和释放。还有一个容易忽略的点多个engine共享同一个CUDA context时可能会出现资源竞争。如果你的服务需要同时加载多个模型建议每个模型用独立的context或者用TensorRT的ICudaEngine的create_execution_context为每次推理创建独立的context。7. 从转换到上线的完整检查清单把整个流程串起来之后我习惯在每次转换完成后过一遍检查清单确保没有遗漏检查项具体内容常见问题环境版本CUDA 10.2 cuDNN 7.6.5 TensorRT 7.0版本不匹配导致import失败ONNX导出opset_version11input/output names显式指定默认opset过高导致解析失败ONNX验证onnx.checker onnxruntime精度对比导出后未验证问题遗留到TRT阶段Engine构建workspace size、fp16_mode、optimization profile动态shape未设profile导致构建失败推理绑定bindings顺序、batch_size匹配、显存分配顺序错误导致输出错位精度对比FP32 vs FP16 vs PyTorch精度损失超出可接受范围性能测试延迟、吞吐、GPU利用率未做压力测试上线后性能不达标这份清单看起来简单但每一条背后都对应着至少一次踩坑经历。比如“bindings顺序”这一条我曾经因为把输入输出顺序搞反导致推理结果完全错乱排查了一整天才发现问题。又比如“optimization profile”这一条动态batch的场景下如果不设置profileengine根本构建不出来。最后分享一个实用技巧把engine构建的参数和推理的配置写成一个YAML文件每次转换时读取配置避免手动改代码出错。同时把每次转换的ONNX文件、engine文件、精度对比结果都归档保存方便回溯和对比。这套流程跑熟之后从PyTorch模型到TensorRT engine的转换基本可以在半小时内完成剩下的时间就可以花在真正的性能调优上了。
阅读完成 · 觉得有帮助？
