5.6.1 计算设备
import torch
from torch import nn
torch.device('cpu'), torch.device('cuda:0') # cuda等价于cuda:0（只有一块显卡没法试别的块号）
(device(type='cpu'), device(type='cuda', index=0))
torch.cuda.device_count() # 查询可用GPU数量
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() # 测试
(device(type='cuda', index=0),
device(type='cpu'),
[device(type='cuda', index=0)])
5.6.2 张量与GPU
x = torch.tensor([1, 2, 3])
x.device # 默认情况下张量是存储在CPU上的
device(type='cpu')
-
存储在GPU上
可以在创建张量时选择位置
X = torch.ones(2, 3, device=try_gpu()) # 只有一个GPU也就不试下面那个了
X
tensor([[1., 1., 1.],
[1., 1., 1.]], device='cuda:0')
-
复制
需要执行运算的数据需要在同一块GPU上，如果不在则可以复制过去。
# Z = X.cuda(1) # 没显卡没法试
X.cuda(0) is X # 如果已存在则不会复制
True
-
旁注
谨慎复制，并行化的瓶颈在于数据传输而不是运算速度。
5.6.3 神经网络与GPU
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu()) # 模型也可以指定位置
net(X)
tensor([[0.3092],
[0.3092]], device='cuda:0', grad_fn=<AddmmBackward0>)
net[0].weight.data.device # 所需要的数据和参数在同一设备时才可用该设备运算
device(type='cuda', index=0)
练习
（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')
大计算量任务CPU运算时长: 23190.1 ms
大计算量任务GPU运算时长: 60.0 ms
小计算量任务CPU运算时长: 1.0 ms
小计算量任务GPU运算时长: 0.0 ms
（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')
逐个记录耗费时间：48 ms
最终记录耗费时间：10 ms
（4）测量同时在两个 GPU 上执行两个矩阵乘法与在一个 GPU 上按顺序执行两个矩阵乘法所需的时间。
略