本文最后更新于:2024年3月19日 晚上

Pytorch 模型转换为 Onnx 时,可能出现 if 节点,本文记录我遇到的这种情况的解决方案。

问题描述

Pytorch 模型导出到 Onnx 时,输入尺寸固定时没有问题,当开启动态轴(dynamic_axes)时,Onnx 中可能出现莫名其妙的 if 节点(前提是网络中没有显式写入 if 等逻辑)。

此类节点会导致 TensorRT 模型无法正常运行。

形如上图,以 Gather -> Equal -> if 结构为基本单元。

问题原因

经过排查,网络中没有写入 if 逻辑,是由于特征层中使用了 squeeze(1) 之类的算子,导致 Onnx 动态输入时需要为自己增加一些判断。

总之,问题出在 squeeze(1)

我的情况是这样,有其他情况后续补充

解决方案

squeeze 的部分换成切片式的写法即可解决类似问题。



文章链接:
https://www.zywvvd.com/notes/study/deep-learning/deploy/onnx-if-node/onnx-if-node/


“觉得不错的话,给点打赏吧 ୧(๑•̀⌄•́๑)૭”

微信二维码

微信支付

支付宝二维码

支付宝支付

Pytorch 转 Onnx 动态输入时出现 if 节点问题解决方案
https://www.zywvvd.com/notes/study/deep-learning/deploy/onnx-if-node/onnx-if-node/
作者
Yiwei Zhang
发布于
2024年3月13日
许可协议