先说结论这个报错不是因为你代码写得有问题而是因为你在把 PyTorch 模型导出成 ONNX 时模型里某个操作触发了aten::index_put_算子而 PyTorch 的 ONNX 导出器对这个算子的支持非常有限。我前阵子做一个需要部署到服务端的文本分类模型时也卡在这个报错上整整一个晚上。报错原文大概是这样的RuntimeError: Only consecutive 1-d tensor indices are supported in exporting aten::index_put to ONNX如果你也遇到了千万别急着重写模型。这篇文章我就用自己实际踩坑的经历把这个报错从根因到替代方案完整梳理一遍最后还会给你一份排查清单照着做基本能解决大部分类似问题。1. 先搞明白报错在说什么ONNX 导出过程的花式拒绝我最早看到这个报错时第一反应是去搜aten::index_put是什么意思。这个算子的全称是 ATen 里的index put操作底层对应的就是你写x[mask] value或者x[indices] value这类按索引赋值的代码。举个最简单的例子x torch.zeros(5) mask torch.tensor([True, False, True, False, True]) x[mask] 1.0这两行代码在 PyTorch 内部就会触发index_put_带下划线表示 in-place 操作。平时训练、推理都没问题但在导出 ONNX 时导出器需要把 PyTorch 算子逐个映射到 ONNX 算子。问题就出在这个映射上。ONNX 里其实是有ScatterElements和ScatterND这类算子的理论上可以表达按索引赋值。但 PyTorch ONNX 导出器对index_put_的支持只覆盖了一个非常狭窄的场景连续的 1 维索引。什么叫连续 1 维索引就是x[2:5] value这种切片式赋值索引是一段连续区间。报错里提到的 consecutive 1-d tensor indices指的就是x[start:end]或者x[[start, start1, ..., end-1]]这类可以用连续区间表示的下标。而你一旦写了x[mask] value这种布尔掩码赋值或者x[[1, 5, 7]] value这种任意位置组合的赋值导出器就直接撂挑子不干了。1.1 为什么模型里会莫名其妙出现这个算子这是最坑的地方。大部分情况下你不会直接在模型前向代码里写x[mask] value但它会通过其他高层 API 间接触发。我那次就是在一个自注意力模块里用了logits.masked_fill_(mask, -1e9)本意是把注意力分数里某些位置填充成一个极小的数。你猜怎么着masked_fill_在底层也会走index_put_的路径。类似的 API 包括但不限于Tensor.masked_fill()/masked_fill_()Tensor.index_fill()/index_fill_()Tensor.scatter()/scatter_()自定义损失函数里的output[target] 0这类写法某些第三方库的 mask 操作比如transformers里部分填充逻辑所以你会发现一个问题代码里明明没写任何危险操作导出时照样报这个错。原因就是高层 API 隐藏了底层算子的真实面貌报错时你根本不知道是哪一行代码闯的祸。1.2 这个报错的时间线为什么以前没事现在出事还有一个现象值得注意同一个模型不同 PyTorch 版本下导出结果可能完全不同。我试过的版本差异大致是这样PyTorch 版本ONNX 导出行为1.8 及更早部分场景直接报错或导出成不正确的图1.9 ~ 1.12严格检查遇到非连续索引直接报 RuntimeError2.0 以上TorchScript 路径行为不稳定有的能导出但推理结果错误2.1dynamo 路径默认导出器仍受限但暴露了更多内部细节这个报错在 1.x 时代就存在到现在也没彻底消失根本原因是 ONNX 的算子集和 PyTorch 的算子集本来就不是一一对应的。PyTorch 的算子粒度很细index_put_这种写入 索引组合 的语义在 ONNX 里没有天然对应的单一算子只能靠导出器做模式匹配。匹配不上就报错简单粗暴。所以你真正要做的不是研究怎么让导出器支持index_put_而是把模型里的索引赋值操作改写成 ONNX 能理解的形式。2. 定位肇事代码三步锁定元凶行光知道原理还不够你得知道具体是哪一行代码在导出时炸的。我那次排查花了大半晚上最后总结出了一套定位流程效率高很多这里分享给你。2.1 第一步用最小复现脚本确认触发条件不要一上来就在完整模型上排查先写一个最小脚本确认报错能否稳定复现同时验证你的修改方向。比如import torch import torch.onnx class IndexPutModel(torch.nn.Module): def forward(self, x, mask): # 触发 index_put_ 的操作 x x.clone() x[mask] 0.0 return x model IndexPutModel() dummy_input torch.randn(8, 10) dummy_mask torch.tensor([True, False, True, False, True, False, True, False]) torch.onnx.export( model, (dummy_input, dummy_mask), test.onnx, opset_version13, input_names[x, mask], )在我印象里这个脚本稍微改一改比如把x[mask] 0.0换成x[0:3] 0.0就能正常导出。用这个最小脚本反复验证比在完整模型上瞎猜高效得多。2.2 第二步用二分注释法缩小范围如果最小复现脚本没报错说明问题出在模型内部的某个模块里。这时我惯用的方法非常土但非常有效先把模型整体的导出脚本写好然后注释掉模型代码的一半导出试试如果还报错再注释掉剩下的一半的一半。具体操作把model.forward()里每个大模块单独导出一次找到报错的那个模块。进入报错模块逐个子模块嵌套导出继续二分。找到具体模块后再检查模块内的数据流搜索masked_fill、scatter、index_fill、[...] ...这类关键字。最后锁定可疑代码后把输入固定好单独跑一遍导出确认报错确实由它引起。这个方法听着笨但每次都很灵验。尤其是模型结构复杂、嵌套了好几层nn.Module的时候靠肉眼直接看代码经常会漏掉藏在某个子模块里的暗桩。2.3 第三步利用报错堆栈定位有时候报错信息会附带 Python 调用栈虽然不精确但能指出方向。比如堆栈里出现File torch/onnx/utils.py, line XXXX, in _export graph, _, _ model(*args) File torch/onnx/symbolic_opset16.py, line XXX, in index_put如果你的代码没有任何封装、纯nn.Module直接导出那堆栈里会出现model.forward的调用位置顺着找就行。但注意如果模型里用了torch.jit.trace或torch.jit.script堆栈可能被消化掉只显示torch/onnx内部的帧这时就必须回到二分法。补充一个我后来养成的习惯写一个公共导出脚本支持传入模块名字符串自动尝试导出指定子模块这样定位下一轮问题时不用反复改代码。3. ONNX 友好改造的核心方案把索引赋值改成 torch.where 组合定位到肇事代码之后接下来的问题就是怎么改。先说结论在绝大多数场景下用torch.where替代索引赋值是改动最小、语义最接近、导出成功率最高的方案。3.1 torch.where 的替代逻辑假设你原来的代码是x[mask] value它表达的意思是把x中mask为真的位置替换成value其余位置保留原值。这本质上就是一个逐元素的条件选择。而torch.where干的正是这件事x torch.where(mask, value, x)torch.where(condition, a, b)对每个位置做判断condition为真取a为假取b。当a是一个标量时PyTorch 会自动广播到和x相同的形状。ONNX 导出器对torch.where的支持非常成熟因为它直接映射到 ONNX 的Where算子。以我那个masked_fill_(mask, -1e9)为例改写前是logits.masked_fill_(mask, -1e9)改写后是logits torch.where(mask, torch.full_like(logits, -1e9), logits)如果mask的 shape 和logits不一致torch.where的广播机制也能处理只要它们能广播到同一形状。这一点比直接索引赋值还灵活。3.2 几种典型场景的具体改写模板我根据自己实际遇到的情况整理了几个高频场景的改写模板。场景一标量替换原代码x[mask] 0.0改写x torch.where(mask, torch.zeros_like(x), x)注意zeros_like保持了和x完全相同的 dtype 和 device避免类型不匹配。场景二用另一个张量的值替换原代码x[mask] y[mask]改写x torch.where(mask, y, x)只有mask为真的位置才取y的同位置值其余保留x原值。如果y和x形状不同但可广播也能正常工作。场景三二维索引赋值原代码x[:, idx] scalar_value这在 NLP 模型里很常见比如把某个 token 位置的 logit 强制设为极大或极小值。改写时可以先用torch.arange构造一个 maskrow_indices torch.arange(x.shape[1], devicex.device) mask (row_indices idx) # shape: [seq_len] x torch.where(mask.unsqueeze(0), torch.full_like(x, scalar_value), x)如果要对 batch 内每个样本的不同位置赋值那就把 mask 构造成[batch_size, seq_len]的形状torch.where一样能处理。3.3 改写后的边界检查改写不是替换完就完事有几个边界问题你必须检查第一个dtype 一致性。torch.full_like(x, value)里的value会按x的 dtype 解释。如果x是 int64、而你想填的value是负数-1e9那结果会溢出或直接截断。所以填充大负数时先确认x是浮点类型。第二个in-place 语义变化。原来的x[mask] value是 in-place 操作x本身被修改改写后x torch.where(...)会创建一个新张量。如果你的代码是logits.masked_fill_(mask, -1e9)这种带下划线的 in-place 版本改写时记得把后续所有对logits的引用都理解成指向新张量不要在前面保留旧引用否则后续计算用的还是旧值。第三个mask 的 dtype。torch.where的condition必须是 bool 张量。如果你的 mask 是 0/1 的 float 张量有些人用(x 0).float()来生成 mask需要先转成 boolmask_bool mask.to(torch.bool) # 或者 mask 0这个细节很容易被忽略踩过一次你就记住了。4. 其他可用的替代方案masked_fill、scatter、自定义算子torch.where不是唯一的路。根据具体场景下面几种方案也值得掌握。4.1 masked_fill 在某些场景下依然可用你可能会问既然masked_fill_触发index_put_那masked_fill不带下划线呢我实测下来的结论是部分场景可以导出但稳定性没有保证。在特定 opset 版本、特定 mask 形状下masked_fill的导出路径可能走的是Where而不是index_put_这时没问题。但换个 mask 形状可能又打回原形报index_put_的错误。所以我的建议是如果你已经在排查这个报错了就直接用torch.where别在masked_fill和masked_fill_之间反复横跳浪费时间。你的目标是成功导出不是帮 PyTorch 测试导出器的算子覆盖范围。4.2 scatter / index_fill 系列能规避就规避scatter、index_fill、index_copy这几个 API 底层同样会进入index_put_的领地ONNX 导出支持也分场景scatter在 opset 11 里有ScatterElements对应但语义不完全一致容易出岔子。index_fill的导出支持也很窄只处理连续索引等简单情形。index_copy在多数版本里导出直接失败。如果你真的需要按索引把某个位置替换成固定值还是建议回到torch.where方案。4.3 自定义算子最后的手段有些场景确实绕不开index_put_的语义比如你要在 ONNX 图里对一个动态位置做复杂写入。这时就得上自定义算子。完整流程如下在 PyTorch 里用torch.library注册一个自定义算子。编写该算子的 symbolic function告诉 ONNX 导出器映射到哪个 ONNX 算子比如ScatterND。在推理引擎如 ONNX Runtime里注册对应的 custom op 实现。这个方案工程量不小还要在部署端配套开发不到万不得已我不推荐。如果你只是导出模型给 TensorRT、ONNX Runtime 这类常见引擎用torch.where几乎总能帮你绕过去。4.4 新版导出路径 dynamoTrue 的差异PyTorch 2.x 提供了基于 dynamo 的新导出路径写法是torch.onnx.export(model, args, model.onnx, dynamoTrue)这个新路径理论上对更多算子有更完整的支持但我实测下来index_put_的限制仍然以不同形式存在。有时候报错信息会从RuntimeError变成别的比如Unsupported: ONNX export of operator index_put本质没变只是换了层皮。所以在导出用于生产环境的 ONNX 模型时我还是以修改模型代码让它对 ONNX 友好为主不指望导出器突然变得万能。5. 避坑清单ONNX 导出前先自查这几类操作经历了那一个晚上的折腾我后来养成了一个习惯在写任何要导出 ONNX 的模型前先自查代码里是否有下面这几类雷区操作。提前规避比事后排查省太多时间。5.1 高频雷区操作清单操作类型示例代码ONNX 导出风险布尔掩码赋值x[mask] value极高任意索引赋值x[[1,3,5]] value极高二维行列索引赋值x[:, idx] value极高masked_fill_x.masked_fill_(mask, value)高取决于 mask 形状scatter_x.scatter_(dim, index, src)高opset 版本敏感index_fill_x.index_fill_(dim, index, value)高in-place 索引更新x[condition] 1极高如果你写完代码发现上面这些一个都没有恭喜你踩雷概率大幅降低。5.2 我常用的导出前扫描脚本为了不每次都手动检查我写过一个简单的 AST 扫描脚本代码量不大但很实用。原理是遍历 Python 源码的 AST找出所有SubscriptAssign的组合以及masked_fill、scatter、index_fill等函数调用。精简版思路如下import ast class IndexPutVisitor(ast.NodeVisitor): def visit_Assign(self, node): # 检查 x[mask] value 这类赋值 if isinstance(node.targets[0], ast.Subscript): print(f可疑索引赋值: {ast.dump(node)}) self.generic_visit(node) def visit_Call(self, node): # 检查 masked_fill / scatter / index_fill 调用 if isinstance(node.func, ast.Attribute): name node.func.attr if name in {masked_fill, masked_fill_, scatter, scatter_, index_fill, index_fill_, index_put_}: print(f可疑API调用: {name}) self.generic_visit(node) with open(model.py) as f: tree ast.parse(f.read()) IndexPutVisitor().visit(tree)这个脚本的作用不是拦截所有问题而是帮你快速圈定可疑区域。排查时直接跳转到这些行附近检查不用满文件找。5.3 导出前数值验证避免导出成功但结果不对最后再提醒一个容易忽略的点不要以为不报错就等于万事大吉。ONNX 导出有成功但语义漂移的可能。index_put_这个场景尤其容易踩坑因为有些版本虽然导出成功但生成的 ONNX 图里计算逻辑已经和 PyTorch 不一致了。所以我每次改写代码后都会做一遍数值验证import numpy as np import onnxruntime as ort # PyTorch 推理结果 torch_output model(dummy_input).detach().numpy() # ONNX Runtime 推理结果 sess ort.InferenceSession(test.onnx) onnx_output sess.run(None, {x: dummy_input.numpy(), mask: dummy_mask.numpy()})[0] # 对比 print(最大绝对误差:, np.abs(torch_output - onnx_output).max()) print(是否一致:, np.allclose(torch_output, onnx_output, atol1e-6))如果你的操作涉及负数填充、极值填充atol可以放宽一些但绝对误差不超过1e-4。如果误差超了说明改写后的语义和原来有出入需要回头检查广播和 dtype 问题。我第一次改写注意力 mask 时就是因为mask的 shape 和logits不一致但又能隐式广播导致 ONNX 里的Where和 PyTorch 里的index_put_语义产生了微妙差异。那一次要不是跑了数值验证模型上线后推理结果可能悄悄漂移后果不堪设想。5.4 最后的小技巧不要遗忘版本记录排查这类问题过程中我越来越觉得记录环境版本 操作代码很重要。ONNX 导出的行为与 PyTorch 版本、opset 版本、甚至操作系统都强相关。我建议你在项目里维护一个小文件记录PyTorch 版本opset_version导出时用的torch.onnx.export参数报错信息原文最终采用的改写方案以后项目升级或换机器重新部署时这份记录能帮你省下从头排查的时间。这个习惯是我在连续踩了几个 ONNX 相关坑之后才养成的现在想想早该如此。回到开头那个报错。其实它并不难解决难的是搞清楚两件事底层的算子语义是什么以及导出器能接受的替代写法是什么。搞明白了这两点你会发现RuntimeError: Only consecutive 1-d tensor indices...只是一个善意的提醒——它告诉你该让模型代码更贴近 ONNX 的逻辑了。而torch.where这条改造路径到目前为止是我用过的最顺手、最稳妥的一招。