几年前我第一次把训练好的PyTorch模型往推理环境里搬试了各种方案最后落在ONNX转TensorRT这条路上。折腾过的人都知道光把模型训出来根本不叫完事能稳定高效地跑在生产环境里才是真功夫。这几年模型部署的需求越来越多ONNX作为通用中间格式TensorRT作为GPU上的推理加速引擎这对组合几乎成了深度学习落地链路上的默认选项。如果你手头有个训练好的模型想在显卡上榨干推理性能或者正在被“模型能跑但太慢”“换了个环境就报错”之类的问题困扰这篇文章就是冲着你来的。我会把从PyTorch转ONNX、ONNX结构检查、再转TensorRT、动态shape处理到精度调优这条链路上所有值得说的细节都过一遍包括我踩过的坑和最后沉淀下来的通用流程尽量让后来的人少走弯路。1. ONNX与TensorRT到底解决什么问题1.1 先搞清模型格式和推理引擎是两回事很多刚接触部署的同学容易把ONNX和TensorRT混在一起以为“转成ONNX就等于加速了”其实完全不是一回事。ONNX是一种模型描述格式它定义了一套标准化的计算图表示方式。你可以把它理解成一份“通用说明书”不管你的模型是用PyTorch、TensorFlow还是别的框架训练的只要导出成ONNX这份说明书就能被各种推理框架读懂。它解决的是“模型怎么表达”的问题解决不了“怎么跑得快”的问题。TensorRT是运行在GPU上的推理优化引擎。它读入模型的计算图之后会把算子进行融合、精度调整、内核自动调优最终生成一个针对特定GPU型号高度定制的执行计划。它解决的是“怎么在GPU上最高效地执行”的问题。同样一个模型在TensorRT上跑和在PyTorch的CPU推理模式下跑性能差距可能是一个数量级。所以标准链路是训练框架 → ONNX统一描述 → TensorRT优化执行。前者负责互通后者负责加速两者分工明确。1.2 为什么最后选了这对组合训练生态里PyTorch占了绝大多数而PyTorch官方对ONNX导出的支持已经相当成熟大部分算子都能正确映射。反过来TensorRT对ONNX的解析路径也是官方主推的部署方式之一再加上整个社区里关于ONNX转TensorRT的踩坑记录非常多遇到问题基本都能搜到解决方案。相比之下如果你直接用PyTorch转其他推理框架或者跳过ONNX直接自定义导出格式要么是算子覆盖不全要么是性能调优成本高遇到问题还很难排查。ONNX相当于一个标准中转站让问题被隔离在“导出”和“转换”两个阶段分别处理排查起来清晰得多。2. 从PyTorch到ONNX实操全过程2.1 导出前的准备和基础参数在动手导出之前有一件很重要的事把模型切到eval模式。这一点经常被忽略如果模型里有Dropout或BatchNormtrain模式下导出的计算图会把随机行为和归一化统计量一起固化进去导致推理结果异常。PyTorch导出ONNX的核心API就是torch.onnx.export一个最小可跑的示例大概长这样import torch # 假设model是已经训练好的模型 device torch.device(cuda) model model.to(device).eval() dummy_input torch.randn(1, 3, 224, 224, devicedevice) torch.onnx.export( model, dummy_input, model.onnx, opset_version17, input_names[input], output_names[output], )这里有几个参数值得多说一句。dummy_input是必须的它用来trace模型的计算图所以要保证shape和实际推理时的输入一致。opset_version决定了导出使用的算子集版本通常建议选一个较新的稳定版本但别盲目追新因为TensorRT解析ONNX时对太高版本的算子支持不一定跟得上。我们一般在16到19之间比较稳妥具体看你的TensorRT版本支持范围。input_names和output_names给模型输入输出起了名字后面转TensorRT时要通过名字绑定输入输出张量所以起得清晰有意义会方便很多。2.2 dynamic_axes动态轴配置实际部署里很少固定batch size比如在线服务可能一次只来一个请求离线批量处理又希望尽量加大batch。如果导出时把batch维度固定成1后面想换批量大小就得重新导出非常麻烦。启动动态轴的方法是dynamic_axestorch.onnx.export( model, dummy_input, model_dynamic.onnx, opset_version17, input_names[input], output_names[output], dynamic_axes{ input: {0: batch_size, 2: height, 3: width}, output: {0: batch_size}, } )dynamic_axes里的键是张量名字值是一个字典字典的键是维度下标值是给该维度取的别名。配置完之后这个ONNX模型就支持任意batch、任意高宽的输入了。但动态轴不是免费的午餐后面转TensorRT时要给动态维度指定范围而且范围设定会影响优化效果。这一点后面在TensorRT章节细说。2.3 导出报错的常见原因和处理思路导出ONNX时最常见的报错是“算子不支持”比如某个自定义操作或者较新的PyTorch算子没有对应的ONNX映射。遇到这种情况我的处理优先级是换opset_version有时候高版本opset包含了新的算子映射。看模型里是不是有非常规操作比如复杂的控制流尝试把它改写成标准的张量运算。实在不行用torch.onnx.register_custom_op_symbolic为自定义算子注册映射或者干脆把这块逻辑拆出来放到前处理/后处理里。另外导出的模型有时候会报“output mismatch”即PyTorch前向结果和导出的ONNX结果不一致。这种问题大多是某些算子在这两种环境下浮点行为略有差异或者动态轴配置导致导出时构建了错误的计算图。验证手段是在导出后用onnxruntime跑一遍同样的输入对比输出差异判断误差是否在可接受范围内。3. ONNX模型检查与结构优化3.1 导出之后先别急着转TensorRT从PyTorch导出的ONNX往往带有不少冗余信息。比如某些维度标记不明确、中间常量节点过多、甚至会出现一些为了兼容而保留的空操作。这些节点TensorRT虽然能解析但会拖慢转换速度甚至影响最终推理性能。所以我有个习惯任何ONNX模型在转TensorRT之前都会先做一次结构检查和精简。检查用Python脚本加载ONNX模型看整体结构import onnx model onnx.load(model.onnx) onnx.checker.check_model(model) print(onnx.helper.printable_graph(model.graph))onnx.checker负责验证模型合法性printable_graph能输出计算图的可读结构方便你对模型有个直观认识。如果发现shape inference缺失还可以用onnx.shape_inference.infer_shapes补全张量维度信息。3.2 用onnxsim给模型做减法结构精简我一般直接用onnx-simplifier这个开源库它会把模型里的常量折叠、冗余节点消除、算子替换等优化统一做掉。安装和命令行用法都很简单pip install onnxsim python -m onnxsim model.onnx model_sim.onnxonnxsim执行完会输出原始节点数和简化后的节点数通常一个训练好的ResNet模型能精简掉一部分冗余节点。我第一次用的时候把两千多个节点压到了一千多后面转TensorRT的解析速度快了很多。但这里有一个要注意的点onnxsim在某些模型上可能会触发输出精度微小的变化。它在做常量折叠时如果遇到浮点计算理论上相同操作会得到相同结果但实际浮点舍入顺序可能导致极小的误差。简化完最好也跑一遍一致性对比误差在1e-4以内基本都可以接受。3.3 结构优化的一些额外手段除了onnxsim还有几个工具可以按需使用。onnx-graphsurgeon允许你对计算图做更精细的手术比如手动删除某个不需要输出的子树、修改某个节点的属性。TensorRT官方对graphsurgeon支持得很好用它能做一些onnxsim做不了的自定义优化。我通常只在onnxsim做不到的情况下才上graphsurgeon。比如有些模型前处理里带了归一化操作减均值除方差这些东西完全可以落在推理前的前处理代码里而不是浪费在TensorRT的图里解释执行。用graphsurgeon把这类算子从模型里摘掉虽然省不了太多时间但能让整个更干净。4. TensorRT转换部署主路线的选择4.1 三种落地方式怎么选TensorRT提供的主流部署方式有三种trtexec命令行工具、Python API、C API。trtexec适合做快速验证和参数压测一行命令就能从ONNX生成engine文件还会自动输出性能benchmark数据。Python API适合在Python环境下做集成尤其是模型推理部分已经用Python写好不想为了部署单独引入C代码的情况。C API适合对延迟和吞吐要求极高的线上服务省掉了Python解释开销又能更精细地控制显存和stream但开发成本也更高。我自己在生产里见过很多项目Python API足够覆盖大部分需求只有到极致优化阶段才会切换到C。所以我的建议是先用trtexec验证可行性再用Python API做业务集成最后根据瓶颈决定要不要上C。4.2 trtexec命令实战用trtexec从ONNX生成engine再简单不过/opt/tensorrt/bin/trtexec \ --onnxmodel_sim.onnx \ --saveEnginemodel.engine \ --fp16 \ --minShapesinput:1x3x224x224 \ --optShapesinput:8x3x224x224 \ --maxShapesinput:16x3x224x224这里面的参数解释一下。--fp16表示允许使用FP16精度这样很多算子会走半精度计算速度通常能提升一倍以上。--minShapes、--optShapes、--maxShapes分别指定动态轴取值的下限、最优值和上限。TensorRT会按照最优值做kernel选择所以optShapes最好设定成线上最常见的shape。trtexec跑完会在终端输出Throughput、Latency等指标。首次见到这些数字的同事往往会很兴奋但先别急着上线因为benchmark用的输入是随机生成的和真实业务数据分布有差异真实延迟还得在服务里再测一遍。4.3 Python API构建engine如果你的项目要在一个Python服务里加载模型那用Python API构建engine会更顺手。核心流程是创建builder、读取ONNX、构建网络、配置优化参数、序列化engine。一个精简的示例如下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_sim.onnx, rb) as f: ok parser.parse(f.read()) if not ok: print(Failed to parse ONNX) print(parser.get_error(0)) config builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) # 1GB workspace profile builder.create_optimization_profile() profile.set_shape(input, (1, 3, 224, 224), (8, 3, 224, 224), (16, 3, 224, 224)) config.add_optimization_profile(profile) engine builder.build_serialized_network(network, config) with open(model.engine, wb) as f: f.write(engine)构建engine的过程比较耗时尤其模型又大又复杂的时候所以实际项目里不会每次启动服务都现场构建而是把engine序列化保存到硬盘下次加载直接反序列化秒级完成。engine文件绑定的是GPU型号和TensorRT版本换卡或者升级库之后最好重新生成。这个我在生产环境里吃过亏同一个engine文件在A卡上跑得好好的换到B卡上直接报错排查了半天才发现是硬编码了计算能力。加载engine做推理的流程是反序列化engine → 创建execution context → 分配输入输出显存 → 用缓冲区指针执行enqueueV3。动态shape的输入要先通过context.set_input_shape绑定维度。engine runtime.deserialize_cuda_engine(engine_bytes) context engine.create_execution_context() # 假设是批量动态 context.set_input_shape(input, (1, 3, 224, 224)) # 分配device buffer d_input torch.empty(shape, dtypetorch.float32, devicecuda) d_output torch.empty(output_shape, dtypetorch.float32, devicecuda) context.execute_v2([d_input.data_ptr(), d_output.data_ptr()])5. 动态形状、精度校准这些绕不开的硬骨头5.1 FP16不是什么时候都能直接开trtexec里加--fp16很爽快但有些模型直接开FP16会出问题。尤其是检测、分割这类对空间细节敏感的任务FP16下某些层的相对误差会被放大最终输出可能完全不可用。我的操作习惯是先用FP16构建一个engine拿一批真实测试数据跑一遍对比FP32的原始模型输出。如果平均误差在可接受范围内再决定是否正式使用。误差对比有个简单做法算输出Tensor的平均相对误差import numpy as np def relative_error(fp32_out, fp16_out): return np.mean(np.abs(fp32_out - fp16_out) / (np.abs(fp32_out) 1e-6))如果整个输出的平均相对误差在1e-3以下通常问题不大。但对于一些本身对误差特别敏感的模型即便平均值不大个别像素点也可能偏差明显最好在业务指标层面再做一轮验证。5.2 INT8量化准备INT8能带来比FP16更极致的加速但配置成本也高它需要一张校准数据集来计算每个张量的动态范围。TensorRT的INT8流程通常分四步准备校准数据、创建校准器、设置INT8标志、构建engine。最省事的做法是让ONNX转出来的FP32模型先跑一遍校准数据在构建engine时指定一个校准器类。这个我在一些图像分类模型上做过精度损失基本可以控制在1%以内推理速度比FP32快三到四倍都是可能的。但INT8的坑也比FP16深得多比如校准数据的选择直接决定量化效果如果校准数据集和线上数据分布差异大INT8模型在真实数据上可能表现得很差。所以如果没有时间做细致调优我一般建议团队先上FP16后续有性能瓶颈再评估INT8。5.3 多profile与动态shape优化动态shape的优化范围和性能是直接相关的。范围越大TensorRT能做的kernel选择就越保守性能往往越差。所以一个常见的优化手段是不要给动态维度留一个过宽的范围而是固定业务最常用的若干组形态。如果你的业务存在几种截然不同的输入形态比如检测任务里图像有4K和1080P两种分辨率可以尝试给同一个engine添加多个optimization profile。在运行时根据当前请求的输入尺寸切换profile既照顾了不同场景又避免单个profile范围过大导致的性能损耗。# 创建两个profile的示意 profile_1 builder.create_optimization_profile() profile_1.set_shape(input, (1,3,2160,3840), (1,3,2160,3840), (1,3,2160,3840)) config.add_optimization_profile(profile_1) profile_2 builder.create_optimization_profile() profile_2.set_shape(input, (1,3,1080,1920), (1,3,1080,1920), (1,3,1080,1920)) config.add_optimization_profile(profile_2) # 运行时切换profile context.set_optimization_profile_async(1, stream)6. 实际部署中的典型问题和排查实录6.1 算子解析失败这类报错一般长这样“Unsupported ONNX node ...”。原因往往是ONNX里存在TensorRT不支持的算子或属性组合。排查时我一般先用polygraphy工具快速定位是哪一层不兼容然后决定替换算子还是用plugin实现。polygraphy是TensorRT官方配套的调试工具我很推荐每个做部署的朋友都装上。它提供了一堆子命令来跑ONNX和TensorRT的对比比如polygraphy run model.onnx --trt --fp16 --onnxrt --tol 0.01这条命令会把ONNX模型分别在TensorRT和onnxruntime上跑一遍逐层对比输出差异定位精度下降或者切图失败的具体节点。相比自己写脚本怼输入输出这个工具能省太多时间。6.2 动态shape运行时维度错误动态shape场景下最常见的错误是构建engine时没设置profile或者运行时没有预先指定输入shape。你可能会遇到“Binding input does not have a shape”或“suboptimal settings”这类报错。前者是忘了调用set_input_shape后者是当前shape偏离了optShapes太多性能会下降但不影响正确性。还有一类诡异的错误是输出buffer维度错乱原因是动态shape下输出Tensor的实际大小依赖输入维度而你不能只用最大shape去分配缓冲区。这里的处理方式是通过context.get_tensor_shape在推理前获取实际的输出shape再据此分配显存。6.3 精度下降严重时的定位手段如果FP16或INT8模型的精度掉得离谱先别急着调量化策略用polygraphy的layer-wise输出对比看看是哪几层偏差最大。我遇到过一个案例所有误差都集中在模型最后的Detection Head附近于是调整优化策略把这部分层保留为FP32精度立刻恢复。TensorRT允许在构建时通过setting per-layer precision来控制比如禁止某些层使用FP16for layer in network: if layer.name.startswith(detector): layer.precision trt.float32 layer.set_output_type(0, trt.float32)结合精度瓶颈分析这种“定点回退”的手段比全模型降精度更能兼顾性能和效果。6.4 显存占用和延迟抖动TensorRT的workspace大小直接决定它能用多少显存来做算子调优和中间结果缓存。设得太小可能导致构建失败或生成性能较差的kernel设得太大又会让多个并行实例互相抢显存。我一般从模型本身的FP32大小出发给它配2倍到4倍的显存空间再根据整机显存容量压到合适范围。延迟抖动是另一个容易被忽视的问题。GPU上多个推理请求共用同一块显存如果代码里反复创建cuda context或做显存碎片化分配很容易出现偶发的延迟尖刺。稳妥做法是推理线程只初始化一次context和buffer后续请求复用避免任何形式的显存重复分配。7. 工程化流程建议从训练到上线的一条龙7.1 转换流水线怎么设计部署久了你会发现最耗时间的往往不是“怎么转”而是“怎么保证每次转换都可靠”。我沉淀下来的转换流水线大概是这样的训练结束 → 固定模型参数和结构 → 导出ONNX → 检查结构并精简 → 用onnxruntime跑一致性对比 → trtexec构建FP32/FP16/INT8候选engine → 用真实测试集做精度和性能双验证 → 线上服务集成engine → 持续监控。每一步都有对应的检查脚本任何一步失败就终止不往下走。这样能保证交付到线上的engine一定是经过全链路验证的而不是“能跑就行”。整个流程里最容易被跳过的环节是用真实数据验证精度。很多人用随机噪点跑一下觉得输出形状和预期一样就以为没事了结果上线后发现结果一塌糊涂。真实数据集和随机输入的分布差异在FP16/INT8下会通过数值误差暴露出来所以这一步真的省不得。7.2 校验环节最容易忽略的细节做一致性对比时输入数据要覆盖边界情况。比如动态shape模型要分别测最小shape、典型shape、最大shape下的输出。一方面确认维度正确另一方面确认数值误差在每种shape下都在容忍范围内。只测一种shape给出的结论往往不可靠。另外用onnxruntime做基准对比时要注意onnxruntime和PyTorch的算子实现也有差异两者误差本身就存在所以在判定ONNX导出是否合格时我一般把threshold放宽到1e-4级别而在TensorRT的FP16对比里threshold放宽到1e-2级别。不同环节的判定标准不同死磕同一套精度只会浪费时间。7.3 监控与版本管理engine文件不像普通的模型权重那样跨平台通用它跟GPU架构强绑定。所以代码仓库里我会同时维护原始PyTorch权重、简化后的ONNX文件和构建engine的脚本engine文件本身不入库只把它作为构建产物。任何一次GPU驱动升级、TensorRT版本升级或模型更新都触发一次重新构建而不是复用旧engine。在服务上线后我会把推理耗时和显存占用都接到监控里。TensorRT本身提供了一个事件探针机制可以拿到精确的kernel执行时间这个数据比日志里手记前后时间戳要准得多排查延迟问题的时候非常有用。一点个人体会这条链路走到今天最核心的体会是模型转换不是一步到位的魔法操作而是一条需要反复校验的流水线。PyTorch导出ONNX时的参数配置ONNX简化时的精度验证TensorRT构建时的精度策略选择每一层都可能有隐形的坑。但只要你养成了“每个环节都做一致性验证”的习惯绝大多数问题都能在自己可控的范围内提前暴露而不是等上线了才被人追着问“为什么结果不对”。另外一个小建议尽量让组内的同学都熟悉polygraphy和onnxsim这两个工具这两个东西真的能节省大量手动argue模型的时间。遇到精度问题、算子问题、性能问题第一步先用polygraphy跑个对比往往比盯着报错日志猜半天下手更快。动手久了你会越来越有感觉最后基本扫一眼报错信息就能猜出问题出在哪一层。