张量索引这件事说简单也简单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, ::2]视图按切片规则收缩连续区间、步长取样整型数组索引t[[0,2], [1,3]]副本由索引形状决定任意位置散点取值布尔掩码索引t[mask]副本拉平为一维条件筛选专用函数gather/index_select副本按 dim 保持维度按索引张量批量取这张表建议先记住“视图还是副本”这一列因为它直接决定了你后续能不能原地改、改了会不会影响原张量。基础切片返回视图意味着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用来插入新维度等价于unsqueezet[:, 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_select3.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]]dim1 时。注意 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 scattergather 的逆操作写入而非读取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形状由 index 决定、GPU 高效按坐标逐元素写scatter/scatter_add写入语义明确条件筛选布尔掩码写法最自然连续区间基础切片零拷贝、最快选型时先问“我是按维度选还是按坐标选”按维度就用index_select或切片按坐标就用gather/scatter。这个二分法能覆盖九成场景。4. 实操过程与关键环节实现4.1 一个完整的按长度取有效位置的例子序列任务里最常见的索引需求是一个 batch 的序列做了 padding要按每个样本的真实长度取出有效 token 做后续计算。这个需求用布尔掩码最自然但要注意保留维度。import torch # 假设 batch2最大长度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, dim1) 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 索引相关的高频报错速查报错信息常见原因解决index out of range索引值 维度大小检查 index 最大值必要时 clampshape mismatchgather/scatter 的 index 形状不对用 unsqueeze/expand 对齐形状expected LongTensor索引用 float.long()转换a leaf Variable that requires grad is being used in an in-place operation对需要梯度的张量原地索引赋值改用非原地写法或 detach结果形状和预期不符布尔掩码拉平 / 广播配对打印 shape 逐步排查这张表是我自己 debug 时最常翻的尤其是前两条占了索引报错的大半。5.2 索引赋值与梯度为什么有时不能原地改索引不仅能读还能写t[idx] value是合法的但如果t是需要梯度的叶子张量原地赋值会破坏计算图PyTorch 会直接报错。正确做法是用torch.index_put的非原地版本或者先 clone 再改。t torch.randn(5, requires_gradTrue) # 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([], dtypetorch.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、大特征图时对内存很友好。这块我还在持续摸索等有更系统的实测数据再单独整理。