《深度学习框架PyTorch：入门与实战》第6章：向量化思想实战——广播法则、高级索引与 einsum/einops

本文以开源仓库 pytorch-book 第 6 章（Chapter6/Chapter6.md）为骨架，系统讲解 PyTorch 中的向量化计算思想：从向量化与广播法则入手，深入 Tensor 底层存储结构（Storage / size / storage_offset / stride），再到基本索引与高级索引（整数数组索引、布尔数组索引），最后落到 einsum / einops 两种爱因斯坦操作，并用目标检测 IoU、RoI Align、反向 Unique 三个深度学习实际问题演示向量化思想的实战价值。读完本文，你将掌握"用形状分析驱动运算、用索引替代循环"的向量化编程能力，并能直接复现书中全部示例（对应可执行版本见 Chapter6/chapter6.ipynb，验证环境为 PyTorch 2.14.0 + Python 3.12）。

6.1 向量化简介：为什么能快数百倍

向量化计算是一种特殊的并行计算方式。一般来说，程序在同一时间内只执行一个操作，而并行计算可以在同一时间内执行多个操作。向量化计算是指对不同的数据执行同样的一个或一批指令，或者把指令应用到一个数组或向量上，从而将多次循环操作变成一次计算。向量化计算的关键是分析张量的形状：一旦确定了张量形状，如何操作就呼之欲出。

向量化操作可以极大地提高科学运算的效率。尽管 Python 本身是一门高级语言、使用简便，但其中存在着许多低效的操作，例如 for 循环等。因此，在科学计算中应当极力避免使用 Python 原生的 for 循环，尽量使用向量化数值运算。下面通过一个最简单的加法例子进行对比：

import warnings
warnings.filterwarnings("ignore")

import torch as t

# 定义for循环完成加法操作
def for_loop_add(x, y):
    result = []
    for i, j in zip(x, y):
        result.append(i + j)
    return t.tensor(result)

x = t.zeros(100)
y = t.ones(100)
%timeit -n 100 for_loop_add(x, y)
%timeit -n 100 (x + y)   # + 是向量化计算

Out:100 loops, best of 3: 786 µs per loop
    100 loops, best of 3: 2.57 µs per loop

从上面的例子中可以看出，for 循环和向量化计算之间存在数百倍的速度差距。其根本原因在于：for 循环在 Python 解释器中逐元素解释执行（每个元素还要先包装成 Python 对象），而向量化的 + 运算直接调用 PyTorch 底层内建函数（build-in function），这些函数由 C/C++ 实现，内部使用了向量化计算的思想并通过底层优化实现了高效计算。因此在日常编程中，读者应该养成向量化的编程习惯：

• 尽量调用内建函数，而不是自己写 for 循环；
• 避免对较大的 Tensor 进行逐元素的遍历操作；
• 先分析输入输出的形状，再选择对应的向量化操作。

所有示例均可在 Chapter6/chapter6.ipynb 中直接复现（书中 In/Out 格式与 notebook cell 一一对应）。

6.2 广播法则（Broadcast）

广播法则（broadcast）是科学计算中经常使用的一个技巧，它在快速执行向量化计算的同时不会占用额外的内存/显存。NumPy 与 PyTorch 中的广播法则定义如下三条：

1. **所有输入数组都与形状（shape）最大的数组看齐，形状不足的部分在前面加 1 补齐。**例如两个数组形状分别为 (3, 4) 和 (2, 3, 1)，前一个会自动补齐为 (1, 3, 4)。
2. **两个数组要么在某一个维度的尺寸一致，要么其中一个数组在该维度的尺寸为 1，否则不符合广播法则的要求。**例如形状为 (1, 3, 4) 与 (2, 3, 1) 的两个数组可以计算，而 (1, 3, 4) 与 (2, 3, 2) 不能计算。
3. **如果输入数组的某个维度的尺寸为 1，那么计算时沿此维度复制扩充成目标的形状大小。**例如 (1, 3, 4) 与 (2, 3, 1) 两个数组会复制扩充成 (2, 3, 4)。

虽然 PyTorch 已经支持了自动广播法则，但为了更直观、更不容易出错，建议通过以下两种方式的组合手动实现广播法则：

• unsqueeze、view 或者 tensor[None]：为数据某一维度补 1，实现第一条广播法则；
• expand 或者 expand_as：重复数组，实现第三条广播法则。该操作不会复制整个数组，因此不会占用额外的空间。

注意： repeat 可以实现与 expand 类似的功能，但二者本质不同——expand 是在已经存在的 Tensor 上创建一个新的视图（view），repeat 会将相同的数据复制多份，因此会占用额外的空间。

下面用 storage().size() 直接对比 expand 与 repeat 的内存占用：

# 比较expand和repeat的内存占用情况
a = t.ones(1, 3)
print("原始储存占用：" + str(a.storage().size()))
# expand不额外占用内存，只返回一个新的视图
b = a.expand(3, 3)
print("expand储存占用：" + str(b.storage().size()))
# repeat复制了原始张量
c = a.repeat(3, 3)
print("repeat储存占用：" + str(c.storage().size()))

Out:原始储存占用：3
    expand储存占用：3
    repeat储存占用：27

可见 expand 前后 storage 大小不变（仍是 3），而 repeat(3, 3) 把数据真正复制了 9 份（占用 27）。这在处理大 Tensor 时对内存/显存的节约是显著的。

接下来看自动广播与手动广播的对照。自动广播只需要直接运算：

a = t.ones(3, 2)
b = t.zeros(2, 3, 1)

# 自动广播法则
# 第一步：a是2维的，b是3维的，所以先在较小的a前面补1个维度，
#        即：a.unsqueeze(0)，a的形状变成（1，3，2），b的形状是（2，3，1）
# 第二步：a和b在第一和第三个维度的形状不一样，同时其中一个为1，
#        利用广播法则扩展，两个形状都变成了（2，3，2）
(a + b).shape

Out:torch.Size([2, 3, 2])

手动广播则用 None（等价于 view(1, ...)）与 expand 分两步显式完成：

# 手动广播法则，下面两行操作是等效的，推荐使用None的方法
# a.view(1, 3, 2).expand(2, 3, 2) + b.expand(2, 3, 2)
a[None, :, :].expand(2, 3, 2) + b.expand(2, 3, 2)

Out:tensor([[[1., 1.],
         [1., 1.],
         [1., 1.]],

        [[1., 1.],
         [1., 1.],
         [1., 1.]]])

广播法则的工程意义在于：它让不同形状的张量在"逻辑上"对齐后直接参与向量化运算，既不需要显式写循环，也不需要为复制数据付出内存代价。

6.3 索引操作：从基本索引到高级索引

索引（indexing）是高维数组中的常用操作，形如 tensor[indices]。根据 indices 的不同类型，索引可以分为基本索引（Basic Slicing and Indexing）和高级索引（Advanced Indexing）。要理解二者的差异，必须先理解 Tensor 的底层实现。

6.3.1 Tensor 的底层实现：Storage、size、storage_offset 与 stride

Tensor 的底层使用 C/C++ 实现，它可以用于存储数据，还可以定义如何对数据进行操作。Tensor 与 Python list 最大的区别是：Tensor 利用一个连续的内存区域去存储数据，这些数据都未经过封装；而 list 中的每个数据都会被封装成 PyObject 对象，每个对象独立分配内存，离散地存储在内存当中。

在 PyTorch 底层使用 Storage 类来管理 Tensor 的内存区域，它使用一个一维数组来存储数据。Tensor 通过修改内部的 size、storage_offset、stride 等属性让这个一维数组"看起来像一个多维数组"（获得多个 Tensor 的实例），实际上这些 Tensor 指向同一个 Storage 区域。下面举例说明：

