PyTorch张量索引全解析：从基础切片到gather/scatter实战

张量索引这件事，说简单也简单，t[0]、t[:, 1]谁都会写；说复杂也复杂，真到调试一个维度对不上的 bug、或者想从一个大 batch 里精准抠出若干行若干列的时候，很多人就开始靠猜和试。我自己在带新人和 review 代码的过程中发现，索引操作是 PyTorch 里最容易“看起来会了、实际一写就错”的部分之一，因为它牵扯到基础切片、整型数组索引、布尔掩码、gather/scatter、index_select这一整套东西，而且每种方式返回的是视图还是副本、形状怎么变，都有讲究。这篇内容就把张量索引这一块从头到尾捋一遍，既讲清楚每种写法的语义和底层逻辑，也给出可以直接抄的代码和踩坑记录，适合刚上手 PyTorch 的读者打基础，也适合已经能写模型但索引写得不够利索的读者查漏补缺。

1. 先把索引这件事的整体设计想明白

1.1 为什么张量索引值得单独拿出来讲

很多人学 PyTorch 是从“照着教程搭一个网络”开始的，索引只是顺带用一下，比如取logits[:, 0]做二分类、取batch[0]看一条样本。这种用法确实不需要理解太多。但只要你的任务稍微复杂一点，比如做序列标注要按长度取有效位置、做检测要按 mask 筛选正样本、做推荐要按用户 id 取 embedding 行，索引就会变成核心逻辑的一部分，写错了不会报错，只会悄悄给你一个形状不对或者数值不对的结果，然后你在 loss 曲线里找半天原因。

我见过最典型的一个例子：某同学做多标签分类，想按标签为 1 的位置取预测值算 loss，写成了preds[labels == 1]，结果preds是二维的[N, C]，labels == 1也是二维的，这一下把整个张量拉平成一维了，形状全乱，但代码不报错，训练照样跑，只是效果差。这类问题不是“不会 API”，而是没搞清楚索引的语义边界。所以把索引单独拎出来系统讲一遍，是有实际价值的。

从设计角度看，PyTorch 的索引体系其实是在对齐 NumPy 的基础上做了扩展，核心目标有三个：一是让“取子集”这件事写得自然，接近数学记号；二是尽量让常见操作零拷贝，靠视图（view）省内存；三是给自动求导留出正确的梯度路径。理解了这三个目标，你就能明白为什么有些索引用完是视图、有些是副本，为什么有些写法在反向传播时会报“原地操作”的错。

1.2 索引方式的分类与选型逻辑

把 PyTorch 里所有索引写法归个类，其实就四大类，选型的时候按“我要什么、我能不能接受拷贝、我要不要梯度”三个问题走一遍就行。

这张表建议先记住“视图还是副本”这一列，因为它直接决定了你后续能不能原地改、改了会不会影响原张量。基础切片返回视图，意味着t[1:3]和t共享内存，你改前者后者也变；而整型数组和布尔掩码返回的是新张量，改它不影响原数据。这个差异在写数据增强、写 in-place 归一化的时候特别关键。

选型上我的习惯是：能用基础切片解决的绝不用花哨写法，因为切片最直观、性能最好、梯度路径最清晰；需要按任意下标取的时候优先考虑index_select而不是 Python 列表索引，因为前者在 GPU 上更高效且语义明确；需要按条件筛选时用布尔掩码，但一定要意识到它会拉平维度，如果不想拉平就得用masked_select之外的手段或者配合where。

1.3 视图与副本：索引里最容易被忽略的分水岭

这一小节单独展开，因为它是后面所有坑的根源。PyTorch 里一个张量由两部分组成：一块存储（storage）和一份元信息（size、stride、offset）。基础切片之所以能零拷贝，是因为它只改了元信息里的 offset 和 stride，存储还是那块存储。你可以用.storage().data_ptr()验证两个张量是不是共享底层内存。

import torch t = torch.arange(12).reshape(3, 4) v = t[1:3] print(t.data_ptr() == v.data_ptr()) # True，共享存储 print(v.storage_offset()) # 4，说明从第 4 个元素开始 print(v.stride()) # (4, 1)，和原张量一致

