Pytorch 搭建神经网络（7）向量化与广播机制与基本索引与切片

基于《深度学习框架 Pytorch 入门与实践》陈云

参考 Github 的 pytorch-book 项目

笔记和代码存储在我的 GitHub 库中 github.com/isKage/pytorch-notes

1 向量化

向量化计算是一种特殊的并行计算方式。

  • 向量化计算：对不同的数据执行同样的一个或一批指令，或者把指令应用到一个数组或向量上，从而将多次循环操作变成一次计算。

例如：当计算乘积加和时，用 [a1,a2,...,an]T⋅[x1,x2,⋯ ,xn][a_1, a_2, ..., a_n]^T \cdot [x_1, x_2, \cdots, x_n] 替代 ∑i=1n aixi\sum_{i=1}^n\ a_i x_i 可以极大地提高计算效率。

在 Python 中，for 循环是一种及其低效的计算方法，我们往往希望能通过向量化运算并行计算。

  • 例如：利用 for 循环和直接向量相加，耗时差距大约数百倍
1
2
3
4
5
6
7
8
9
10
11
12
13
14
import torch

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

x = torch.randn(100)
y = torch.randn(100)

%timeit -n 100 sum_with_for(x, y)  # for 循环
%timeit -n 100 (x + y) # 向量化计算

综上可知，在进行计算时尽可能使用向量化计算，避免使用 for 循环。

2 广播机制

广播法则（broadcast）：快速执行向量化计算的同时不会占用额外的内存 / 显存。

2.1 NumPy 中的广播法则

    1. 广播后形状一致：所有输入数组都与形状 (shape) 最大的数组看齐，形状不足的部分在前面加 1 补齐
    1. 存在相同大小的维度：两个数组要么在某一个维度的尺寸一致，要么其中一个数组在该维度的尺寸为 1 ，否则不符合广播法则的要求
    1. 通过复制补齐：如果输入数组的某个维度的尺寸为 1 ，那么计算时沿此维度复制扩充成目标的形状大小

2.2 Pytorch 实现广播

PyTorch 支持自动广播法则，也可以手动设置：

  • unsqueeze view tensor[None] ： 为数据某一维度补 1 ，实现广播法则 1

  • expand expand_as 重复数组：实现广播法则 3 (该操作不会复制整个数组，因此不会占用额外的空间)

【注意】repeat 可以实现与 expand 类似的功能： expand 是在已经存在的 Tensor 上创建一个新维度复制数据。故 repeat 会占用额外的空间

2.3 具体使用

2.3.1 自动广播法则

以 a = torch.ones(3, 2); b = torch.zeros(2, 3, 1) 为例计算 a + b 的自动广播过程：

  • 第一步（法则 1）

a 是 2 维的，b 是 3 维的。所以先在较小的 a 前面补 1 个维度 a.unsqueeze(0) -> a.shape = (1, 3, 2)

  • 第二步（法则 2）

此时 a.shape = (1, 3, 2) 而 b.shape = (2, 3, 1) 二者在第二个维度形状一样，都为 3 。满足法则 2，存在相同形状的维度，可以进行广播

  • 第三步（法则 3）

a 和 b 在第一和第三个维度的形状不一样， 利用广播法则扩展，按照最大的形状复杂。最终变为 a.shape = b.shape = (2, 3, 2) 即，a 在第一个维度复制，b 在第三个维度复制

1
2
3
4
5
a = torch.ones(3, 2)
b = torch.zeros(2, 3, 1)

(a + b).shape
# torch.Size([2, 3, 2])

2.3.2 详细计算

1
2
3
4
a = torch.tensor([[1, 2], [3, 4], [5, 6]])
b = torch.tensor([[[-1], [-2], [-3]], [[-4], [-5], [-6]]])
a.shape, b.shape
# (torch.Size([3, 2]), torch.Size([2, 3, 1]))

加和使用自动广播机制得到结果：

1
2
3
4
5
6
7
8
9
10
print(a + b)
'''
tensor([[[ 0,  1],
         [ 1,  2],
         [ 2,  3]],

        [[-3, -2],
         [-2, -1],
         [-1,  0]]])
'''

具体 a 的广播结果：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
a =
[
    [1, 2],
    [3, 4],
    [5, 6]
].shape = (3, 2)
# -(broadcast)-> 
[
    [
        [1, 2],
        [3, 4],
        [5, 6]
    ],
    [
        [1, 2],
        [3, 4],
        [5, 6]
    ]
].shape = (2, 3, 2)