a = t.arange(6).view(3, 2)
b = a.reshape(2, 3)   # 改变形状
c = a.transpose(1, 0) # 转置
d = a[:2, 1]          # 切片
# a、b、c、d四个实例指向同一个storage
id(a.storage()) == id(b.storage()) == id(c.storage()) == id(d.storage())

Out:True

# 发生改变的实际是三个内部属性：size、storage_offset和stride
print(a.size(), a.storage_offset(), a.stride())
print(b.size(), b.storage_offset(), b.stride())
print(c.size(), c.storage_offset(), c.stride())
print(d.size(), d.storage_offset(), d.stride())

Out:torch.Size([3, 2]) 0 (2, 1)
    torch.Size([2, 3]) 0 (3, 1)
    torch.Size([2, 3]) 0 (1, 2)
    torch.Size([2]) 1 (2,)

关于这三个属性的说明：

• size：控制 Tensor 每个维度上的取值范围。切片操作、reshape 操作等都会修改 Tensor 实例的 size 属性。
• storage_offset：Tensor 实例起始元素对应存储区 Storage 的索引。部分 Tensor 实例只使用了一部分 Storage，该属性用来控制每个实例的起始位置。
• stride：一个元组，stride 的第 k 个元素表示 Tensor 的第 k 个维度中两个元素之间的内存间隔。这一概念使得 Tensor 的高级索引运算变得更加高效。

图6-1 Tensor 的底层结构

图 6-1 解释了一个高维 Tensor 是怎样从一个一维数组上存取数据的：通过 storage_offset 和 stride 属性，结合维度信息即可计算出任意位置元素在一维 Storage 上的索引 offset，计算方法如式 (6.1) 所示：

\text{offset}=\sum_{k=0}^{N-1}\text{stride}_{k}\times \text{dim}_{k}+\text{storage_offset}\tag{6.1}

其中，$\text{dim}$ 表示所取元素所在的维度，$N$ 表示 Tensor 的维度大小。下面举例说明如何应用该式进行计算：

a = t.arange(12).view(3, 4)
a, a.size(), a.storage_offset(), a.stride()

Out:(tensor([[ 0,  1,  2,  3],
             [ 4,  5,  6,  7],
             [ 8,  9, 10, 11]]), torch.Size([3, 4]), 0, (4, 1))

如果想要获取第二行第三列的元素 a[1, 2]，那么可以通过式 (6.1) 得出：$\text{offset}=4\times1+1\times2+0=6$，也就是 Storage 中的第 7 个元素，这与 a[1, 2] 得到的值相同。stride 元组的计算也极为简单，它可以视为多维数组向一维数组映射的权重，它的取值只与当前 Tensor 实例的 size 属性有关，计算方法如式 (6.2) 所示：

$$\text{stride}_k=\prod_{j=k+1}^{N-1} \text{DimNum}_{j}\tag{6.2}$$

其中，$\text{DimNum}_{j}=\text{self.shape}[j]$，stride 元组的末尾元素默认为 1。

注意：必须在 tensor.is_contiguous = True 的时候，才可以使用式 (6.2) 计算 stride（非连续 Tensor 的 stride 由操作语义决定，如转置会交换 stride）。

# a的size设为2×3×4×5
a = t.rand(2, 3, 4, 5)
# 可以通过式(6.2)计算得到，stride0=3×4×5=60，stride1=4×5=20，stride2=5，stride3=1
a.stride()

Out:(60, 20, 5, 1)

小结：对 Tensor 的许多操作（reshape、转置、切片等）都可以通过改变 size、stride、offset 等属性实现，更改这些属性前后的 Tensor 共享同一个存储区域，这样的设计模式可以有效节省内存。

6.3.2 基本索引（Basic Slicing and Indexing）

PyTorch 中 Tensor 的索引和 NumPy 数组的索引类似，通过索引操作可以定位到数据的具体位置，也可以进行切片操作。基本索引是将 Python 的切片概念扩展到 N 维的形式，通过基本索引生成的所有 Tensor 都是原始 Tensor 的视图（view），它们和原始 Tensor 共享 Storage，只改变 stride、offset 等属性。基本索引有以下四种形式：

形式	说明
元组序列（Tuple）	Tensor[(IntA, IntB, IntC)] 或 Tensor[IntA, IntB, IntC]，后者是前者的语法糖
切片对象（Slicing）	Tensor[start:stop:step]，对一个维度全选时直接用 :
省略号（...）	代表一个或多个维度的切片，等效于一个或多个 :
None	与 NumPy 中的 newaxis 相同，起到增加一个维度的作用

① 元组序列

a = t.arange(1, 25).view(2, 3, 4)
a

Out:tensor([[[ 1,  2,  3,  4],
         [ 5,  6,  7,  8],
         [ 9, 10, 11, 12]],

        [[13, 14, 15, 16],
         [17, 18, 19, 20],
         [21, 22, 23, 24]]])

# 提取位置[0, 1, 2]的元素
# 等价于a[(0, 1, 2)]（保留括号的元组形式）
a[0, 1, 2]

Out:tensor(7)

# 固定前两个维度，提取第三个维度上的切片（省略第三个参数）
# 等价于a[(1, 1)]，a[(1, 1, )]，a[1, 1, :]
a[1, 1]

Out:tensor([17, 18, 19, 20])

注意：a[0, 1, 2] 与 a[[0, 1, 2]]、a[(0, 1, 2),] 并不等价，后面两个不满足基本索引的条件（既不是元组序列又不是切片对象），它们属于高级索引的范畴（见 6.3.3 节）。

② 切片对象 : 与省略号 ...

• : 常用于对一个维度进行操作，基本语法是 start:end:step。单独使用 : 代表全选该维度；start 和 end 为空分别表示从头开始和一直到结束；step 的默认值是 1。
• ... 用于省略任意多个维度，可以用在切片的中间，也可以用在首尾。一般一个索引中只使用一次 ...；至于它到底代表多少个 :，取决于数组的维度——最终索引元素的总个数要与数组维度数一致。

a = t.rand(64, 3, 224, 224)
print(a[:, :, 0:224:4, :].shape) # 第三个维度间隔切片
# 省略start和end代表整个维度
print(a[:, :, ::4, :].shape)

Out:torch.Size([64, 3, 56, 224])
    torch.Size([64, 3, 56, 224])

# 使用...代替一个或多个维度，建议一个索引中只使用一次
a[..., ::4, :].shape
# a[...,::4,...].shape  # 如果将最后一个维度也改为...，那么在匹配维度时将混乱出错

Out:torch.Size([64, 3, 56, 224])

切片对象还有几个值得注意的细节：

• 若 idx 的数目小于 Tensor 的维度（Tensor.ndim），默认在后面补多个 : 或一个 ...（这一点对基本索引和高级索引都适用）。
• 整数索引与切片索引的维度差异：当索引为整数 i 与切片对象 i:i+1 时，虽然取到的值相同，但结果的维度不同——直接使用整数 i 会比使用切片对象少一个维度。例如对三维 Tensor，a[:, 1] 的形状是 torch.Size([2, 2])，而 a[:, 1:2] 的形状是 torch.Size([2, 1, 2])。
• 关于负步长：NumPy 支持 step 为负数的反向索引，但 PyTorch 的切片语法不支持，可改用 t.flip 或高级索引（如 a[t.arange(8, 2, -3)]）实现。

③ None 索引