而整型数组索引t[[0, 2]]走的是另一条路，它必须新建一块存储把选中的行拷过去，因为选中的行在原存储里不连续，没法用 stride 表达。所以它返回副本，data_ptr不同。

注意：判断视图还是副本不要靠记忆，直接看data_ptr或者用v._base判断，_base不为 None 说明它是某个张量的视图。写库函数或者写需要原地修改的代码前，先确认这一点能省掉大量诡异 bug。

理解了这个分水岭，后面讲gather、讲布尔索引、讲梯度回传时，很多“为什么”就自然通了。

2. 基础切片与高级索引的细节拆解

2.1 基础切片：冒号、步长与省略号

基础切片就是 Python 列表切片那套语法搬到多维上，每一维用start:stop:step描述，维与维之间用逗号分隔。规则很朴素：左闭右开、步长可负、省略 start 默认 0、省略 stop 默认到维度末尾。真正需要留心的是“省略号”和“None”这两个特殊符号。

...（省略号）表示“剩下的所有维度”，在维度多的时候特别好用。比如你有一个[B, C, H, W]的特征图，想取所有 batch、所有通道的第 0 行，可以写feat[..., 0, :]，不用数到底几个冒号。None（或者torch.newaxis）用来插入新维度，等价于unsqueeze，t[:, None]会把[N]变成[N, 1]。这两个符号组合起来能写出很紧凑的代码，但也容易写错维度顺序，我的建议是复杂场景下宁可多写两个unsqueeze，可读性优先。

t = torch.arange(24).reshape(2, 3, 4) print(t[..., 0].shape) # [2, 3]，取最后一维第 0 个 print(t[:, None].shape) # [2, 1, 3, 4]，插入新维 print(t[0, :, ::2].shape) # [3, 2]，第 0 个 batch、所有行、步长 2

步长为负的时候要小心边界，t[::-1]是反转，但如果 start/stop 写得不合适会得到空张量而不报错。我一般反转就用torch.flip，语义更清楚，也不容易踩空切片的坑。

2.2 整型数组索引：散点取值的形状规则

整型数组索引（也叫 advanced indexing）是用一个整型张量或列表作为下标，比如t[[0, 2, 1]]。它和切片的根本区别在于：切片选的是连续区间，整型数组选的是任意位置集合，所以结果一定是副本。

形状规则是这块最容易错的地方，我把它总结成一句话：索引张量的形状决定输出形状，被索引的维度被“替换”掉。当你在多个维度上同时给整型数组时，这些数组会按广播规则对齐，然后逐元素配对取值。

t = torch.arange(12).reshape(3, 4) # 单维索引：选第 0、2 行，输出 [2, 4] print(t[[0, 2]].shape) # 多维配对索引：取 (0,1)、(2,3) 两个点，输出 [2] rows = torch.tensor([0, 2]) cols = torch.tensor([1, 3]) print(t[rows, cols]) # tensor([1, 11]) # 广播配对：rows 形状 [2,1]，cols 形状 [1,3]，输出 [2,3] print(t[rows[:, None], cols[None, :]].shape)

最后那个广播例子是重点，它等价于取一个 2x3 的笛卡尔积子矩阵，做批量采样、做注意力里的位置对齐时经常用到。很多人第一次见t[rows[:, None], cols]会懵，其实就是把两个一维索引张量先广播成同形状再配对。

实操心得：多维整型索引时，先把每个索引张量的形状在纸上写出来，按广播规则对齐，输出形状就是广播后的形状。这一步花十秒，能省掉半小时的 debug。

2.3 布尔掩码索引：条件筛选与维度拉平的陷阱

布尔掩码索引用一个 bool 张量做下标，选出为 True 的位置。它最实用的场景是“按条件筛样本”，比如features[labels == 1]取出所有正类样本。语义上它等价于先把掩码nonzero拿到坐标，再做整型数组索引，所以结果同样是副本。

关键陷阱在于：布尔掩码索引会把结果拉平成一维（当掩码和原张量同形状时）。t[mask]里mask和t形状一致，输出就是所有 True 位置的值组成的一维张量，原来的维度结构全丢了。如果你想要“按行筛选、保留列结构”，掩码必须是“按维度的”，比如t[mask_rows]里mask_rows是一维的、长度等于行数，这样输出才是[选中行数, 列数]。

