动手学深度学习（Pytorch版）代码实践 -卷积神经网络-18GPU

18GPU

import torch
from torch import nn

#使用CPU
torch.device('cpu')

#使用GPU
#cuda:0和cuda是等价的
#是一个上下文管理器，用于在其上下文范围内临时设置默认CUDA设备
torch.cuda.device('cuda')

#如果有多个GPU,torch.device(f'cuda:{i}'),表示使用第i块GPU（从0开始）
#是一个设备对象，用于在创建或移动张量时明确指定目标设备
torch.device('cuda:1')

# 查询可用gpu的数量
print(torch.cuda.device_count())

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

def try_all_gpus():
    """返回所有可用的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')]
# 如果 devices 列表不为空，返回 devices 列表。
# 如果 devices 列表为空，返回一个包含单个CPU设备的列表 [torch.device('cpu')]。

print(try_gpu())
print(try_gpu(10))
print(try_all_gpus())

#张量与GPU

x = torch.tensor([1, 2, 3])
print(x.device)
# cpu
# 默认情况下，张量是在CPU上创建的。

#张量存储在GPU上
Y = torch.ones(2, 3, device=torch.device('cuda'))
X = torch.ones(2, 3, device=try_gpu())
print(Y)
print(X)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:0')

# 假设至少有两个GPU
# Y = torch.rand(2, 3, device=try_gpu(1))
# print(Y)
# tensor([[0.4860, 0.1285, 0.0440],
#         [0.9743, 0.4159, 0.9979]], device='cuda:1')

#复制
# Z = X.cuda(1)
# print(X)
# print(Z)
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:0')
# tensor([[1., 1., 1.],
#         [1., 1., 1.]], device='cuda:1')

#相加
print(X + Y)
# tensor([[2., 2., 2.],
#         [2., 2., 2.]], device='cuda:0')

# 假设变量Z已经存在于第二个GPU上。 如果我们还是调用Z.cuda(1)
# 不会复制并分配新内存
# Z.cuda(1) is Z
# True

# 神经网络与GPU
# 神经网络模型可以指定设备
net = nn.Sequential(nn.Linear(3, 1)) # 在cup上建立
net = net.to(device=try_gpu()) # 将模型参数放在GPU上

print(net(X))
# tensor([[-1.4124],
#         [-1.4124]], device='cuda:0', grad_fn=<AddmmBackward>)

#查看模型参数在哪里
print(net[0].weight.data.device)
# cuda: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
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88