【Bug已解决】Shape of Gather output is wrong making it unusable as K input to TopK operator 解决方案一、现象长什么样模型里有这么一段用Gather从一个常量里取出“要保留的 top-k 个数k”再把这个结果喂给TopK的第二个输入K 输入。在某些 ONNX Runtime 版本/配置下Gather的输出形状不对导致TopK直接报错或选出错误数量import onnxruntime as ort # Gather 取出 k期望输出形状 [1] 或标量但拿到 [1, 1] 或 [N, 1] sess ort.InferenceSession(gather_topk.onnx, providers[CPUExecutionProvider]) # TopK 报错K input 的形状不被接受期望 0-D 或 1-D 且 size1最小信号Gather 输出形状: [1, 1] 或 [N, 1]多了一维 TopK 期望 K 输入: 标量 或 [1] - 形状不兼容 - TopK 失败 / 选出错误数量注意这不是 Gather 算错值而是输出张量的 rank/形状不符合 TopK 对 K 输入的约束导致下游用不了。二、背景ONNX 的TopKopset 17 起K 变成图的第二个输入对 K 输入有严格要求K 必须是一个0-D标量或 1-D 且元素个数为 1的张量表示“取前 k 个”。为什么要这么严因为 K 是控制流意义上的“超参数”必须是单值。Gather的语义是output input[indices]输出形状 input.shape[:axis] indices.shape input.shape[axis1:]。如果indices本身带了一个额外维度比如indices形状是[1, 1]而不是[1]Gather的输出就会把那个多余维度带出来变成[1, 1]。问题就出在这里很多模型用Gather从一个 1-D 常量里取一个数但indices在导出时被塑造成了[1, 1]或别的带多余维度的形状于是Gather输出k成了[1, 1]。这个形状送到TopK的 K 输入ORT 的校验器拒绝或某些版本静默接受但选出错。三、根因根因是Gather的输出形状继承了indices的多余维度而该多余维度没有被消除导致输出形状不满足TopK对 K 输入“标量或 [1]”的约束indices 带多余维度导出时k的索引常量被塑成[1, 1]多了一维常见于框架导出对 scalar 的处理Gather按规则把indices.shape原样带进输出得到[1, 1]。缺少 Squeeze/Reshape模型图里没有在Gather之后接Squeeze/Reshape把[1, 1]压成[1]于是脏形状直接进TopK。TopK 校验拒绝ORT 在构造TopK节点时校验 K 输入形状发现不是标量/[1]要么报错要么在部分版本里把[1,1]当成[1]但 axis 推断错位选出错误数量。不是 Gather 值错Gather取到的k值本身是对的只是“包装”它的张量形状多了维。所以这不是数值错而是形状rank不满足下游算子约束属于图构造/形状推断的衔接问题。四、最小可运行复现下面用 NumPy 模拟“Gather 输出带多余维度导致形状不兼容 TopK”import numpy as np def gather_shape(data_shape, indices_shape, axis0): 按 ONNX Gather 规则推导输出形状。 return tuple(data_shape[:axis] tuple(indices_shape) data_shape[axis1:]) def topk_accepts_k_shape(k_shape): TopK 接受的 K 输入形状标量(空)或 [1]。 if len(k_shape) 0: return True if len(k_shape) 1 and k_shape[0] 1: return True return False if __name__ __main__: # 常量 k 来源形状 [3]indices 本应是 [1]但导出成 [1,1] bad_indices (1, 1) good_indices (1,) bad_out gather_shape((3,), bad_indices) # (1, 1) - 不兼容 good_out gather_shape((3,), good_indices) # (1,) - 兼容 print(Gather 输出(坏):, bad_out, TopK 接受?, topk_accepts_k_shape(bad_out)) print(Gather 输出(好):, good_out, TopK 接受?, topk_accepts_k_shape(good_out)) assert topk_accepts_k_shape(bad_out) is False assert topk_accepts_k_shape(good_out) is True跑出来(1,1)不被 TopK 接受、(1,)接受。这复现了“Gather 输出多一维导致 TopK 用不了”的机制。五、解决方案第一层最小直接修复最小修复在Gather之后加Squeeze/Reshape把 K 输入压成标量或[1]或者导出时把indices塑成正确的 1-D。import onnx from onnx import helper, TensorProto # 修复前Gather 输出 [1,1] 直接进 TopK报错 # 修复后Gather - Squeeze(axis[0,1] 或全 squeeze) - [1] 或标量 - TopK # 用 onnx 修改图在 Gather 与 TopK 之间插入 Squeeze def fix_graph(model_path, out_path): model onnx.load(model_path) # 伪代码找到 Gather 节点在其输出后插入 Squeeze把 [1,1] 压成 [1] # squeeze_node helper.make_node(Squeeze, [gather_out], [k_fixed], axes[0]) # 再把 TopK 的 K 输入从 gather_out 改成 k_fixed onnx.save(model, out_path)对 ORT 仓库侧也可让Gather的形状推断在遇到标量语义的 indices 时自动给出匹配形状但更稳的是在图里显式Squeeze。这一层立刻让 TopK 拿到合法 K 输入。六、解决方案第二层结构性改进把“哪些算子的输出形状必须满足下游约束如 TopK 的 K 输入”收口成唯一的配置对象OrtGatherTopkShapePolicy图校验与导出读它from dataclasses import dataclass, field from typing import Tuple, Dict dataclass(frozenTrue) class OrtGatherTopkShapePolicy: Gather-TopK 形状衔接的单一事实来源。 # TopK 接受的 K 输入形状 topk_k_accepted_shapes: Tuple[Tuple[int, ...], ...] ((), (1,)) # 需要 Squeeze 的多余维度indices 带来的 squeeze_axes: Tuple[int, ...] (0, 1) # 修复方式在 Gather 后插 Squeeze或导出时把 indices 塑成 1-D fix_strategy: str insert_squeeze_after_gather # 需要校验的算子对 checked_pairs: Tuple[str, str] (Gather, TopK) def is_k_shape_valid(self, shape: Tuple[int, ...]) - bool: return shape in self.topk_k_accepted_shapes def describe(self) - str: return Gather 取 k 后必须 Squeeze 成标量/[1] 才能喂给 TopK POLICY OrtGatherTopkShapePolicy() def validate_k_input(shape: Tuple[int, ...], policy: OrtGatherTopkShapePolicy POLICY) - bool: return policy.is_k_shape_valid(shape)所有图导出与校验读同一份POLICYGather-TopK 的形状衔接被固化不会再出现多一维。七、解决方案第三层断言 / CI 守护把“Gather 输出形状满足 TopK 约束”做成断言。下面用 pytest 风格守护复用第四节逻辑import numpy as np def test_topk_rejects_extra_dim(): assert topk_accepts_k_shape((1, 1)) is False assert topk_accepts_k_shape((1,)) is True assert topk_accepts_k_shape(()) is True def test_gather_shape_inference(policy): # 坏 indices 产生不兼容形状 assert policy.is_k_shape_valid(gather_shape((3,), (1, 1))) is False assert policy.is_k_shape_valid(gather_shape((3,), (1,))) is True def test_squeeze_axes_defined(policy): assert len(policy.squeeze_axes) 1 def test_checked_pair_is_gather_topk(policy): assert policy.checked_pairs (Gather, TopK)这四组断言锁住(1) TopK 拒绝多余维度、接受标量/[1](2) Gather 形状推断正确识别坏/好形状(3) Squeeze 轴已定义(4) 校验的算子对是 Gather-TopK。CI 跑通即代表形状衔接被守护。八、排查清单遇到 TopK 报 K 输入形状错看 Gather 输出形状是不是[1,1]/[N,1]等多余维度。看 indices 形状导出时k的索引是不是被塑成了[1,1]。插 Squeeze/Reshape在 Gather 后把 K 压成标量或[1]。或修导出让 indices 直接是 1-D[1]。查 TopK 版本opset 17 的 K 是第二输入形状约束更严。统一策略对象用OrtGatherTopkShapePolicy固化。CI 守护断言 Gather 输出形状满足 TopK 约束。九、小结Shape of Gather output is wrong making it unusable as K input to TopK operator的根因是Gather取k时indices常量带了多余维度如[1,1]Gather按规则把indices.shape带入输出得到形状[1,1]而TopK的 K 输入要求标量或[1]于是形状不兼容TopK 报错或选出错误数量。最小修复是在Gather之后加Squeeze/Reshape把 K 压成标量或[1]或导出时把 indices 塑成 1-D结构性改进是用唯一的OrtGatherTopkShapePolicy固化形状衔接CI 用四组断言守护“TopK 拒绝多余维度、Gather 形状推断正确、Squeeze 轴定义、校验对是 Gather-TopK”。记住Gather 的输出形状会继承 indices 的维度喂给有形状约束的下游算子前必须先 Squeeze。