在 PyTorch 的源码中，None 索引经常被使用。None 索引可以直观地表示维度的扩展，在广播法则中充当 1 的作用。使用 None 索引本质上与使用 unsqueeze 函数是等价的，都能起到扩展维度的作用；但在维度较多、或需要对多个维度先扩展再进行矩阵计算时，None 索引会更加清晰直观：

a = t.rand(2, 3, 4, 5)
# 在最前面加一个维度，下面两种写法等价
print(a.unsqueeze(0).shape)
print(a[None, ...].shape)

Out:torch.Size([1, 2, 3, 4, 5])
    torch.Size([1, 2, 3, 4, 5])

# 在原有的四个维度中均插入一个维度，成为（2,1,3,1,4,1,5）
# unsqueeze方法，每成功增加一个维度，都需要重新计算下一个需要增加的维度位置
b = a.unsqueeze(1)
b = b.unsqueeze(3)
b = b.unsqueeze(5)
b.shape

Out:torch.Size([2, 1, 3, 1, 4, 1, 5])

# None索引方法，直接在需要增加的维度上填写None即可
a[:, None, :, None, :, None, :].shape

Out:torch.Size([2, 1, 3, 1, 4, 1, 5])

下面举一个使用 None 索引的实战例子。假设 Tensor a 是一组图像的 feature maps 经过全连接层后的结果，维度是 $\text{batch\_size} \times \text{features}$，现在需要构建每张图像不同维度特征之间的乘积矩阵，形如 $\text{batch\_size} \times \text{features} \times \text{features}$。

最直观的想法是矩阵乘法：将 $n \times 1$ 维的矩阵与 $1 \times n$ 维的矩阵相乘，结果就是 $n \times n$ 维的矩阵。因此，将 Tensor a 的第二个维度扩展并转置得到 $a^T$，再使用矩阵乘法即可：

# 假设batch_size为16，features为256
a = t.arange(16 * 256).view(16, 256)
a1 = a.unsqueeze(1)      # shape：16×1×256
a_T = a1.transpose(2, 1) # shape：16×256×1
a_matrix = a_T @ a1      # @表示矩阵的乘法
a_matrix.shape

Out:torch.Size([16, 256, 256])

而使用 None 索引和广播法则解决这个问题将更加直观——完全不需要显式构造转置：

b = a[:, :, None] * a[:, None, :]
c = a[:, None, :] * a[:, :, None]
# 读者可以思考一下，b和c有什么区别？
b.shape, c.shape

Out:(torch.Size([16, 256, 256]), torch.Size([16, 256, 256]))

# 没有触发异常，说明两种方式得到的结果是一样的
assert t.equal(b, c)
assert t.equal(a_matrix, b)

从结果可以看出，虽然生成 b 和 c 的写法不一样，但最终结果是一样的。这是因为使用 None 索引进行维度扩展时，扩展的维度是 1；当维度中有 1 存在时，Tensor 相乘会触发广播法则，在计算时自动填充，所以最终结果一致。这正是"广播法则 + None 索引"组合威力的直观体现。

6.3.3 高级索引（Advanced Indexing）

高级索引是更能体现向量化思想的操作。与基本索引相比，高级索引的触发条件不同，返回值也有所不同——高级索引的返回值是原始 Tensor 的拷贝（copy）而不是视图（view），这是与基本索引最重要的区别。常见的高级索引遵循以下规律：

1. 索引是整数类型或者布尔类型的 Tensor；
2. 索引是 List，如 Tensor[[0, 1]]（注意与基本索引 Tensor[(0, 1)] 区分）；
3. 索引是元组序列，但里面至少包含一个整数类型或布尔类型的 Tensor（或 List），如 Tensor[1, [1, 2, 3]]；
4. 特例：为保持向后兼容，如果索引对象是任何包含切片对象、省略号对象或 None 对象的非张量序列（如 List），那么会触发基本索引而非高级索引（这种情况较少见，也不建议使用）。

a = t.arange(12).view(2, 2, 3)
a[0, 1].shape          # 基本索引，等价于a[(0, 1)]，返回视图
a[[0, 1]].shape        # 高级索引，索引是List，返回拷贝
a[t.tensor([0, 1])].shape  # 高级索引，索引是Tensor
a[([0, 1], 1)].shape   # 高级索引，索引是Tuple但包含List/Tensor
a[[0, None]].shape     # 特殊情况：None触发基本索引，等价于a[0, None]

Out:torch.Size([3])        # a[0, 1]
    torch.Size([2, 2, 3])  # a[[0, 1]]
    torch.Size([2, 2, 3])  # a[t.tensor([0, 1])]
    torch.Size([2, 3])     # a[([0, 1], 1)]
    torch.Size([1, 2, 3])  # a[[0, None]]

整数数组索引（Integer Array Indexing）

整数数组索引根据 $N$ 维索引选取数组中的任意项。一般情况下需要先确定输入输出 Tensor 的形状，因为所有的整数索引都有一个相对固定的模式：$\text{tensor}[\text{index}_1, \text{index}_2, ..., \text{index}_N]$。

• 其中 $N$ 的大小必须小于等于这个 Tensor 的维度（tensor.ndim）。
• 如果经过索引操作后得到的 Tensor 形状是 $M_1 \times M_2 \times ... \times M_K$，那么所有索引 $\text{index}_1 ... \text{index}_N$ 的形状都必须是 $M_1 \times M_2 \times ... \times M_K$，同时输出的第 $[m_1, m_2, ..., m_K]$ 个元素是 $\text{tensor}[\text{index}_1[m_1, ..., m_K], \text{index}_2[m_1, ..., m_K], ..., \text{index}_N[m_1, ..., m_K]]$。
• 如果 index 的形状不完全相同但满足广播法则，那么它们会自动对齐成一样的形状，从而完成整数数组索引操作；对于不能广播或不能得到相同形状的索引，无法进行整数数组索引操作。

相同形状的索引：

a = t.arange(12).view(3, 4)
# 相同形状的index索引
# 获取索引为[1,0]、[2,2]的元素
a[t.tensor([1, 2]), t.tensor([0, 2])]

Out:tensor([ 4, 10])

不同形状、满足广播的索引（[1,2] 扩展为列、[0,2] 扩展为行，得到 2×2 结果）：

# 不相同形状的index索引，满足广播法则
# 获取索引为[1,0]、[2,0]、[1,2]、[2,2]的元素
a[t.tensor([1, 2])[None, :], t.tensor([0, 2])[:, None]]

Out:tensor([[ 4,  8],
        [ 6, 10]])

高级索引与基本索引混合使用时，基本索引（如切片对象、省略号、None 等）会将高级索引切分成多个区域。假设高级索引 idx1, idx2, idx3 的形状都是 $M \times N$，tensor[idx1, idx2, idx3] 的输出形状即为 $M \times N$。如果将部分高级索引替换为基本索引，那么会先计算高级索引部分的维度，然后补齐输出结果的维度，通常有以下几种情况：

• 所有高级索引都处于相邻的维度（如 tensor[idx1, :, :] 或 tensor[:, idx2, idx3]）：直接将高级索引所在区域的维度转换成高级索引的维度，Tensor 的其他维度按基本索引正常计算；
• 基本索引将多个高级索引划分到不同区域（如 tensor[idx1, :, idx3]）：统一将高级索引的维度放在输出 Tensor 维度的开头，剩下部分补齐基本索引的维度——因为所有高级索引并不相邻，无法确定高级索引的维度应该替换 Tensor 的哪些维度，因此统一放到开头位置。

下面举例说明：

a = t.arange(24).view(2, 3, 4)
idx1 = t.tensor([[1, 0]]) # shape 1×2
idx2 = t.tensor([[0, 2]]) # shape 1×2

# 所有的高级索引相邻
a[:, idx1, idx2].shape

