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