t = torch.arange(12).reshape(3, 4) mask_full = t > 5 print(t[mask_full].shape) # [6]，拉平了 mask_rows = torch.tensor([True, False, True]) print(t[mask_rows].shape) # [2, 4]，保留列结构

这个区别在做数据筛选时极其重要。我踩过的坑是：想按样本筛选，结果掩码写成了和特征同形状的二维，输出被拉平，后面接全连接层直接维度报错。正确做法是先对标签或某个标量做条件得到一维掩码，再用它索引特征。

另外torch.masked_select和布尔索引效果类似，也是拉平返回一维，区别是它支持广播。需要保留维度的时候，用torch.where(cond, t, other)更合适，它按元素选择、形状不变。

2.4 组合索引：切片与数组混用的优先级

实际代码里经常是切片和整型数组混着用，比如t[1:3, [0, 2]]。这种混合索引的规则是：先按位置把切片和数组分开，切片作用在对应维度上收缩，数组维度按广播配对，最后把结果拼起来。PyTorch 的处理顺序和 NumPy 一致，但有个细节要注意——当整型数组出现在切片前面时，维度顺序可能和你预期不同。

t = torch.arange(24).reshape(2, 3, 4) # 切片在前、数组在后：输出 [2, 2, 4] print(t[:, [0, 2]].shape) # 数组在前、切片在后：输出 [2, 2, 4]，但语义是选第 0、1 个 batch 的第 2 列 print(t[[0, 1], :, 2].shape)

我的经验是：混合索引超过两个维度就别硬写了，拆成两步、中间变量命名清楚，可读性和正确率都高得多。性能上多一次索引的拷贝开销，在绝大多数场景下可以忽略，不值得为省这一次拷贝去写一行没人看得懂的索引。

3. 专用索引函数：gather、scatter 与 index_select

3.1 index_select：按维度批量取，语义最清晰

torch.index_select(input, dim, index)干的事很明确：在指定维度上，按index里的下标取子集，其他维度保持不变。它比 Python 列表索引好在两点：一是语义显式（明确说了在哪个 dim 上选），二是 GPU 上实现更高效，三是index必须是 LongTensor，类型检查能帮你早发现错误。

t = torch.arange(12).reshape(3, 4) idx = torch.tensor([0, 2]) print(torch.index_select(t, 0, idx).shape) # [2, 4]，选第 0、2 行 print(torch.index_select(t, 1, idx).shape) # [3, 2]，选第 0、2 列

它和t[idx]在单维场景下结果一样，但index_select不会触发广播配对那套复杂规则，写起来更“安全”。做 embedding 查表、做 batch 内按索引取样本时，我基本都用它。

3.2 gather：按坐标逐元素取，形状由索引决定

gather是索引里最“反直觉”的一个，但也是做序列任务绕不开的。它的语义是：对输出里的每个位置，去 input 的对应位置按 index 取值。公式化地说，out[i][j] = input[i][index[i][j]]（dim=1 时）。注意 index 的形状必须和输出形状一致，且 index 的每个值不能超过 input 对应维度的大小。

input = torch.tensor([[10, 20, 30], [40, 50, 60]]) index = torch.tensor([[2, 0], [1, 2]]) out = torch.gather(input, 1, index) # out = [[30, 10], [50, 60]]

这个例子要这么读：第 0 行，index 是[2, 0]，就去 input 第 0 行取第 2 个和第 0 个，得[30, 10]；第 1 行同理。gather的典型用途是“按预测的类别索引取对应 logit”，比如做分类时logits.gather(1, labels.unsqueeze(1))取出每个样本真实类别的分数，比 one-hot 相乘省内存。

注意：gather的 index 形状必须和输出一致，且 dtype 必须是 int64。很多人第一次用会传一个一维 index 想取一列，结果报形状不匹配，正确做法是把 index 扩成二维。

3.3 scatter：gather 的逆操作，写入而非读取

scatter是把 src 里的值按 index 写到 output 的对应位置，out[i][index[i][j]] = src[i][j]。它在做 one-hot 编码、做标签平滑、做“把预测结果写回原位置”时很有用。scatter_add则是累加版本，同一个位置被多次写入时求和而不是覆盖，做计数、做直方图时用它。

