[pytorch、学习] - 4.6 GPU计算

参考

4.6 GPU计算

到目前为止,我们一直使用CPU进行计算。对复杂的神经网络和大规模数据来说,使用CPU来计算可能不够高效。
在本节中,将要介绍如何使用单块NIVIDA GPU进行计算

4.6.1 计算设备

PyTorch可以指定用来存储和计算的设备,如果用内存的CPU或者显存的GPU。默认情况下,PyTorch会将数据创建在内存,然后利用GPU来计算。
用torch.cuda.is_available()查看GPU是否可用:

import torch
from torch import nn

torch.cuda.is_available()

# 查看GPU数量
torch.cuda.current_device()

PS: GPU的数序号从0开始计算

# 根据索引查看GPU名字
torch.cuda.get_device_name(0)

4.6.2 Tensor的GPU计算

默认情况下,Tensor会被存在内存上。因此,之前我们每次打印Tensor的时候看不到GPU相关标识

x = torch.tensor([1,2,3])
x

使用.cuda()可以将CPU上的Tensor转换(复制)到GPU上。如果有多块GPU,我们用.cuda(i)来表示第i块GPU

x = x.cuda()

我们可以通过Tensor的device属性来查看该Tensor所在的设备

x.device

我们可以在创建的时候就指定设备

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

x = torch.tensor([1, 2, 3], device=device)
print(x)

x= torch.tensor([1,2,3]).to(device)

print(x)

如果对在GPU上的数据进行运算，那么结果还是存放在GPU上

y = x**2
y

需要注意的是，存储在不同位置中的数据是不可以直接进行计算的。即存放在CPU上的数据不可以直接与存放在GPU上的数据进行运算，位于不同GPU上的数据也是不能直接进行计算的

4.6.3 模型的GPU计算

同Tensor类似,PyTorch模型也可以通过 .cuda转换到GPU上。我们可以通过device属性来查看存放模型的设备

net = nn.Linear(3, 1)
list(net.parameters())[0].device

将模型转到GPU上

net.cuda()
list(net.parameters())[0].device

模型和参数需要同时在GPU上才能计算

x =torch.rand(2, 3).cuda()
net(x)