Out:torch.Size([2, 1, 2])

# 手动计算输出形状
# a的第一个维度保留，后两个维度是索引维度
a.shape[0], idx1.shape

Out:(2, torch.Size([1, 2]))

a = t.arange(120).reshape(2, 3, 4, 5)
# 中间两个维度替换成高级索引的维度
a[:, idx1, idx2, :].shape

Out:torch.Size([2, 1, 2, 5])

# 高级索引被划分到不同区域
# 高级索引的维度放在输出维度的最前面，剩下的维度依次补齐
a[idx1, :, idx2].shape

Out:torch.Size([1, 2, 3, 5])

a[:, idx1, :, idx2].shape

Out:torch.Size([1, 2, 2, 4])

再看一个更复杂的混合情况：

a = t.arange(108).reshape(2, 3, 3, 3, 2)
idx1 = t.tensor([[1, 0]])
idx2 = t.tensor([[0, 1]])
a[:, idx1, :2, idx2, 0].shape

Out:torch.Size([1, 2, 2, 2])

在这个例子中，:、:2 和 0 都是基本索引，它们将高级索引分到了两个区域。首先将高级索引的维度放在输出 Tensor 维度的开头，然后计算基本索引、将结果放置在高级索引维度之后，最终得到索引后的结果。

另外，当基本索引和高级索引的总数小于 Tensor 的维度数（$N < \text{tensor.ndim}$）时，会自动在最后补上 ... 操作：

a = t.arange(24).view(2, 3, 4, 1)
a[:, t.tensor([[1, 0]])].shape

Out:torch.Size([2, 1, 2, 4, 1])

# 等价于在末尾补上...:
a[:, t.tensor([[1, 0]]), ...].shape

Out:torch.Size([2, 1, 2, 4, 1])

整数数组索引的本质是：将索引数组对应位置的数字组合成源 Tensor 实例的索引，再根据索引值和 Tensor 的 size、storage_offset、stride 属性计算出 Storage 空间上的真实索引，最后返回结果。如果索引数组的个数小于 Tensor 的维度数，那么缺少的部分需要对整个轴进行完整的切片，再重复上述过程。归纳为两种情况：

• 索引数组的个数 N 等于 Tensor 的维度数 tensor.ndim：索引输出的形状等价于 index.shape，输出的每一个元素等价于 tensor[index_1[i], index_2[i], ..., index_N[i]]；
• 索引数组的个数 N 小于 Tensor 的维度数 tensor.ndim：类似于切片操作，把这个切片当做索引操作的结果。

a = t.arange(12).view(3, 4)
print(a)
print(a[[2, 0]])  # 索引数组个数小于a的维度数
# 索引数组个数等于a的维度数
# 获取索引为[1,3]、[2,2]、[0,1]的元素
print(a[[1, 2, 0], [3, 2, 1]])

Out:tensor([[ 0,  1,  2,  3],
        [ 4,  5,  6,  7],
        [ 8,  9, 10, 11]])
    tensor([[ 8,  9, 10, 11],
        [ 0,  1,  2,  3]])
    tensor([ 7, 10,  1])

# 输出形状取决于索引数组的形状
# 获取索引为[0,1]、[2,3]、[1,3]、[0,1]的元素
idx1 = t.tensor([[0, 2], [1, 0]])
idx2 = t.tensor([[1, 3], [3, 1]])
a[idx1, idx2]

Out:tensor([[ 1, 11],
        [ 7,  1]])

# 输出主对角线的元素
a[[0, 1, 2], [0, 1, 2]]

Out:tensor([ 0,  5, 10])

# 获取四个角的元素
idx1 = t.tensor([[0, 0], [-1, -1]])
idx2 = t.tensor([[0, -1], [0, -1]])
a[idx1, idx2]

Out:tensor([[ 0,  3],
        [ 8, 11]])

注意：不能先对索引数组进行组合，再进行索引操作。 如果只有一个索引数组，那么该数组会被视为第一个维度上的索引。例如下面这种"先组合再索引"的写法就是错误示范：idx = t.tensor([idx1, idx2]); a[idx]——即使不报错，结果也不是想要的，因为只索引了第一个维度，后面的维度直接进行了切片。

此外，在索引数组中还可以使用 :、... 和 None 等索引方式，进行更加复杂的索引操作：

print(a[[2, 0], 1:3])
print(a[[1, 0], ...])
print(a[[1, 0], None, [2, 3]])

Out:tensor([[ 9, 10],
        [ 1,  2]])
    tensor([[4, 5, 6, 7],
        [0, 1, 2, 3]])
    tensor([[6],
        [3]])

布尔数组索引（Boolean Array Indexing）

在高级索引中，如果索引数组的类型是布尔型，那么就会使用布尔数组索引。布尔类型的数组对象可以通过比较运算符产生：

a = t.arange(12).view(3, 4)
idx_bool = t.rand(3, 4) > 0.5
idx_bool

Out:tensor([[ True,  True,  True,  True],
        [False, False, False, False],
        [False, False, False, False]])

a[idx_bool] # 返回idx_bool中为True的部分

Out:tensor([0, 1, 2, 3])

布尔数组索引常用于对特定条件下的数值进行修改。例如，对一个 Tensor 中的所有正数进行乘 2 操作。最直观的方法是写一个双层 for 循环，遍历整个 Tensor：

# 利用for循环
a = t.tensor([[1, -3, 2], [2, 9, -1], [-8, 4, 1]])
for i in range(a.shape[0]):
    for j in range(a.shape[1]):
        if a[i, j] > 0:
            a[i, j] *= 2
a

Out:tensor([[ 2, -3,  4],
        [ 4, 18, -1],
        [-8,  8,  2]])

使用布尔数组索引只需一行，语义也更清晰：

# 利用布尔数组索引
a = t.tensor([[1, -3, 2], [2, 9, -1], [-8, 4, 1]])
a[a > 0] *= 2
a

Out:tensor([[ 2, -3,  4],
        [ 4, 18, -1],
        [-8,  8,  2]])

布尔索引还可以与聚合函数组合实现条件筛选，例如返回所有行和小于 3 的行：

# 返回Tensor中所有行和小于3的行
a = t.tensor([[1, -3, 2], [2, 9, -1], [-8, 4, 1]])
row_sum = a.sum(-1)
a[row_sum < 3, :]

Out:tensor([[ 1, -3,  2],
        [-8,  4,  1]])

小试牛刀：用高级索引实现卷积

在深度学习中，最常用的操作是卷积。除了调用 PyTorch 封装好的函数，也可以自行编写函数实现卷积功能。根据卷积的定义，只需遍历整个图像，依次获取与卷积核相乘的子块，相乘求和后即可得到卷积结果。

方法一：img2col 思路（含 for 循环的基准实现）。将整张图像提前转换成与卷积核相乘的子块，再把每个子块的维度展平，此时形状从 $(C_{\text{in}}, H_{\text{in}}, W_{\text{in}})$ 变为 $(C_{\text{in}} \times K \times K, H_{\text{out}}, W_{\text{out}})$；同时卷积核形状变为 $(C_{\text{out}}, C_{\text{in}} \times K \times K)$，最终通过矩阵乘法得到卷积结果：

import torch.nn as nn