具体 b 的广播结果：

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
b = 
[
    [
        [-1],
        [-2],
        [-3]
    ],
    [
        [-4],
        [-5],
        [-6]
    ]
].shape = (2, 3, 1)
# -(broadcast)-> 
[
    [
        [-1, -1],
        [-2, -2],
        [-3, -3]
    ],
    [
        [-4, -4],
        [-5, -5],
        [-6, -6]
    ]
].shape = (2, 3, 2)

所以计算过程：

1
2
3
4
5
6
7
8
9
10
11
12
13
a + b =
[
    [
        [1 - 1, 2 - 1],
        [3 - 2, 4 - 2],
        [5 - 3, 6 - 3]
    ],
    [
        [1 - 4, 2 - 4],
        [3 - 5, 4 - 5],
        [5 - 6, 6 - 6]
    ]
].shape = (2, 3, 2)

2.3.3 repeat 复制占用额外空间

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# 比较 expand 和 repeat 的内存占用情况
a = torch.ones(1, 3) # shape = (1, 3)
print(str(a.storage().size()))
''' 3 '''

# expand 不额外占用内存，只返回一个新的视图
b = a.expand(3, 3)  # shape = (3, 3)
print(str(b.storage().size()))
''' 3 '''

# repeat 复制了原始张量
c = a.repeat(3, 3)  # shape = (3, 3)
print(str(c.storage().size()))
''' 27 '''

2.3.4 手动广播

可以先扩展成相同维度数，再扩展成相同形状。

1
2
3
4
5
6
7
8
9
10
11
12
# 手动广播
a = torch.ones(3, 2)
b = torch.zeros(2, 3, 1)

# 1. unsqueeze + expand
a.unsqueeze(0).expand(2, 3, 2) + b.expand(2, 3, 2)

# 2. view + expand
a.view(1, 3, 2).expand(2, 3, 2) + b.expand(2, 3, 2)

# 3. None + expand 【推荐】
a[None, :, :].expand(2, 3, 2) + b.expand(2, 3, 2)

3 基本索引与切片

Pytorch 的基本索引与切片与 Numpy 类似，例如：

  • 元组序列 ：在索引中直接使用一个元组序列对 Tensor 中数据的具体位置进行定位，也可以直接使用数字，即去除元组括号
  • 切片对象 ：切片形如[start:end:step]，对一个维度进行全选时可以直接使用 :
  • ... ：在索引中常用省略号来代表一个或多个维度的切片
  • None ：表示增加一个维度

3.1 元组序列

使用 tensor[x, x, x] 或 tensor[(x, x, x)] 索引

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
a = torch.Tensor([i for i in range(24)]).view(2, 3, 4)
print(a)
'''
[
	[
		[ 0.,  1.,  2.,  3.],
		[ 4.,  5.,  6.,  7.],
		[ 8.,  9., 10., 11.]],
    [
    	[12., 13., 14., 15.],
    	[16., 17., 18., 19.],
    	[20., 21., 22., 23.]
    ]
]
'''
1
2
3
4
5
6
7
# 提取位置 [0, 1, 2] 的元素 
# 等价于 a[(0, 1, 2)]
a[0, 1, 2]
'''
tensor(6.)
# (第一维度形状 2 的第 0 个, 第二维度形状 3 的第 1 个, 第三维度形状 4 的第 2 个) 所以为 6.
'''
1
2
3
4
5
6
# 第三个维度全取
# 等价于 a[(1, 1)]，a[(1, 1, )]，a[1, 1]
a[1, 1, :]
'''
tensor([16., 17., 18., 19.])
'''

3.2 ... 和 : 运算符

  • : ：使用格式 [start:end:step] 。单独使用 : 代表这个维度全取，start 和 end 为空分别表示从头开始和一直到结束，step 的默认值是 1
  • ... ：用于省略任意多个维度，可以用在切片的中间，也可以用在首尾

3.2.1 : 运算符代表全取

常见格式 tensor[:, :, start:end:step]

1
2
3
4
5
6
7
8
9
10
11
a = torch.rand(64, 3, 224, 224)

a[:, :, 0:224:4, :].shape # 第一、二、四维度全取，第三个维度取 0 到 223 间隔 4 个一取

# 省略 start 和 end 代表整个维度
a[:, :, ::4, :].shape  # 第一、二、四维度全取，第三个维度间隔 4 个一取，从开始取到结尾

"""
均为
torch.Size([64, 3, 56, 224])
"""

3.2.2 ... 代替多个维度

若 a.shape = (64, 3, 224, 224) 使用 a[..., ::4, :] 代表第一、二维度均取，相当于用 a[:, :, ::4, :]

