前阵子有个朋友问我a[[0,1],[1,2]]按他的理解应该返回四个元素结果 PyTorch 只给了两个。这个问题我太熟了当年我第一次写这类索引的时候盯着输出发了半天呆。PyTorch 的多轴索引multi-axis indexing看起来只是把 Python 里的a[i][j]换成a[i, j]但一旦引入整数数组、切片、布尔掩码的组合返回值的数量和形状遵循的是一套从 NumPy 继承下来的规则。规则没吃透索引报错就成家常便饭——超界、形状不匹配、布尔掩码对不上轴轮着来。这篇文章就专门聊多轴索引和它引发的报错既有最基础的轴概念也有整数数组索引的配对规则还有切片和高级索引混在一起时的维度顺序坑。适合刚接触张量操作的萌新也给写过不少数据预处理代码、却依然被输出形状搞晕的进阶选手当一面镜子。我会把每个知识点配一个小例子你照着在终端里敲一遍比背十遍文档都管用。1. 轴与基础索引先建立坐标系1.1 张量的三层嵌套盒子想理解多轴索引先得把 Tensor 看成嵌套盒子。拿形状(2, 3, 4)的张量举例它就像一栋两层楼每层有三排货架每排货架上放了四个格子。索引a[1, 2, 0]的意思很直白去第二层、第三排、第一个格子取东西。这里的 0、1、2 三个数字就是三个轴上的坐标每个数字都代表“沿着对应轴走几步”。很多人一开始会混淆“轴”和“维度”的称呼。在 PyTorch 里它们基本可以混用dim0是第一个轴dim1是第二个轴。关键是记住索引时你给的每一个位置参数从左到右分别对应第 0 轴、第 1 轴、第 2 轴……少给一个参数张量就会自动把后面的维度整体保留下来。例如形状(2, 3, 4)的张量写a[1]返回的是(3, 4)的子张量写a[1, 2]返回的是形状(4,)的一维张量。这个“参数越少返回维度越多”的习惯和多轴索引报错里的too many indices有直接关系。实际项目里最常见的场景是处理图像批次一个形状为(N, C, H, W)的 Batch Tensora[0]拿到第一张图a[:, 0]拿到所有图的第一个通道。单轴索引理解透了多轴索引只是把多个轴上的坐标一次性写出来但它背后有一套稍复杂的配对和广播逻辑。1.2 索引类型决定维度去留能把索引玩明白先要搞清楚 PyTorch 里有四类基本索引它们对输出形状的影响完全不一样整数索引a[1, 2]会“去掉”对应的维度。切片索引a[1:3]会“保留”该维度但宽度变成切片范围。整数数组索引a[[0, 2]]会“替换”该维度宽度变成数组长度。布尔索引a[a 0]会把对应维度“压扁”成一个不确定长度的一维结果。用生活经验来记整数像点名点一个人点到谁谁就出列切片像划一条范围整个队伍还站在原地整数数组像拿着几份名单挨个点名回到队伍里时按名单顺序重新排队布尔掩码像筛子留下的元素直接摊平放一排。我用一个小例子演示import torch a torch.arange(24).reshape(2, 3, 4) print(a[:, 1].shape) # 结果是 (2, 4)第1维的3被整数索引去掉 print(a[1].shape) # 结果是 (3, 4)第0维的2被整数索引去掉 print(a[1:2, 1].shape) # 结果是 (1, 4)切片保留了一个长度为1的维度如果你看到a[:, 1]的返回形状不是预期的(2, 4)而不是(2, 3, 4)说明你还没适应“整数去维”的规则。这个规则在多轴混合索引里会被反复用到。1.3 为什么要多轴一起索引有人会说我一个个轴慢慢取不行吗当然行但慢而且代码丑。更重要的是多个轴组合索引天然就是向量化操作不需要写 Python 循环。比如你有一个(100, 20, 15, 15)的批量数据想随机抽 32 个样本的第 3~8 个通道的某块区域一个表达式就能写完而且计算图能够把索引操作当成普通算子进行反向传播。数据加载和采样是这里的高频场景。很多自定义 Dataset 里都有类似batch[idx][:, start:end]的写法每次都要先生成idx整数数组再去取行切片和列切片。把多轴索引用熟之后这类代码可以一步到位也能避免不小心把数据复制了好几遍。2. 整数数组索引配对规则决定形状2.1 两个整数数组不是笛卡尔积回到开头那个问题。很多人第一次看到a[[0, 1, 2], [1, 2, 0]]会以为它会返回所有坐标组合一共 9 个元素。但实际上 PyTorch 的做法是把两个整数数组当成两份名单位置一一对应取(0, 1)、(1, 2)、(2, 0)这三个坐标。它更像是zip而不是两层 for 循环。a torch.arange(12).reshape(3, 4) print(a) # tensor([[ 0, 1, 2, 3], # [ 4, 5, 6, 7], # [ 8, 9, 10, 11]]) print(a[[0, 1, 2], [1, 2, 0]]) # tensor([1, 6, 8])验证一下a[0, 1]是 1a[1, 2]是 6a[2, 0]是 8。如果真想拿笛卡尔积那就得显式让它做广播后面会讲或者老老实实写循环。这个差异是索引报错里shape mismatch的高发源头——两份名单长度不一样时PyTorch 会直接跟你说索引张量不能广播到一起。2.2 整数数组与标量混合更常见的写法是一个轴用数组另一个轴用固定整数。比如选中第 0、1、2 行再固定要第 2 列print(a[[0, 1, 2], 2]) # tensor([2, 6, 10])这里整数标量2不会参与配对它只被当作“所有选中的行都取第 2 列”。输出形状完全由整数数组的长度决定。对应到图像处理里就是“选几张图但通道都选同一个”比如batch[[0, 5, 7], 0]只取这三张图的第一个通道。2.3 用广播制造组合效果如果你确实想要两个整数数组的笛卡尔积办法是用view或None把两个索引数组拆成不同形状让广播机制帮你扩展。rows torch.tensor([0, 2]).view(2, 1) cols torch.tensor([1, 3]).view(1, 2) result a[rows, cols] print(result.shape) # (2, 2) print(result) # tensor([[1, 3], # [9, 11]])rows形状(2, 1)cols形状(1, 2)广播后得到(2, 2)的坐标网格对应(0,1)、(0,3)、(2,1)、(2,3)。这里的原理和算术运算的广播一样形状在某个维度上是 1 就自动扩展。如果两个数组形状完全不同且不能广播就会触发IndexError: shape mismatch。2.4 三维及更多轴场景扩展到三维张量配对规则不变只是每个输出元素需要三个轴上的坐标一起来定。对b torch.arange(24).reshape(2, 3, 4)来说b torch.arange(24).reshape(2, 3, 4) idx0 torch.tensor([0, 1]) idx1 torch.tensor([1, 2]) idx2 torch.tensor([2, 3]) print(b[idx0, idx1, idx2]) # tensor([6, 23])这里(0,1,2)对应b[0, 1, 2] 6(1,2,3)对应b[1, 2, 3] 23。每个输出结果都是三个索引数组“逐元素对位”得到的。如果你的某个轴索引是标量那更简单这个坐标对每个输出元素都相同。实际跑这种代码时我习惯先把结果形状写在注释里再去预测值减少大脑的内存负担。3. 混合切片与高级索引三个维度顺序的坑3.1 切片保维度数组索引替换维度真实代码里切片和数组索引往往同时出现。比如a[:, [0, 2], :]表示“所有第一维第二维只取 0 和 2第三维全要”。切片冒号保留对应轴的大小整数数组把它替换成数组长度。b torch.arange(24).reshape(2, 3, 4) print(b[:, [0, 2], :].shape) # (2, 2, 4)只看结果很好理解第 0 维大小 2 保留第 1 维被替换成 2第 2 维大小 4 保留。新手最容易在这里翻车的是以为数组索引像切片一样“又宽了一点”但实际上它替换掉的是整个维度的取值方式返回的张量不是视图也不是副本的直觉能直接推断的。3.2 被切片隔开的高级索引会前移多轴索引最深的一个坑当两个整数数组索引之间夹着一个切片时整数数组索引生成的维度会被移到结果的最前面。这是 NumPy 兼容行为PyTorch 完整继承了下来但文档里讲得非常含蓄。b torch.arange(24).reshape(2, 3, 4) print(b[[0, 1], :, [1, 2]].shape) # 结果是 (2, 3)不是 (3, 2)直觉上第 0 维数组长度 2中间切片保留 3第 2 维数组长度 2看起来像是应该得到(2, 3, 2)或(3, 2)但真正的答案是两个高级索引被打散不连续它们产生的公共输出维度被搬到了最前面于是只剩(2, 3)。这里的2是数组索引长度3是中间切片保留下来的维度。我在项目里因为这个规则吃过亏当时从 NLP 模型里按句子索引和 token 索引抽 embedding写出来hidden[[batch_ids], :, [token_ids]]结果跑到后面对不上形状。后来打印标量验证才发现维度顺序变了。建议你在遇到多轴混合索引时先做个最小样例把形状试出来别死记。3.3 相邻与不相邻写法天差地别对比下面两组写法print(b[:, [0, 1], [2, 3]].shape) # (2, 2) print(b[[0, 1], :, [2, 3]].shape) # (2, 3)前者的两个整数数组索引紧挨着没有被切片隔开所以结果按原顺序放在后面前面第 0 维切片保留大小 2后面数组输出长度 2得到(2, 2)。后者的两个整数数组索引中间隔了切片数组输出维度被提前所以得到(2, 3)。不光形状变了值的排列方式也完全变了。前者的输出是b[:, 0, 2]和b[:, 1, 3]这两条向量的堆叠后者的输出是b[0, :, 2]和b[1, :, 3]这两条向量的堆叠。想靠大喊“它应该差不多”是过不去的必须按规则一步步拆。上面的小例子建议你亲手跑一遍把中间结果打出来看。4. 布尔索引与多轴组合筛选和定位的正确姿势4.1 布尔掩码先压扁布尔索引在数据筛选里几乎天天用。a[a 3]会把满足条件的元素全部挑出来返回一维张量顺序是按行优先排的。a torch.tensor([[1, 2, 3], [4, 5, 6]]) print(a[a 3]) # tensor([4, 5, 6])有人写a[a 3]之后期待得到一个二维结果结果发现变成一维就开始慌了。其实这是特性不是 bug布尔掩码不管原张量是几维筛选后都先压成一维除非你用torch.nonzero保留坐标信息。4.2 掩码只作用在部分维度很多时候你只想沿部分维度做筛选。比如形状(2, 3, 4)的张量想根据前两个维度上的条件筛选同时保留最后一个维度全部内容就可以让掩码形状与张量的前两个维度一致b torch.arange(24).reshape(2, 3, 4) mask b[:, :, 0] 4 print(mask.shape) # (2, 3) print(b[mask].shape) # (5, 4)这里的mask是(2, 3)布尔张量它只作用于第 0 和第 1 维筛选出 5 个位置所以结果第一维变成 5第三维完整保留 4。这个模式在做“对整行或整块数据做局部筛选”时非常有用比先把所有坐标nonzero出来再挨个取要友好得多。4.3 布尔索引和维度没对齐的报错布尔索引最常见的报错是维度和轴大小对不上。比如a torch.arange(12).reshape(3, 4) mask torch.tensor([True, False, True, True]) # 长度是4 try: a[mask] except IndexError as e: print(e)报错信息会明确告诉你boolean index did not match indexed array along dimension 0; dimension is 3 but corresponding boolean dimension is 4。看到这个信息时第一反应不是改 mask 的数据分布而是检查它到底对应哪个轴。很多时候你以为 mask 在筛“列”但实际上你把它写在了行轴的位置上。5. 多轴索引报错速查表与真实排查实录5.1 报错速查表下面是我在实际调试中最常遇到的几类报错整理成速查表可以直接对着看。报错特征文本触发原因修复方向index X is out of bounds for dimension Y with size Z索引值超过该轴大小检查索引数组是否越界必要时做clamp或先打印shapeshape mismatch: indexing tensors could not be broadcast together多个整数数组索引形状无法广播统一索引数组长度或使用view调整形状boolean index did not match indexed array along dimension X布尔掩码长度与该轴大小不一致确认掩码对应哪个轴改成对应轴的真实长度too many indices for tensor of dimension N给的索引坐标数量超过张量的轴数去掉多余坐标或先检查张量是不是比你以为的少一维only integer tensors of a single element can be converted to an index用了包含多个元素的 Tensor 当标量位置改成int(idx)或改用idx.item()invalid index of a 0-dim tensor给 0 维标量张量做索引标量张量不要用[0]直接用ten.item()或float(ten)其中too many indices最阴险经常出现在你用之前某个预处理函数的输出张量去索引时函数返回的维度比预期少了一个。遇到这种报错我会直接绕开原表达式把张量拆成最小样例再测。5.2 最小样例复现法排查索引报错不能靠在大工程里瞎试。我的标准流程是把报错涉及的张量和索引原样抄出来用一个尽量小的 Demo 复现。比如你在(64, 128, 256)上出错了那就先把它换成(2, 3, 4)形状结构保持一致但每一步都能肉眼算清楚。b torch.arange(24).reshape(2, 3, 4) i0 torch.tensor([0, 1]) i2 torch.tensor([2, 3]) # 先打印中间结果的形状 print(b.shape, i0.shape, i2.shape) print(b[i0, :, i2].shape)这一步能筛掉至少一半的问题因为报错往往不是“值算错”而是“形状推导错了”。一旦形状对了再回来用真实数据跑心里就踏实很多。5.3 拆成多步索引当你面对复杂的混合索引心里没底时拆开写是最稳妥的办法。比如想把每个满足条件的样本的第 0、2 个通道拿出来你完全可以先写selected b[mask]得到一个中间张量再去selected[:, [0, 2]]。这样做虽然多了两步中间变量但每一步都清晰也不容易触发高级索引被切片隔开的维度前移问题。性能上绝大多数场景这点拷贝开销可以忽略。千万不要在同一个表达式里又用布尔掩码、又用整数数组、又穿插切片去挑战极限。能跑出来的人很多跑不出来的人也很多你不需要在这上面证明自己。6. 实战项目中的避坑指南6.1 需要坐标网格时优先用 meshgrid想要多个整数数组的笛卡尔积除了上一章提到的view广播更推荐直接用torch.meshgrid。这是语义最清晰的做法。rows torch.tensor([0, 2]) cols torch.tensor([1, 3]) grid_r, grid_c torch.meshgrid(rows, cols, indexingij) a torch.arange(12).reshape(3, 4) print(a[grid_r, grid_c]) # tensor([[1, 3], # [9, 11]])meshgrid返回两个和输出同形状的坐标张量索引时完全对位。它比view广播的做法直观也能避免你记不住广播规则。6.2 越界前先检查而不是直接 clamp有人喜欢用torch.clamp把所有索引值限制在合法范围内但这会静默改变你的采样语义。比如你想取第 4 行结果张量只有 3 行clamp之后悄悄变成第 2 行模型可能还在用错误数据训练。更安全的做法是先做一次断言或打印idx torch.tensor([0, 5, 2]) n a.shape[0] assert idx.max() n, f索引越界: max{idx.max()}, dim_size{n}如果只是偶尔越界我更推荐用torch.where把越界索引替换成 0并记录哪些位置发生了替换再根据场景决定要不要丢弃这批数据。6.3 想取稀疏采样位置用 nonzero布尔索引返回的一维结果会丢失“原坐标在哪”的信息。如果你既想筛选又想保留每个坐标对应的行列位置用torch.nonzerob torch.arange(24).reshape(2, 3, 4) mask b 10 coords torch.nonzero(mask) # shape (num_true, 3) print(coords.shape)coords的每一行都是一个原始坐标。你可以用coords[:, 0]、coords[:, 1]分别当作用在各个轴上的整数数组索引也可以直接逐行取数。当年我在做点云数据处理时就靠这个方式代替多层循环把几千个候选点一次性筛出来。6.4 把“预期形状”写下来再写代码有段时间我反复在索引 bug 上栽跟头后来想了个笨办法写代码前先在注释里写出最终输出形状和每个轴是怎么来的。比如# 输入: x (64, 128, 5) # 索引: x[rows, :, feats] # rows: (32,), feats: (32,) # 预期: 高级索引被切片隔开输出形状 (32, 128)这样写完之后哪怕运行报错我只要拿报错信息和注释一比立刻知道是维度顺序理解错了还是索引长度对不上了。这个方法看起来傻但很有效特别适合刚上手 PyTorch 的同学。6.5 索引代码保持可读性在深度学习项目里索引经常出现在Dataset或者数据增强逻辑里。代码一旦进入迭代期可读性比“酷”重要得多。一个建议是如果你的索引表达式超过三层混合试着把复杂的布尔筛选和数组索引拆成两个函数分别起个好名字比如filter_by_label和sample_by_index。这样别人读你的代码时不用先心理默算一遍维度项目交接和排障都轻松不少。我自己平时写多轴索引最终习惯是“先画形状再写代码最后补注释”。每次遇到奇怪输出先把中间结果打出来再回头推规则。靠这个流程我至少减少了一半以上的索引报错。索引是张量操作的地基多花点时间把它弄扎实后面的模型调试才会更顺。