def Conv_base(img, filters, stride, padding):
    '''
    img: 输入图像 channel×height×width
    filters: 卷积核 input_channel×output_channel×height×width
    stride: 卷积核的步长
    padding: 边缘填充的大小
    '''
    Cin, Hin, Win = img.shape
    _, Cout, K, _ = filters.shape

    # 计算卷积输出的大小
    Hout = ((Hin + 2 * padding - K) / stride).long() + 1
    Wout = ((Win + 2 * padding - K) / stride).long() + 1

    # 首先构建一个输出的样子
    col = torch.zeros(Cin, K, K, Hout, Wout)
    # 通过padding的值将imgs进行扩充
    imgs = nn.ZeroPad2d(padding.item())(img)
    for h in range(Hout):
        for w in range(Wout):
            h1 = int(h * stride.item())
            w1 = int(w * stride.item())
            col[..., h, w] = imgs[:, h1:h1+K, w1:w1+K]
    col = col.view(Cin*K*K, Hout*Wout)
    # 将卷积核变形
    filters = filters.transpose(1, 0).view(Cout, Cin*K*K)
    out_img = (filters @ col).view(Cout, Hout, Wout)
    return out_img

img = t.arange(36).view(1, 6, 6)
filters = t.ones(1, 1, 3, 3) / 9
stride, padding = t.tensor(1.), t.tensor(0)
output = Conv_base(img, filters, stride, padding)
print("进行卷积操作的图像为：\n", img[0])
print("卷积核为：\n", filters[0][0])
print("卷积后的结果为：\n", output)

Out:进行卷积操作的图像为：
    tensor([[ 0,  1,  2,  3,  4,  5],
        [ 6,  7,  8,  9, 10, 11],
        [12, 13, 14, 15, 16, 17],
        [18, 19, 20, 21, 22, 23],
        [24, 25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34, 35]])
卷积核为：
    tensor([[0.1111, 0.1111, 0.1111],
        [0.1111, 0.1111, 0.1111],
        [0.1111, 0.1111, 0.1111]])
卷积后的结果为：
    tensor([[ 7.0000,  8.0000,  9.0000, 10.0000],
        [13.0000, 14.0000, 15.0000, 16.0000],
        [19.0000, 20.0000, 21.0000, 22.0000],
        [25.0000, 26.0000, 27.0000, 28.0000]])

方法二：整数数组索引实现卷积（完全向量化）。卷积运算还可以从整数数组索引的角度考虑：在确定输入形状和卷积核大小后，输出 Tensor 的形状也就固定了，剩下的计算都可以通过整数数组索引解决。

假设输入形状是 $(C_{\text{in}}, H_{\text{in}}, W_{\text{in}})$，卷积核形状是 $(C_{\text{in}}, C_{\text{out}}, K, K)$，输出形状是 $(C_{\text{out}}, H_{\text{out}}, W_{\text{out}})$。卷积核可以变形为 $(C_{\text{out}}, C_{\text{in}} \times K \times K)$；如果输入能变形为 $(C_{\text{in}} \times K \times K, H_{\text{out}}, W_{\text{out}})$，那么通过矩阵乘法可直接得到输出。此时构造两个形状为 $(K, K, H_{\text{out}}, W_{\text{out}})$ 的索引 $\text{index}_1$ 和 $\text{index}_2$，执行 $\text{input}[:, \text{index}_1, \text{index}_2]$ 即可完成"取子块"的变形。其中索引由卷积核下标的索引（meshgrid(K1, K1)）与输出 Tensor 下标的索引（meshgrid(H, W)）通过广播相加组合而成：

# 定义输入图像的参数
Cin, Hin, Win = 1, 6, 6
img = t.arange(Cin*Hin*Win).view(Cin, Hin, Win).float()
# 定义卷积核的大小和输出通道数
K, Cout, stride = 3, 1, 1
filter = t.ones((Cin, Cout, K, K)).float() / (Cin * K * K)

