《动手学深度学习 Pytorch版》 5.6 GPU

5.6.1 计算设备

import torch
from torch import nn

torch.device('cpu'), torch.device('cuda:0')  # cuda等价于cuda:0（只有一块显卡没法试别的块号）
1
2
3
4

(device(type='cpu'), device(type='cuda', index=0))
1

torch.cuda.device_count()  # 查询可用GPU数量
1

1
1

def try_gpu(i=0):  #@save
    """如果存在，则返回gpu(i)，否则返回cpu()"""
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def try_all_gpus():  #@save
    """返回所有可用的GPU，如果没有GPU，则返回[cpu(),]"""
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

try_gpu(), try_gpu(10), try_all_gpus()  # 测试
1
2
3
4
5
6
7
8
9
10
11
12
13

(device(type='cuda', index=0),
 device(type='cpu'),
 [device(type='cuda', index=0)])
1
2
3

5.6.2 张量与GPU

x = torch.tensor([1, 2, 3])
x.device  # 默认情况下张量是存储在CPU上的
1
2

device(type='cpu')
1

1. 存储在GPU上

   可以在创建张量时选择位置

X = torch.ones(2, 3, device=try_gpu())  # 只有一个GPU也就不试下面那个了
X
1
2

tensor([[1., 1., 1.],
        [1., 1., 1.]], device='cuda:0')
1
2

2. 复制

   需要执行运算的数据需要在同一块GPU上，如果不在则可以复制过去。

# Z = X.cuda(1)  # 没显卡没法试
X.cuda(0) is X  # 如果已存在则不会复制
1
2

True
1

3. 旁注

   谨慎复制，并行化的瓶颈在于数据传输而不是运算速度。

5.6.3 神经网络与GPU

net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())  # 模型也可以指定位置

net(X)
1
2
3
4

tensor([[0.3092],
        [0.3092]], device='cuda:0', grad_fn=)
1
2

net[0].weight.data.device #  所需要的数据和参数在同一设备时才可用该设备运算
1

device(type='cuda', index=0)
1

练习

（1）尝试一个计算量很大的任务，比如大矩阵的乘法，看看CPU和GPU的速度差异。再尝试一个计算量很小的任务呢？

import time

start_CPU=time.time()
for i in range(1000):
    A = torch.rand(1000, 1000)
    B = torch.rand(1000, 1000)
    C = torch.matmul(A, B)
end_CPU=time.time()

start_GPU=time.time()
for i in range(1000):
    A = torch.rand(1000, 1000, device=try_gpu())
    B = torch.rand(1000, 1000, device=try_gpu())
    C = torch.matmul(A, B)
end_GPU=time.time()

print(f'大计算量任务CPU运算时长: {round((end_CPU - start_CPU)*1000, 2)} ms')
print(f'大计算量任务GPU运算时长: {round((end_GPU - start_GPU)*1000, 2)} ms')

start_CPU=time.time()
A = torch.rand(5, 5)
B = torch.rand(5, 5)
C = torch.matmul(A, B)
end_CPU=time.time()

start_GPU=time.time()
A = torch.rand(5, 5, device=try_gpu())
B = torch.rand(5, 5, device=try_gpu())
C = torch.matmul(A, B)
end_GPU=time.time()

print(f'小计算量任务CPU运算时长: {round((end_CPU - start_CPU) * 1000, 2)} ms')
print(f'小计算量任务CPU运算时长: {round((end_GPU - start_GPU) * 1000, 2)} ms')

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

大计算量任务CPU运算时长: 23190.1 ms
大计算量任务GPU运算时长: 60.0 ms
小计算量任务CPU运算时长: 1.0 ms
小计算量任务GPU运算时长: 0.0 ms
1
2
3
4

---

（2）我们应该如何在GPU上读写模型参数？

使用 net.to() 函数迁移模型到 GPU 上即可。

---

（3）测量计算 1000 个 100*100 矩阵乘法所需的时间，并记录输出矩阵的弗罗贝尼乌斯范数，一次记录一个结果，而不是在GPU上保存日志并进传输最终结果。

A = torch.rand(100, 100, device=try_gpu())

start = time.time()
for i in range(1000):
    A = torch.mm(A, A)
    B = torch.norm(A)  # 逐个记录
end = time.time()
print(f'逐个记录耗费时间：{round((end - start) * 1000)} ms')

A = torch.rand(100, 100, device=try_gpu())
start = time.time()
for i in range(1000):
    A = torch.mm(A, A)
B = torch.norm(A)  # 最终记录
end = time.time()
print(f'最终记录耗费时间：{round((end - start) * 1000)} ms')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

逐个记录耗费时间：48 ms
最终记录耗费时间：10 ms
1
2

---

（4）测量同时在两个 GPU 上执行两个矩阵乘法与在一个 GPU 上按顺序执行两个矩阵乘法所需的时间。

略