out = torch.zeros(2, 3) index = torch.tensor([[0, 2], [1, 1]]) src = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) out.scatter_(1, index, src) # out = [[1, 0, 2], [0, 7, 0]]，第 1 行两个 3 和 4 都写到位置 1，累加得 7

这里用的是scatter_原地版本，注意它要求 out 已经存在且形状正确。scatter的坑在于“同一位置多次写入”的语义：普通scatter是覆盖，结果取决于写入顺序，不确定；要累加必须用scatter_add。我见过有人用scatter做投票统计，结果因为覆盖丢票，换成scatter_add才对。

3.4 三个函数的选型对照

选型时先问“我是按维度选还是按坐标选”，按维度就用index_select或切片，按坐标就用gather/scatter。这个二分法能覆盖九成场景。

4. 实操过程与关键环节实现

4.1 一个完整的按长度取有效位置的例子

序列任务里最常见的索引需求是：一个 batch 的序列做了 padding，要按每个样本的真实长度取出有效 token 做后续计算。这个需求用布尔掩码最自然，但要注意保留维度。

import torch # 假设 batch=2，最大长度=4，真实长度分别是 3 和 2 seq = torch.tensor([[1, 2, 3, 0], [4, 5, 0, 0]]) lengths = torch.tensor([3, 2]) # 构造位置掩码：位置 < 长度 为 True pos = torch.arange(seq.size(1)).unsqueeze(0) # [1, 4] mask = pos < lengths.unsqueeze(1) # [2, 4] valid = seq[mask] # [5]，拉平成一维 print(valid) # tensor([1, 2, 3, 4, 5])

这里seq[mask]得到的是拉平后的有效 token，正好适合送进一个对所有 token 统一处理的层（比如 token 级分类）。如果后续需要按样本聚合，就得配合scatter_add或者用index_select逐样本处理。这个例子说明：布尔掩码拉平不一定是坏事，关键看下游要什么形状。

4.2 按索引张量批量取 embedding 行

推荐或 NLP 里经常要按一批 id 取 embedding。假设 embedding 表是[V, D]，id 是[B]，要得到[B, D]。

V, D, B = 1000, 16, 8 emb = torch.randn(V, D) ids = torch.randint(0, V, (B,)) # 方式一：index_select out1 = torch.index_select(emb, 0, ids) # [8, 16] # 方式二：直接索引 out2 = emb[ids] # [8, 16] print(torch.equal(out1, out2)) # True

两种写法结果一致，emb[ids]更简洁，index_select更显式。性能上在大表、大 batch 时index_select略稳，因为它不涉及广播规则判断。我一般写模型用emb[ids]，写工具函数用index_select，兼顾可读和明确。

4.3 用 gather 取真实类别分数并算交叉熵

分类任务里想手动算“真实类别的预测分数”，gather是标准做法。

logits = torch.randn(4, 5) # 4 个样本、5 类 labels = torch.tensor([1, 3, 0, 4]) # 取每个样本真实类别的分数 true_scores = logits.gather(1, labels.unsqueeze(1)).squeeze(1) # [4] print(true_scores.shape) # 手动算交叉熵（仅演示，实际用 F.cross_entropy） log_probs = torch.log_softmax(logits, dim=1) nll = -log_probs.gather(1, labels.unsqueeze(1)).squeeze(1).mean() print(nll)

这里labels.unsqueeze(1)把[4]变成[4, 1]，满足 gather 对 index 形状的要求。.squeeze(1)再把结果还原成[4]。这一串操作是模板，记住就行。踩过的坑是忘了unsqueeze，直接gather(1, labels)报形状不匹配，因为 gather 要求 index 和输出同形状，而输出是[4, 1]。

4.4 用 scatter_add 做按索引聚合

假设有一批特征和对应的组 id，要把同组的特征求和。这是scatter_add的经典用法。

feat = torch.randn(6, 3) # 6 条特征 group = torch.tensor([0, 1, 0, 2, 1, 0]) # 每条属于哪个组 num_groups = 3 agg = torch.zeros(num_groups, 3) agg.scatter_add_(0, group.unsqueeze(1).expand(-1, 3), feat) print(agg.shape) # [3, 3]，每组特征求和

