torch笔记十三 | GPU计算

1.GPU设备

print(torch.cuda.is_available())        # 查看GPU是否可用
print(torch.cuda.device_count())        # 查看GPU数量
print(torch.cuda.current_device())      # 查看GPU索引号
print(torch.cuda.get_device_name(0))    # 根据索引号得到GPU名称

2.Tensor on GPU

import torch

# 1.tensor在CPU和CPU之间转换
x = torch.randn(3)  # 默认在CPU创建的tensor
print(x)            # tensor([-2.1135, -0.6386,  0.1393])
x = x.cuda(0)       # .cuda()将CPU上的Tensor转换到GPU上
print(x)            # tensor([ 0.3179,  1.3020, -1.8062], device='cuda:0')
print(x.device)     # .device属性查看该Tensor所在的设备

# 2.创建tensor时指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x = torch.randn(3, device=device)   # 方法1
y = torch.randn(3).to(device)       # 方法2
print(x, y)

# 注意：
# 在GPU上的数据进行运算，那么结果还是存放在GPU上
# 存储在不同位置中的数据是不可以直接进行计算的

3.model on GPU

# 1.model.cuda(0)
net = torch.nn.Linear(3, 1)
net.cuda(0)         # .cuda转换到GPU上
net_device = list(net.parameters())[0].device   # 检查模型的参数的device属性
print(net_device)

# 2.tensor和model要在同一设备
x = torch.tensor([1, 2, 3], dtype=torch.float).cuda()
print(net(x))