张量进到网络里之后，我们不可能每次都对整块数据做运算。训练时要按条件挑出某一批样本，注意力里要取某个位置的向量，调试时要看某一行某一列的值对不对。这些动作都靠索引完成。索引写得不对，轻则取错数，重则形状悄悄变了，后面接的层直接报错。

这一节我们把最常用的五类索引一次讲透，分别是基本行列索引、列表索引、范围切片、布尔索引、多维索引。每一类都拿一张固定的 4 行 5 列矩阵走一遍输入、调用和输出，讲清楚每种写法到底挑出了哪些位置。读之前，我们需要认识二维张量，知道第一维是行、第二维是列，以及冒号在切片里的意思。

为了让例子能复现，我们先固定一张矩阵。它按行从 1 填到 20，后面所有操作都在这张表上做：

import torch

data = torch.arange(1, 21).reshape(4, 5)
print(data)
# tensor([[ 1,  2,  3,  4,  5],
#         [ 6,  7,  8,  9, 10],
#         [11, 12, 13, 14, 15],
#         [16, 17, 18, 19, 20]])

1. 基本索引

最简单的索引，是用一个整数定位。data[0] 取第 0 行，data[:, 0] 取第 0 列。这里的冒号表示这一维全要，行维写冒号就是所有行，列维给 0 就是只要第 0 列。逗号把两个维度的选择隔开，这是后面所有写法的基础。

print(data[0])        # 第 0 行
# tensor([1, 2, 3, 4, 5])

print(data[:, 0])     # 所有行的第 0 列
# tensor([ 1,  6, 11, 16])

我们注意一下，取一行得到的是一维向量，取一列也是一维向量，原来的 4 乘 5 在取单点之后，降成了 5 个元素的一条线。这就是索引和切片的关键区别：整数索引会把对应那维直接消掉，冒号切片则会把那维保留下来。

2. 列表索引

基本索引一次只能选一个位置，想同时挑好几个行或列，就要用列表。这一部分最容易搞混，因为两个维度上的列表配对方式不同，挑出来的结果差别很大。我们先看成对写法，data[[0, 1], [1, 2]] 表示行列表和列列表一一配对，只取 (0,1) 和 (1,2) 两个点。

# 行列表 [0,1] 和列列表 [1,2] 配对，取 (0,1) 和 (1,2) 两个点
print(data[[0, 1], [1, 2]])
# tensor([2, 8])

# 行列表写成列向量形状 [[0],[1]]，列列表是 [1,2]，二者广播成一个 2 乘 2 的子矩阵
print(data[[[0], [1]], [1, 2]])
# tensor([[2, 3],
#         [7, 8]])

第一种写法，两个列表同形，按位置两两配对，结果是一维的两个数，2 来自第 0 行第 1 列，8 来自第 1 行第 2 列。第二种写法，把行列表竖起来写成 2 行 1 列，列列表还是一行两列，PyTorch 自动广播成 2 行 2 列的网格，把第 0、1 行和第 1、2 列围成的子矩阵整块取出来。差别就在于，行列表是横着写还是竖着写。

3. 范围切片

要挑连续的一段，我们用 start:end 这种切片写法。规则和 Python 列表一致，包含 start，不包含 end，end 留空就是到末尾。这种写法用冒号占位，被切的那一维会保留下来，不会像整数索引那样消失。

# 前 3 行的前 2 列
print(data[:3, :2])
# tensor([[ 1,  2],
#         [ 6,  7],
#         [11, 12]])

# 第 2 行到最后的前 2 列
print(data[2:, :2])
# tensor([[11, 12],
#         [16, 17]])

:3 表示 0、1、2 三行，不含第 3 行。2: 表示从第 2 行一直取到最后，也就是第 2、3 两行。切片常用于裁掉 padding、取图像的中心区域这类操作。因为保留了维度，切出来还是二维矩阵，可以直接接着做运算。