关键在group.unsqueeze(1).expand(-1, 3)：scatter_add 要求 index 和 src 同形状，所以要把组 id 沿特征维扩展。这个写法在 pooling、图神经网络的消息聚合里非常常见，值得背下来。

5. 常见问题与排查技巧实录

5.1 索引相关的高频报错速查

这张表是我自己 debug 时最常翻的，尤其是前两条，占了索引报错的大半。

5.2 索引赋值与梯度：为什么有时不能原地改

索引不仅能读还能写，t[idx] = value是合法的，但如果t是需要梯度的叶子张量，原地赋值会破坏计算图，PyTorch 会直接报错。正确做法是用torch.index_put的非原地版本，或者先 clone 再改。

t = torch.randn(5, requires_grad=True) # t[[0, 2]] = 0.0 # 报错：leaf Variable 不能原地改 t2 = t.clone() t2[[0, 2]] = 0.0 # 可以，t2 不是叶子

这个限制在做 mask 填充、做 label 修正时经常遇到。我的习惯是：凡是涉及“按索引改值”且张量参与梯度计算，一律先 clone，宁可多一份内存也不去和自动求导机制较劲。

5.3 性能上的几个实测经验

索引操作的性能差异在 CPU 上不明显，在 GPU 上放大。我实测下来几条经验：基础切片几乎零成本，因为它只改元信息；index_select和gather会触发拷贝，但都是高度优化的 kernel，比 Python 循环快几个数量级；布尔掩码索引因为要先算nonzero再拷贝，开销比index_select略大，能用一维掩码就别用二维全掩码。

还有一个容易忽略的点：频繁的小规模索引会打断 GPU 的 kernel 合并，导致大量小 kernel 启动开销。如果在一个循环里反复做索引，考虑把索引合并成一次批量操作，或者把数据先搬到 CPU 处理完再搬回去（数据量小的时候反而更快）。

实操心得：怀疑索引是性能瓶颈时，用torch.profiler看 kernel 时间，别靠猜。我见过有人把index_select换成切片以为能提速，结果因为切片后接的操作不连续反而更慢。

5.4 几个容易写错的边界情况

第一，空索引。t[torch.tensor([], dtype=torch.long)]返回形状[0, ...]的空张量，不报错，但下游如果没处理空 batch 会崩。做数据筛选时一定要考虑“筛完为空”的情况。

第二，负索引。PyTorch 支持负索引，t[-1]是最后一行，但整型数组里的负值也会被解释为从末尾数，t[[-1, -2]]是合法的。这个特性有时很方便，有时很坑，因为如果你本意是“无效标记 -1”，它会被当成最后一行而不是报错。我的建议是索引张量里永远不要出现负数，无效值先过滤掉。

第三，布尔掩码的 dtype。掩码必须是torch.bool，用 0/1 的整型张量当掩码不会报错但语义不对（会被当成整型数组索引），这是个静默错误，特别隐蔽。

6. 把索引写对写快的几条个人习惯

写索引代码时，我养成了几个固定动作，分享出来供参考。第一个动作是“先打印形状”：任何超过两维的索引，写之前先把涉及的张量 shape 打出来，索引写完再打一次结果 shape，和预期对不上立刻停。第二个动作是“复杂索引拆两步”：混合索引、广播配对这类写法，能拆就拆，中间变量起个有意义的名字，可读性换来的正确率提升远超那点性能损失。第三个动作是“掩码只做一维”：需要按样本筛选就用一维掩码，需要按元素筛选才用全形状掩码，并且明确知道结果会被拉平。

还有一个习惯是给索引操作写小测试。索引逻辑一旦写进模型，出错很难定位，不如在写的时候就拿一个小张量手动算一遍期望值，用torch.equal断言。我现在的习惯是每个涉及 gather/scatter 的函数旁边都留一个if __name__ == "__main__"的小用例，改代码时跑一下，比事后 debug 省心得多。

最后说一个后续可以扩展的方向：如果你把索引和torch.Tensor.expand、torch.broadcast_to结合起来，能写出很多“零拷贝的批量视图”，这在处理大 batch、大特征图时对内存很友好。这块我还在持续摸索，等有更系统的实测数据再单独整理。