最近在部署一个分割模型(DeepLabV3+,backbone是ResNet50),用TensorRT加速。按教程先转ONNX,再转engine。结果发现转出来的ONNX用onnxruntime跑CPU,比原PyTorch模型直接跑还慢20%左右。我确认了opset版本(12),也试了dynamic_axes,输入输出都命名了,但速度就是上不去。是不是我漏了某些优化pass?或者量化、剪枝这类操作必须在ONNX之前做?另外,转出来的图里有一堆我没见过的op(比如Resize、Pad这种),会不会是这些op在ORT里没有高效实现?有没有有经验的老哥分享下你们从PyTorch到ONNX的常规优化流程,或者踩过的坑?谢谢了。
PyTorch转ONNX后推理速度反而变慢了,是我的导出姿势不对吗?
全部回复
共 17 条说实话你这个对比基准本身就有点吃亏,PyTorch跑CPU是有MKL-DNN这类底层库优化的,而ONNX Runtime如果没装对执行 Provider(比如默认CPU EP没开特定的图优化),某些算子的调度效率确实还不如PyTorch原生路径。Resize和Pad这种op在ONNX里很常见,但ORT对它们的实现分版本差别挺大,有时候转成静态shape反而能触发更好的内核融合。
我建议你先别纠结ONNX的CPU速度,直接跳到TensorRT那步看engine延迟,因为部署目标本来就是GPU。如果TRT也慢,那大概率是模型里的某些自定义操作(比如ASPP里的空洞卷积)在转换时被拆成了低效子图,这时候可以试下torch2trt或者直接导出成TRT的plugin格式,绕过ONNX的中间层抽象。
另外你提到的量化剪枝确实应该在导出前做,因为ONNX只是计算图描述,它不负责帮你做数值压缩。我自己的经验是,如果追求极致推理,别指望ONNX那条链路能自动优化出什么花样,通常我会先做torch.jit.trace,再用手动改图的方式把一些固定shape的resize和pad替换成更简单的实现。你可以看看onnxsim这个工具,它能做常量折叠和op融合,至少能把图里的冗余节点清掉一批。还有就是检查下你的推理线程数设没设对,ORT默认线程数和物理核数不匹配时,慢20%是很正常的现象。
ort跑cpu跟pytorch比没啥参考价值,pytorch本身就有不少融合优化,onnx导出后很多算子被拆碎了反而吃亏。你看到的Resize、Pad这类op确实容易成瓶颈,尤其Resize在onnx里默认用float坐标,ort对它的实现未必比得上pytorch原生的插值。真要对比速度,建议直接用trt engine跑gpu,cpu这环基本可以跳过。另外onnx转trt时注意下trt的版本和onnx parser兼容性,有时候模型图里有不支持的模式会回退到性能很差的实现,你可以用trtexec加--verbose看看有没有warning。量化剪枝这些确实最好在导出前做,但你这情况更像图优化没吃透,可以试试onnxsim精简一下图,再把batch维固定住看看。
说实话你这个对比基准本身就不太公平,PyTorch跑GPU,ONNX用CPU测,这俩根本没可比性。如果你真想验证转换有没有引入额外开销,应该拿同一张卡同一个batch size去比PyTorch和ORT的GPU推理,或者干脆直接对比TensorRT engine和PyTorch的耗时。至于Resize和Pad这些op,确实在ORT CPU上有些实现不是走最优化路径的,尤其是当输入分辨率不是固定值的时候,动态shape会触发很多额外的内存拷贝和算子调度开销,建议你固定输入尺寸再试一次,很多情况下能快回来不少。
另外关于优化pass这事,onnxruntime本身有graph优化选项,默认是开basic的,但你得确认一下是不是用了ENABLE_ALL,有时候某些op融合需要手动开启。不过更关键的还是你导出时的opset和算子选择,像DeepLabV3+这种带空洞卷积的结构,建议导出前把backbone里的BatchNorm全部fold掉,顺便用torch.onnx.export的operator_export_type切到ONNX_ATEN_FALLBACK试试看能不能精简一些冗余节点。量化这些确实得在onnx之前做,但你现在还没到那步,先把静态shape的engine跑通看下真实瓶颈在哪,我怀疑你转完TRT之后速度应该能反超PyTorch不少,如果还慢那大概率是TRT版本对某些算子的支持问题。
ONNX Runtime对动态shape支持差,固定尺寸重导出试试,能快不少。
说实话CPU上对比ORT和PyTorch没啥意义,PyTorch走的是自己的推理路径还有MKL加持,ORT在CPU上的优化本来就不一定占优,你直接拿TensorRT的engine跟原始PyTorch比才有参考价值。另外Resize、Pad这些op在ONNX里是标准op没错,但ORT对它们的实现确实不如TensorRT的plugin高效,建议你转engine的时候开trt的fp16或者int8,顺便把onnx-simplifier跑一遍,很多冗余shape计算能被消掉。
你这速度变慢大概率不是导出姿势的问题,而是比较基准选错了,ONNX只是个中间表示,真正加速得看后端。我之前转过Mask R-CNN也遇到类似情况,后来发现是动态shape导致ORT每次推理都要重新做内存规划,你把input固定成实际部署的尺寸试试,静态shape下ORT的CPU性能通常会好不少。量化剪枝那些确实得在导出前做,但你这个场景先别急着上,把TensorRT的engine跑起来看下实际延迟再说。
ONNX Runtime在CPU上跑得比原生PyTorch慢其实挺常见的,别急着怀疑自己导出姿势有问题。PyTorch的CPU后端用了MKL-DNN(现在叫oneDNN),对卷积、BN融合这些做了非常激进的图优化,而ORT的CPU EP虽然也支持这些,但默认的图优化级别可能没开满,你可以试试把graph_optimization_level设成ORT_ENABLE_ALL,再把intra_op_num_threads手动对齐物理核心数,有时候提升很明显。至于那些Resize、Pad的op,确实是个坑,ORT对某些插值模式的Resize实现不如PyTorch的native kernel高效,尤其是双线性上采样在分割模型里到处都是,累积起来就很可观了。量化剪枝这些最好在PyTorch侧做完再导出,ONNX的量化工具链虽然能用但坑更多,而且动态量化对分割这种密集预测任务收益有限。如果你最终目标是TensorRT,其实可以跳过ONNX Runtime的CPU benchmark,直接看engine在GPU上的表现,ONNX只是个中间格式,CPU慢不代表GPU慢。另外建议用onnxsim先过一遍常量折叠和算子融合,再把opset升到13或17,有些op的新版本实现会好很多。
ONNX在CPU上跑本来就不一定比PyTorch快,ORT的CPU EP对Resize、Pad这类动态shape算子确实优化一般,尤其你如果没固定输入尺寸,它会走很多fallback。建议先用onnxsim做一遍常量折叠和算子融合,再试试固定batch和H/W,把dynamic axes去掉。量化剪枝这些最好在PyTorch侧做完再导出,ONNX里再搞很容易翻车,而且TensorRT才是真正提速的地方,ORT CPU慢20%其实挺正常的。
学到了,感谢分享!
ONNX在CPU上比PyTorch慢挺正常的,ORT的CPU EP对Resize、Pad这类动态shape的op确实优化一般,尤其你分割模型里插值操作多,很容易拖后腿。转TensorRT的话其实不用太纠结ONNX在CPU上的速度,直接看engine在GPU上的表现就行。另外量化剪枝建议在PyTorch端做完再导出,ONNX里再搞这些会很折腾。你可以试试固定输入尺寸导出,把dynamic_axes去掉,往往能快不少。
ONNX在CPU上慢太正常了,ORT对CPU的优化本来就有限,尤其Resize和Pad这类动态shape的op,实现效率跟PyTorch的cudnn kernel没法比。你真正该看的是转成TensorRT engine后的速度,ONNX只是中间格式,别拿它跟原模型比CPU推理。另外DeepLabV3+里那些插值和padding,在TRT里也容易掉到slow path,建议导出时固定输入尺寸,关掉dynamic_axes试试。量化剪枝一般放在ONNX之后做,比如用TRT的INT8校准,提前在PyTorch里搞反而容易丢精度。
ONNX转出来op变多挺正常的,ORT对Resize这类确实优化一般,你可以试试用onnx-simplifier先过一遍再转engine。
CPU上ORT不一定比PyTorch快,先看有没有开图优化和线程数,不然白搭。
这个其实挺常见的,尤其分割模型里Resize和Pad这类op确实容易拖后腿,ORT的CPU EP对动态shape的Resize支持不算特别好,经常退化成比较慢的实现。你opset 12没问题,但ONNX Runtime跑CPU本身就不是强项,拿它跟PyTorch的CPU比其实参考意义有限,PyTorch底层也是oneDNN,而且很多op是融合过的。真正该看的是转成TensorRT engine之后的延迟,ONNX那一步慢20%不用太焦虑。另外dynamic_axes设了之后,shape变成动态的,ORT和TRT都会丢掉不少针对固定shape的优化,如果部署时输入尺寸固定,建议直接写死shape导出,别开动态轴。量化剪枝这些最好在PyTorch侧先做完再导出,ONNX层面的量化工具链虽然能用,但踩坑概率高不少。还有一点,导出时试试torch.onnx.export里的do_constant_folding,能把一些BN和Conv折进去,图会干净一些。Resize、Pad这些其实可以手动改写成固定尺寸的插值,或者干脆在网络里换成更友好的上采样方式,TRT对固定scale的Resize支持就好很多了。
ONNX在CPU上比PyTorch慢其实挺常见的,尤其分割模型这种带大量Resize和插值的结构。PyTorch的CPU推理走的是MKLDNN+自己的JIT融合,很多op都有针对性优化,而ORT的CPU EP虽然也有优化,但对某些动态shape或者非标准op的处理就没那么激进了。你看到的Resize、Pad大概率就是问题所在,DeepLabV3+里的双线性插值在ONNX里经常被拆成多个op,ORT执行起来开销不小。另外dynamic_axes如果没实际需要就别开,动态维度会让ORT没法做静态内存规划和kernel特化,速度掉得很明显。量化剪枝这些确实建议在PyTorch侧先做完再导出,ONNX的量化工具链对分割模型支持一般,容易掉点。你可以先用onnxsim做一下常量折叠和算子融合,再看看ORT的profiling看时间花在哪几个节点上。如果目标就是TensorRT,其实可以试试torch-tensorrt直接编译,绕开ONNX这一层,省心不少。
ONNX在CPU上比原生PyTorch慢其实挺常见的,不一定是你导出姿势的问题。PyTorch的CPU后端有MKL和oneDNN加持,很多算子融合和线程调度是默认开好的,而onnxruntime的CPU EP虽然也有优化,但默认配置下不一定能跑到最优,你可以试试调intra_op_num_threads和开graph_optimization_level为ORT_ENABLE_ALL。至于那些Resize、Pad,确实是PyTorch导出时把一些语义拆成了细粒度算子,ORT里对Resize的实现分很多种模式,如果coordinate_transformation_mode没匹配好,性能会差不少,可以拿netron看一下这些op的属性。量化剪枝这些一般建议在ONNX图上做,PyTorch侧做QAT再导出也行,但纯PTQ放到ONNX后用ORT的量化工具更顺手。另外CPU上跑ONNX和TensorRT完全是两条路,你要上TRT的话其实不用太纠结ONNX在CPU上的速度,直接看engine在GPU上的表现更实际。如果非要在CPU上比,可以试试导出时把opset升到13以上,或者用torch.jit.trace配合optimize_for_inference先过一遍。
有没有更详细的教程推荐?
写得挺好,建议补充一些性能数据。