def Conv(img, filter, stride=1, padding=0):
    '''
    img: 形状为 channel_in×height×width
    filter: 形状为 channel_in×channel_out×kernel×kernel
    '''
    Cin, Hin, Win = img.shape
    Cout, K = filter.shape[1], filter.shape[2]
    # 计算卷积输出图像的参数，默认stride=1，padding=0
    Hout = ((Hin + 2 * padding - K) / stride).long() + 1
    Wout = ((Win + 2 * padding - K) / stride).long() + 1

    # 卷积核下标的索引
    K1 = t.arange(-(K//2), K//2+1)
    idx11, idx12 = t.meshgrid(K1, K1)
    # 输出Tensor下标索引
    H = t.linspace(K//2, K//2+stride*(Hout-1), Hout).long()
    W = t.linspace(K//2, K//2+stride*(Wout-1), Wout).long()
    idx21, idx22 = t.meshgrid(H, W)
    # 两种索引的组合形式
    idx1 = idx11[:, :, None, None] + idx21[None, None, :, :]
    idx2 = idx12[:, :, None, None] + idx22[None, None, :, :]

    # 改变filter的形状，便于接下来的矩阵相乘
    filter = filter.transpose(0, 1).reshape(Cout, Cin*K*K)
    # 输入图像经过整数数组索引后改变成适合矩阵乘法的形状
    img = img[:, idx1, idx2].reshape(Cin*K*K, Hout*Wout)
    # 矩阵相乘得到卷积后的结果
    res = (filter @ img).reshape(Cout, Hout, Wout)
    return res

Conv(img, filter, stride)

Out:tensor([[[ 7.0000,  8.0000,  9.0000, 10.0000],
         [13.0000, 14.0000, 15.0000, 16.0000],
         [19.0000, 20.0000, 21.0000, 22.0000],
         [25.0000, 26.0000, 27.0000, 28.0000]]])

两种方法的输出完全一致。区别在于：整数数组索引版本的卷积只分析输入和输出的形状，中间变换全部通过索引操作完成，彻底消灭了 for 循环，实现更加简洁、目的性更强。

进一步地，可以利用这个向量化卷积模拟平均池化：构造一个 3×3 的全 1 卷积核（除以 K²），并把步长设为卷积核大小 K，即可得到长宽都缩为原来的 1/3 的池化结果（仓库中的示例图像见 Chapter6/imgs/input.png 与 Chapter6/imgs/output.png）：

import torch as t
from PIL import Image
import numpy as np

# 读取一张图像，模拟池化操作
img = Image.open('./imgs/input.png')
img = t.tensor(np.array(img)).float()
# 将img的形状从h×w×c转化为c×h×w
img = img.transpose(0, 2).transpose(1, 2)
Cout, K, Cin = 3, 3, 3
filter_pool = t.zeros(Cin, Cout, K, K)
# 初始化卷积核
filter_pool[t.arange(Cin), t.arange(Cin), :, :] = 1./K/K
# 利用卷积去模拟池化，将步长设置为卷积核的大小即可
out = Conv(img, filter_pool, stride=K)
# 将输出结果转换为h×w×c的形状，用于显示
out = out.transpose(1, 0).transpose(1, 2).long()
Image.fromarray(np.array(out, dtype=np.uint8))

可以看出，经过池化操作的图像长宽都变成了原来的三分之一——这正是 stride=K 的卷积与池化在数学上的等价关系。

6.3.4 einsum / einops：爱因斯坦操作

在高级索引之外，还有一类特殊的方法：爱因斯坦操作。下面介绍两种常用的爱因斯坦操作 einsum 和 einops，它们被广泛用于向量、矩阵和张量的运算，可以用非常简单的方式表示较为复杂的多维 Tensor 之间的运算。

einsum：爱因斯坦求和

在数学界，有一个由爱因斯坦提出的求和约定，能够有效处理坐标方程。einsum 基于这个法则，省略求和符号和默认成对出现的下标，从而完成对向量、矩阵和张量的运算：

# 转置操作
import torch as t
a = t.arange(9).view(3, 3)
b = t.einsum('ij->ji', a) # 直接交换两个维度
print(a)
print(b)

Out:tensor([[0, 1, 2],
        [3, 4, 5],
        [6, 7, 8]])
    tensor([[0, 3, 6],
        [1, 4, 7],
        [2, 5, 8]])

# 求和操作
a = t.arange(36).view(3, 4, 3)
b = t.einsum('ijk->', a) # 所有元素求和
b

Out:tensor(630)

# 输入a是三维张量，下标是i、j和k
# 输入b是三维张量，下标是j、i和m
# 输出c是二维张量，下标是k和m
# 操作：a和b中相同的下标i和j是求和下标，结果保留了k和m
a = t.arange(36).view(3, 4, 3)
b = t.arange(24).view(4, 3, 2)
c = t.einsum('ijk,jim->km', a, b)
c

Out:tensor([[2640, 2838],
        [2772, 2982],
        [2904, 3126]])

# 多个张量之间的混合运算
a = t.arange(6).view(2, 3)
b = t.arange(3)
# 矩阵对应维度相乘，b进行了广播
t.einsum('ij,j->ij', a, b)

Out:tensor([[ 0,  1,  4],
        [ 0,  4, 10]])

# 直观表达矩阵的内积和外积
a = t.arange(6).view(2, 3)
b = t.arange(6).view(3, 2)
c_in = t.einsum('ij,ij->', a, a)   # 内积，结果是一个数
c_out = t.einsum('ik,kj->ij', a, b) # 外积，矩阵乘法的结果
print(c_in)
print(c_out)

Out:tensor(55)
    tensor([[10, 13],
        [28, 40]])

einsum 的核心是用字符串下标描述张量运算的输入输出契约：出现多次的下标（如 ij,j->ij 中的 j）默认按求和约定对齐，出现在输出中的下标被保留，不出现在输出中的下标被求和。这让复杂运算（转置、求和、缩并、广播）在一行内完成表达。

einops：rearrange 与 reduce

除了爱因斯坦求和，其他的爱因斯坦操作都封装在 einops 中。它支持 NumPy、PyTorch、Chainer、TensorFlow 等多种框架的数据格式。在爱因斯坦操作中，多次转置操作不再使用 tensor_x.transpose(1, 2).transpose(2, 3)，而是用更直观的方式：rearrange(tensor_x, 'b c h w -> b h w c') 代替。

einops 有很多复杂操作，这里讲解最常见、最直观的用法，并分析如何在深度学习框架中高效使用（更详细的内容可参考 einops 官方文档）。

from einops import rearrange, reduce
a = t.rand(16, 3, 64, 64) # batch × channel × height × weight

# 转置操作
rearrange(a, 'b c h w -> b h w c').shape

Out:torch.Size([16, 64, 64, 3])

# 融合部分维度（flatten）
y = rearrange(a, 'b c h w -> b (h w c)')
y.shape

Out:torch.Size([16, 12288])

在视觉 Transformer 中的应用（ViT patch 化）。Transformer 起源于 NLP，近年来不少研究者尝试将其应用于计算机视觉领域。爱因斯坦操作凭借其便捷、直观的特点在视觉 Transformer 中得到了广泛应用。以 Vision Transformer（ViT）为例：假设输入是 256×256 的彩色图像，需要将其划分成 8×8=64 个块，每个块有 32×32×3=3072 个像素，用 einops 实现如下：

img = t.randn(1, 3, 256, 256)
x = rearrange(img, 'b c (h p1) (w p2) -> b (p1 p2) (h w c)', p1=8, p2=8)
x.shape

Out:torch.Size([1, 64, 3072])

这个操作把高度和宽度维度各自"劈开"成 (h, p1) 与 (w, p2)，再重组为 token 序列，一行代码完成了传统实现中需要 reshape + transpose + contiguous 多步操作才能完成的工作。

在 Sequential 中直接搭建网络。搭建神经网络时，会频繁地使用 flatten 操作（如卷积层生成 feature maps 后需要压缩维度以便输入全连接层）。传统写法用 view 完成，而 view 只能写在 forward 中；网络结构复杂时通常用 Sequential 整合模块，此时 einops 操作展现出优势——通过 einops.layers.torch.Rearrange 可以把维度转换直接嵌入网络，在一个 Sequential 中构建整个网络：

from torch import nn
import torch.nn.functional as F

# 传统写法：view只能在forward中使用
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

# einops写法：将维度转换嵌入Sequential
import torch.nn as nn
from einops.layers.torch import Rearrange

ConvImproved = nn.Sequential(
    nn.Conv2d(1, 10, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Conv2d(10, 20, kernel_size=5),
    nn.MaxPool2d(kernel_size=2),
    nn.ReLU(),
    nn.Dropout2d(),
    Rearrange('b c h w -> b (c h w)'),
    nn.Linear(320, 50),
    nn.ReLU(),
    nn.Dropout(),
    nn.Linear(50, 10),
    nn.LogSoftmax(dim=1)
)

多种 concat 方式与 Space-to-Depth / Depth-to-Space。rearrange 可以在括号中指定重组顺序，实现各种拼接方式：

# 多种concat方式
y_h = rearrange(a, 'b c h w -> (b h) w c') # concat维度h
y_c = rearrange(a, 'b c h w -> h w (b c)') # concat通道c
y_h.shape, y_c.shape

Out:(torch.Size([1024, 64, 3]), torch.Size([64, 64, 48]))

在很多网络结构中，需要提取通道间或空间像素间的信息，完成通道的部分维度和空间的部分维度之间的转化，einops 可以直观地完成（这也是视觉 Transformer、sub-pixel 卷积等结构中的常见操作）：

# Space to Depth
b = t.rand(16, 32, 64, 64)
s2d = rearrange(b, 'b c (h h0) (w w0) -> b (h0 w0 c) h w', h0=2, w0=2)
# Depth to Space
d2s = rearrange(b, 'b (c h0 w0) h w -> b c (h h0) (w w0)', h0=2, w0=2)
print("Space to Depth: ", s2d.shape)
print("Depth to Space: ", d2s.shape)

Out:Space to Depth:  torch.Size([16, 128, 32, 32])
    Depth to Space:  torch.Size([16, 8, 128, 128])

注意：使用 rearrange 进行维度间的相互转化时，一定要注意维度的顺序。 如果提取的顺序有误，即使形状一致，数值也会不同：

res1 = rearrange(b, 'b c (h h0) (w w0) -> b (h0 w0 c) h w', h0=2, w0=2)
res2 = rearrange(b, 'b c (h0 h) (w0 w) -> b (h0 w0 c) h w', h0=2, w0=2)
print(res1.shape == res2.shape)
print(res1.equal(res2))

Out:True
    False

res1 与 res2 形状相同，但由于提取维度时的顺序（(h h0) 与 (h0 h) 的元素排列）不同，数值并不相同。另外，rearrange 还可以通过括号内下标交换实现通道 shuffle：

# 下面两种操作都可以打乱通道进行shuffle
y1 = rearrange(b, 'b (c1 c2 c) h w -> b (c2 c1 c) h w', c1=2, c2=4)
y2 = rearrange(b, 'b (c0 c) h w -> b (c c0) h w', c0=4)

reduce：求和、池化与归一化。除了 rearrange，常见的 einops 操作还有 reduce，它常用于求和、求均值等操作，也用于搭建卷积神经网络中的池化层：

# 对空间像素求和
y = reduce(a, 'b c h w -> b c', reduction='sum')
y.shape # 对h和w维度求和

Out:torch.Size([16, 3])

# 池化操作：最大池化
max_pooling = reduce(a, 'b c (h h0) (w w0) -> b c h w',
                     reduction='max', h0=2, w0=2)
max_pooling.shape

Out:torch.Size([16, 3, 32, 32])

# 全局平均池化
global_avg_pooling = reduce(a, 'b c h w -> b c', reduction='mean')
global_avg_pooling.shape

Out:torch.Size([16, 3])

# 通道归一化：将求取均值的维度结果设为1，运算时即可广播
# 对每张图像进行通道归一化
y1 = a - reduce(a, 'b c h w -> b c 1 1', 'mean')
# 对整个batch的图像进行通道归一化
y2 = a - reduce(a, 'b c h w -> 1 c 1 1', 'mean')

einops 的所有操作都支持反向传播，可以有效地嵌入到深度学习模型框架中（这在需要端到端训练的模型里非常关键）：

x0 = t.rand(16, 3, 64, 64)
x0.requires_grad = True
x1 = reduce(x0, 'b c h w -> b c', reduction='max')
x2 = rearrange(x1, 'b c -> c b')
x3 = reduce(x2, 'c b -> ', reduction='sum')

x3.backward()
x0.grad.shape

Out:torch.Size([16, 3, 64, 64])

梯度能够正确地回传到原始输入 x0，形状与原 Tensor 一致，证明 einops 的算子全程可微。

6.4 小试牛刀：使用向量化思想解决实际问题

本节将向量化思想应用于深度学习中三个经典问题，进一步体会向量化与高级索引的实战价值。

6.4.1 Box_IoU：目标检测的评价指标

Box_IoU 是目标检测任务中最基本的评价指标，即模型预测的检测结果框（predicted bbox）与原始标记框（ground truth）之间的交并比，如图 6-4 所示，其中 $(\text{Lx}_i, \text{Ly}_i)$ 和 $(\text{Rx}_i, \text{Ry}_i)$ 分别表示回归框的左上角和右下角坐标。

目标检测中的 IoU 评价指标可以用公式 (6.3) 计算：

$$\text{IoU}=\frac{\text{predicted bbox} \cap \text{ground truth}}{\text{predicted bbox} \cup \text{ground truth}}\tag{6.3}$$

其中分子部分表示两个框的交集，分母部分表示两个框的并集，二者的比值即为 IoU 的结果。最直观的计算方法是遍历每个 predicted bbox 与 ground truth，逐个计算 IoU：

# 框的左上和右下的坐标
Lx1, Ly1, Rx1, Ry1 = 0, 3, 3, 1 # 回归框
Lx2, Ly2, Rx2, Ry2 = 2, 2, 5, 0 # ground truth

area_reg = abs(Lx1 - Rx1) * abs(Ly1 - Ry1) # 预测框的面积
area_gt = abs(Lx2 - Rx2) * abs(Ly2 - Ry2)  # ground truth的面积

w = min(Rx1, Rx2) - max(Lx1, Lx2)
h = min(Ry1, Ry2) - max(Ly1, Ly2)

if w <= 0 or h <= 0:
    print("两个框没有交集")
    IoU = 0
else:
    inter = h * w
    IoU = inter / (area_reg + area_gt - inter) # 获取IoU的值
IoU

Out:0.09090909090909091

在实际检测问题中，对同一个 ground truth 会预测得到一组框，依次计算每个框与 ground truth 的 IoU 稍显麻烦。这里可以用向量化思想，同时计算一组预测框与 ground truth 的 IoU——关键在于利用广播法则构造"交集边长"，再用 min 和 abs 等向量化算子一次算完：

# 默认所有的predicted bbox框都和ground truth有交集
# 模拟四个候选框和一个gt框
Reg = t.tensor([[0, 3, 3, 1], [1, 4, 3, 0], [1, 2, 5, 0], [-1, 2, 4, 0]])
GT = t.tensor([2, 2, 5, 0])

# 通过广播法则计算交集框的边长
dist = (Reg - GT.reshape(1, 4)[:, [2, 3, 0, 1]]).abs()
# 计算这一组候选框分别与gt框的交集面积
inter = t.min(dist[:, [0, 2]] * dist[:, [1, 3]], dim=1)[0]

# 计算每个返回框的面积
Reg_edge = (Reg[:, [0, 1]] - Reg[:, [2, 3]]).abs()
area_reg = Reg_edge[:, 0] * Reg_edge[:, 1]

# 计算Ground Truth框的面积
area_gt = (GT[0] - GT[2]).abs() * (GT[1] - GT[3]).abs()

# 计算这一组返回框与ground truth框之间的IoU值
IoU = inter.float() / (area_reg + area_gt - inter)
IoU

Out:tensor([0.0909, 0.1667, 0.7500, 0.3333])

一行输出得到 4 个候选框对同一 gt 框的 IoU 值（0.0909、0.1667、0.7500、0.3333）。注意这里的技巧：GT.reshape(1, 4)[:, [2, 3, 0, 1]] 将 gt 坐标重排成 [Rx, Ry, Lx, Ly] 的顺序，与 Reg（顺序为 [Lx, Ly, Rx, Ry]）相减取绝对值后，[:, [0, 2]] 对应宽方向差、[:, [1, 3]] 对应高方向差，再两两相乘并取 min 即得交集面积。整个过程完全避开了 Python 循环。

6.4.2 RoI Align：基于双线性插值的区域特征池化

使用过 Faster RCNN 的读者对 RoI Pooling 和 RoI Align 算法一定十分熟悉：这两种算法实现了对具有不同特征大小的输入区域，输出具有相同大小的输出特征。

• RoI Pooling 的目的是将不同尺寸的 RoI（Region of Interest）区域映射到固定大小的 feature map 上进行池化操作，以便后续进行分类和回归。但由于 feature map 大小固定，RoI Pooling 中需要进行两次量化：一是将 RoI 对应特征图与原始特征图的网格单元对齐，也就是将候选框边界量化为整数点坐标值；二是将上一步得到的量化 RoI 特征进一步细化为量化的空间单元（bin），也就是对每一个单元的边界进行量化。经过两次量化后，候选框已经与最开始的回归位置有一定偏差，即不匹配问题（misalignment），这对小目标来说非常致命。
• RoI Align 算法全程保持浮点数，通过双线性插值算法计算最后的池化结果，解决了 RoI Pooling 中的不匹配问题。

RoI Align 算法的流程如下：

1. 遍历每个候选区域，将网格单元平均划分为 $H \times W$ 个子网格区域（注意：与 RoI Pooling 不同，此处无需进行量化操作）；
2. 对每个区域选择 4 个采样点，进一步划分为 4 个子区域，取得每个子区域的中点；
3. 使用双线性插值算法计算这 4 个中点的像素值大小；
4. 通过池化操作对每个子区域进行聚合，从而得到最终的 feature map。

双线性插值算法的原理如图 6-5 所示：根据周围四个坐标点 A、B、C、D 的坐标和权重，得到点 P 的插值结果。首先根据 A、B、C、D 计算出点 M 和点 N 的坐标与权重，然后根据 M 和 N 得到最终的插值结果。

具体计算方法如式 (6.4) 所示：

$$\begin{array}{l} F(M)=\frac{x-x_{1}}{x_{2}-x_{1}} F(A)+\frac{x_{2}-x}{x_{2}-x_{1}} F(B) \\ F(N)=\frac{x-x_{1}}{x_{2}-x_{1}} F(C)+\frac{x_{2}-x}{x_{2}-x_{1}} F(D) \\ F(P)=\frac{y-y_{2}}{y_{1}-y_{2}} F(N)+\frac{y_{1}-y}{y_{1}-y_{2}} F(M) \end{array}\tag{6.4}$$

其中 $F(·)$ 表示 · 的权重。由于 A、B、C、D 均为相邻像素的中间位置坐标，所以 $x_2-x_1=1$、$y_2-y_1=1$，式 (6.4) 可以转化为权重的形式，如式 (6.5) 所示：

$$\begin{aligned}F(P) &=\left(y-y_{1}\right)\left(x-x_{1}\right) F(A)+\left(y-y_{1}\right)\left(x_{2}-x\right) F(B) \\ &+\left(y_{2}-y\right)\left(x-x_{1}\right) F(C)+\left(y_{2}-y\right)\left(x_{2}-x\right) F(D)\end{aligned}\tag{6.5}$$

用向量化思想实现 RoI Align。先分析输入输出形状：假设经过卷积生成的 feature map 形状是 $C \times H \times W$，对应的 bboxes 形状是 $N \times 4$（$N$ 代表 bbox 个数，4 是左上和右下的坐标值），目标生成的池化结果是 $C \times N \times n \times n$。显然，需要在左上角和右下角坐标的基础上构造出 $n \times n$ 的结果，然后对 feature map 进行整数数组索引，得到一个 $C \times N \times n \times n$ 的结果，最后相乘完成 RoI Align。由于第 3 步用到了双线性插值，还需要在计算中扩展一个用于计算插值的维度，最后消去即可：

def RoI_Align(BBoxes, feature_map, n):
    N = BBoxes.shape[0] # BBox的个数
    C = feature_map.shape[0]
    BBoxes = BBoxes.float()
    feature_map = feature_map.float()
    # 获取BBox边上等分点的坐标
    Xs = BBoxes[:, [0]] + t.arange(n+1).float() / n * (BBoxes[:, [2]] - BBoxes[:, [0]])
    Ys = BBoxes[:, [1]] + t.arange(n+1).float() / n * (BBoxes[:, [3]] - BBoxes[:, [1]])

    bins = t.linspace(0, n * n - 1, n * n)
    idx_x, idx_y = (bins // n).long(), (bins % n).long()

    # 获取每个进行pooling的bin的中心的坐标
    x = Xs[:, idx_x] + (Xs[:, [1]] - Xs[:, [0]]) / 2
    y = Ys[:, idx_y] + (Ys[:, [1]] - Ys[:, [0]]) / 2
    x_floor = x.floor().long().unsqueeze(-1)
    y_floor = y.floor().long().unsqueeze(-1)

    # 获取最近四个点的相对坐标，左上(0,0)，左下(1,0)，右下(1,1)，右上(0,1)
    corner_x = t.tensor([0, 1, 1, 0])
    corner_y = t.tensor([0, 0, 1, 1])
    idx_feature_x, idx_feature_y = x_floor + corner_x, y_floor + corner_y
    # 获取feature map中的实际值
    points = feature_map[:, idx_feature_x, idx_feature_y]

    # 利用双线性插值法计算权重ratio
    ratio_x = (x - x.floor()).unsqueeze(-1)
    ratio_y = (y - y.floor()).unsqueeze(-1)
    index1 = corner_x[None, None, :].float()
    index2 = corner_y[None, None, :].float()
    ratio = t.abs(index1 - ratio_x) * t.abs(index2 - ratio_y)

    # 计算最终pooling的结果
    res = t.einsum('cnkp,nkp->cnkp', points, ratio)
    RoIAlign = t.einsum('cnkp->nck', res).reshape(N, C, n, n)
    return RoIAlign

# pooling后的大小为n*n
BBoxes = t.tensor([[50, 40, 100, 120], [30, 40, 90, 100]])
feature_map = t.arange(40000).view(1, 200, 200)
# 最终得到2组roi的结果
RoI_Align(BBoxes, feature_map, n=7).shape

Out:torch.Size([2, 1, 7, 7])

这段实现非常浓缩，值得逐行解读：

• Xs、Ys：利用广播，为每个 bbox 一次性生成 n+1 个等分点的坐标（形状 N × (n+1)）；
• bins 与 idx_x / idx_y：把 n×n 个 bin 展开成一维下标，再映射回 (行, 列)；
• x / y：每个 bin 的中心点坐标（注意原实现用相邻等分点差的一半近似 bin 中心）；
• x_floor / y_floor 与 corner_x / corner_y：取出中心点周围 4 个整数格点坐标，通过 feature_map[:, idx_feature_x, idx_feature_y] 完成整数数组索引一次性取出所有采样点值（形状 C × N × n² × 4）；
• ratio：按式 (6.5) 计算双线性插值权重；
• 两次 einsum：先逐点加权（cnkp,nkp->cnkp），再聚合求和（cnkp->nck）并 reshape 成 N × C × n × n。

书中还给出了可视化示例（对应仓库图片 Chapter6/imgs/ROI.png、Chapter6/imgs/output1.png、Chapter6/imgs/output2.png）：在一幅图中框出八个字母作为 RoI Align 处理对象，将 n 调大以便观察，可以看到每个框内的内容被正确映射到固定大小的池化结果中。这套"整数数组索引取点 + einsum 加权聚合"的思路，就是现代检测框架中 RoI Align 高效实现的核心。

6.4.3 反向 Unique：用高级索引还原张量

PyTorch 中有一个 unique 函数，它返回输入 Tensor 中不同的元素组成的 unique list，同时返回输入 Tensor 对应于这个 unique list 的索引。那么，拿到 unique list 和对应的索引后，能否还原出输入的 Tensor 呢？

最简单的思路是遍历索引逐个生成输入 Tensor 对应位置的元素再组合，过程比较繁琐。根据整数数组索引的思路：这个索引的 size 和目标 Tensor 的 size 是一致的，因此可以直接使用整数数组索引对原始 Tensor 进行构建：

# 随机生成一组形状为(10, 15, 10, 5)、1~9数字组成的张量
a = t.randint(1, 10, (10, 15, 10, 5))
# 获取输出的unique list和索引
output, inverse_indices = t.unique(a, return_inverse=True)
# 通过整数数组索引 还原原始tensor
a_generate = output[inverse_indices]

a_generate.equal(a)

Out:True

还原的 Tensor 值与原始值一致，即 output[inverse_indices] 这一行高级索引就完整完成了反向 unique 操作，避免了耗时较长的循环遍历。这个技巧在需要"去重 + 还原"的数据预处理流水线（如词汇表映射、聚类标签还原）中非常实用。

6.5 小结

本章对 PyTorch 中的向量化计算与高级索引进行了详细介绍：

• 向量化是高维数据处理时提升计算效率的关键手段——同样的加法，向量化比 for 循环快数百倍；
• 广播法则在零内存/显存开销的前提下让不同形状的张量直接参与运算，配合 unsqueeze / None / expand 可以手动、直观地完成维度对齐；
• Tensor 底层结构（Storage + size + storage_offset + stride）揭示了视图与拷贝的本质区别，也解释了基本索引（返回视图）与高级索引（返回拷贝）的差异；
• 整数数组索引与布尔数组索引是向量化取数的核心工具，可用于实现卷积、池化、条件修改、数据还原等操作；
• einsum 与 einops 用字符串契约描述张量运算，在视觉 Transformer、Sequential 网络搭建、池化与归一化等场景中兼具简洁性与可读性，且全程可微、可直接嵌入模型训练；
• 三个实战案例（Box_IoU、RoI Align、反向 Unique）展示了"先分析输入输出形状，再选择向量化算子"的通用方法论。

建议读者对照 Chapter6/chapter6.ipynb 逐 cell 复现全部示例，并在解决实际问题时尝试使用向量化思想进行编程，从而切实提高程序的运行效率。文中所有代码基于 PyTorch 2.x 验证通过；示例所涉及的原始图像素材均位于仓库 Chapter6/imgs 目录下。