1
2
3
4
5
6
7
# 使用 ... 代替一个或多个维度，建议一个索引中只使用一次

a[..., ::4, :].shape # 第一、二维度都取，用 ... 替代了 :, :,

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

【建议】... 在一个索引里只出现一次，否则会出现难以匹配

1
2
3
4
5
a[..., ::4, ...].shape  # 如果将最后一个维度也改为 ... 那么在匹配维度时将混乱出错

"""
torch.Size([64, 3, 224, 56])  此处第一个 ... 代表了第一、二、三维度，而第二个 ... 没有匹配
"""

3.3 None 索引

None 索引可以简单地理解为维度的扩展，在广播法则用于补充维度。即相当于 unsqueeze() 函数。不过与 unsqueeze() 相比，使用 None 可以更为直观，因为它直接在索引中表示了扩展的维度。

简单而言，【推荐】使用 None 来扩展维度

  • 增加一个维度

在增加一个维度时，None 的优势并不明显

1
2
3
4
5
6
7
print(x.unsqueeze(0).shape)  # 使用 unsqueeze 在第 0 位置补充维度
print(x[None, ...].shape)  # 直接指定 0 号位置补充维度 (... 代表后面所有维度)

"""
torch.Size([1, 3, 224, 224])
torch.Size([1, 3, 224, 224])
"""
  • 增加多个维度

如果要增加多个维度，则需要多次 unsqueeze() 并且每次还要计算索引位置，十分复杂

1
2
3
4
5
6
7
8
9
10
11
x = torch.randn(3, 3, 3)

# 变为 [1, 3, 1, 3, 1, 3]
x = x.unsqueeze(0) # [1, 3, 3, 3]
x = x.unsqueeze(2) # [1, 3, 1, 3, 3]
x = x.unsqueeze(4) # [1, 3, 1, 3, 1, 3]

x.shape
'''
torch.Size([1, 3, 1, 3, 1, 3])
'''

若使用 None 直接指定维度位置，则方便很多

1
2
3
4
5
6
7
8
9
x = torch.randn(3, 3, 3)

# 变为 [1, 3, 1, 3, 1, 3]
x = x[None, :, None, :, None, :]

x.shape
'''
torch.Size([1, 3, 1, 3, 1, 3])
'''

3.4 综合使用

【案例】计算批量 batch_size = 16 的样本的矩阵乘法： [256, 1] @ [1, 256] 最后应该得到 shape = [16, 256, 256] （其中 @ 代表矩阵乘法）

1
2
3
4
5
a = torch.arange(16 * 256).view(16, 256)
a.shape
'''
torch.Size([16, 256])
'''

3.4.1 手动扩展并计算

1
2
3
4
5
6
7
8
9
b = a.unsqueeze(1)  # b.shape = [16, 1, 256]
c = b.transpose(2, 1)  # c.shape = [16, 256, 1]
print((b @ c).shape)
print((c @ b).shape)  

'''
第一个 [16, 1, 256] @ [16, 256, 1] 得到的是 torch.Size([16, 1, 1])
第二个 [16, 256, 1] @ [16, 1, 256] 得到的是 torch.Size([16, 256, 256])
'''

这很符合我们平常的矩阵计算原则，但不是我们希望看到的，因为相乘的顺序直接导致了结果形状的不同。

3.4.2 自动广播技术

1
2
3
4
5
6
7
8
9
b = a[:, None, :]  # b.shape = [16, 1, 256]
c = a[:, :, None]  # c.shape = [16, 256, 1]
print((b * c).shape)
print((c * b).shape)

'''
两个结果均为
torch.Size([16, 256, 256])
'''

通过广播机制：形如 [16, 1, 256] 与 [16, 256, 1] 相乘时，首先因为通过 None 扩展了，故维度相同（满足法则一）；其次存在有形状的维度，第一个维度均为 16 （满足法则二）；最后相乘时，对形状 1 进行复制，复制到 256 再与同维度的数相乘（使用法则三）

所以，通过广播机制，得到的结果是相同的，无需考虑顺序问题【推荐】

Pytorch 的张量 Tensor 在进行乘法时：

  • 维度、形状不同：检查是否能进行广播，从而计算
  • 维度、形状均相同：直接进行逐元素计算，例如：
1
2
3
4
5
6
7
8
# 补充：逐元素计算
a = torch.arange(16 * 256).view(16, 256)  # a.shape = [16, 256]

(a * a).shape
'''
和 a 相同，仍然是 
[16, 256]
'''