Pytorch 转 Onnx 动态输入时出现 if 节点问题解决方案

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/