4. 布尔索引

前面几种写法都是按位置挑，实际工程里更多是按条件挑。布尔索引先拿张量和一个比较符算出一张真假表，再拿这张表去筛数据。比如筛出第三列大于 5 的所有行，或者挑出第二行大于 5 的所有列。

# 第三列大于 5 的行。第三列是 [3, 8, 13, 18]，大于 5 的是第 1、2、3 行
print(data[data[:, 2] > 5])
# tensor([[ 6,  7,  8,  9, 10],
#         [11, 12, 13, 14, 15],
#         [16, 17, 18, 19, 20]])

# 第二行大于 5 的列。第二行是 [6, 7, 8, 9, 10]，五个数都大于 5
print(data[:, data[1] > 5])
# tensor([[ 1,  2,  3,  4,  5],
#         [ 6,  7,  8,  9, 10],
#         [11, 12, 13, 14, 15],
#         [16, 17, 18, 19, 20]])

第一步 data[:, 2] > 5，先算出一个长度为 4 的布尔向量 [False, True, True, True]，拿它当行索引，PyTorch 就把对应为 True 的行挑出来。列方向同理，data[1] > 5 算出五个全是 True 的布尔向量，于是五列全保留。这种写法在数据清洗和按样本筛选时用得极多，比手动写行号灵活得多。

5. 多维索引

二维的规则搬到三维张量上完全通用，只是多了一层维度。我们造一个 (3, 2, 5) 的三维张量，把它想成 3 张 2 行 5 列的小矩阵摞在一起，再分别用 0 号切片切三个方向。

data3 = torch.arange(30).reshape(3, 2, 5)
print(data3.shape)          # torch.Size([3, 2, 5])

# 取第 0 张矩阵，剩下两维
print(data3[0])
# tensor([[0, 1, 2, 3, 4],
#         [5, 6, 7, 8, 9]])

# 每张矩阵的第 0 行
print(data3[:, 0])
# tensor([[ 0,  1,  2,  3,  4],
#         [10, 11, 12, 13, 14],
#         [20, 21, 22, 23, 24]])

# 每张矩阵的第 0 列
print(data3[:, :, 0])
# tensor([[ 0,  5],
#         [10, 15],
#         [20, 25]])

三个切片分别咬住不同的维度。data3[0] 在第 0 维取单点，把三张矩阵摞中的第一张拿出来，剩 (2, 5)。data3[:, 0, :] 等价于 data3[:, 0]，在中间维取第 0 行，把每张矩阵的第一行抽出来摞成 (3, 5)。data3[:, :, 0] 在最后一维取第 0 列，得到 (3, 2)。规律和二维一样，整数索引消掉那一维，冒号切片保留那一维。

6. 常见误区

第一个坑，是把成对列表索引当成取子矩阵。data[[0, 1], [1, 2]] 只挑两个散点，不是第 0、1 行第 1、2 列的整块。想要子矩阵，得把行列表竖起来写成 [[0], [1]]，或者干脆分开两步写 data[[0, 1]][:, [1, 2]]。

第二个坑，是切片的 end 边界。Python 切片左闭右开，:3 不含第 3 行。调试时取出来的行数总比预期少一行，先检查是不是把 end 多算了一个。

第三个坑，是布尔索引把维度压成一维。data[data[:, 2] > 5] 筛出来还是行方向的二维块，但如果条件作用在一维向量上，结果就是一维的，后面接期望二维张量的层会形状报错。取完索引顺手 print 一下 shape，是最省心的习惯。

索引，就是按位置或按条件从张量里挑数据。整数索引取单点并消掉那一维，冒号切片取连续段并保留维度，列表索引按行列配对挑散点或子矩阵，布尔索引先算真假表再筛行筛列，多维张量只是把这套规则套到更多维度上。写的时候记住三句话：冒号保维、整数消维、切片左闭右开，再配合 print 核对 shape，索引这关就过了。