一、背景

在模型开发阶段，我们经常会用到 Gather 算子 来做张量的索引操作
在我的模型中，输入是一个五维张量，形状为：(3, 2, 4, 5, 6)
模型结构里使用了Gather(axis=0)，把这个输入在第 0 维上拆成三份，每一份的形状是：(2, 4, 5, 6)
在 PC 上推理没有任何问题，但当模型部署到 J5 开发板上时遇到了算子支持问题

报错信息：

二、问题描述

在 J5 的算子支持约束下，Gather 算子对输入张量维度有限制。
我这里的输入是五维： (3, 2, 4, 5, 6)

直接使用 Gather 会报错，模型无法正常运行。

换句话说，硬件推理引擎不支持这种用法，必须想办法用受支持的算子来实现相同的逻辑。

三、思路分析

回顾 Gather 的作用：

Gather(axis=0, index=i) 等价于 在第 0 维取出第 i 份子张量。
输入 (3, 2, 4, 5, 6)，取 i=0,1,2，得到三份 (2, 4, 5, 6)。

如果从功能上看，其实就是 在第 0 维上做切片。
因此完全可以用 Slice + Squeeze 的组合来替代：

Slice
- 参数：starts=[i], ends=[i+1], axes=[0], steps=[1]
- 输出形状：(1, 2, 4, 5, 6)
Squeeze
- 去掉第 0 维（长度为 1）
- 得到目标形状 (2, 4, 5, 6)

这样就和原来的 Gather 完全等价。

四、解决方法

通过这个脚本就可以将onnx模型中的所有 Gather 算子 用 Slice + Squeeze 替换

五、效果展示

下面给出 Gather 替换前后的张量切片效果对比：

原始 Gather	替换为 Slice+Squeeze

说明：

左图为原始模型使用 Gather(axis=0) 对输入 (3, 2, 4, 5, 6) 的切片输出，每次取第 0 维的一份 (2, 4, 5, 6)。
右图为替换后模型使用 Slice + Squeeze 实现同样的功能，输出形状与原始完全一致。
验证通过 ONNX Runtime 测试，保证替换后模型输出数值等价。

J5 Slice + Squeeze 替换 Gather 算子调试记录

一、背景

二、问题描述

三、思路分析

四、解决方法

五、效果展示