张量与GPU、模型与GPU

import torch
from torch import nn
print(torch.cuda.is_available())
print(torch.cuda.device_count())

# 获取指定GPU，如果指定编号的GPU存在，则返回gpu(i)，否则返回cpu
def try_gpu(i=0):
    if torch.cuda.device_count() >= i + 1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

# 获取所有GPU，如果没有GPU，则返回CPU
def try_all_gpus():
    devices = [torch.device(f'cuda:{i}')
             for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]
print(try_gpu())
# 输出：
# cuda:0
print(try_gpu(10))
# 输出：
# cpu
print(try_all_gpus())
# 输出:
# [device(type='cuda', index=0)]

# 在CPU上创建张量
cpu_x = torch.tensor([1,2,3])
# 查看张量在什么设备上
print(cpu_x.device)
# 输出:
# cpu

# 在GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu())
# 查看张量在什么设备上
print(gpu_x.device)
# 输出:
# cuda:0

# 在指定GPU上创建张量
gpu_x = torch.ones(2, 3, device=try_gpu(0))
# 查看张量在什么设备上
print(gpu_x.device)

# 输出:
# cuda:0

# 定义一个神经网络模型,并将模型参数放在GPU上
net = nn.Sequential(nn.Linear(3, 1))
net = net.to(device=try_gpu())

# 当输入为同一个GPU上的张量时,模型将在同一个GPU上计算结果
print(net(gpu_x))
# 输出:
# tensor([[-0.1049],
#         [-0.1049]], device='cuda:0', grad_fn=<AddmmBackward0>)

# 查看模型参数位于哪个设备上
print(net[0].weight.data.device)
# 输出:
# cuda:0