4.6 GPU计算

发表于 更新于 分类于 ，

计算设备

PyTorch可以指定用来存储和计算的设备，如使用内存的CPU或者使用显存的GPU。默认情况下，PyTorch会将数据创建在内存，然后利用CPU来计算。

用torch.cuda.is_available()查看GPU是否可用:

1
2
3
4
import torch
from torch import nn

torch.cuda.is_available() # 输出 True

查看GPU数量：

1
torch.cuda.device_count() # 输出 1

查看当前GPU索引号，索引号从0开始：

1
torch.cuda.current_device() # 输出 0

根据索引号查看GPU名字:

1
torch.cuda.get_device_name(0) # 输出 'GeForce GTX 1050'

Tensor的GPU计算

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

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

输出：

1
tensor([1, 2, 3])

使用.cuda()可以将CPU上的Tensor转换（复制）到GPU上。如果有多块GPU，用.cuda(i)来表示第 \(i\) 块GPU及相应的显存（\(i\)从0开始）且cuda(0)和cuda()等价。

1
2
x = x.cuda(0)
print(x)

输出：

1
tensor([1, 2, 3], device='cuda:0')

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

1
print(x.device)

输出：

1
device(type='cuda', index=0)

可以直接在创建的时候就指定设备。

1
2
3
4
5
6
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

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

输出：

1
tensor([1, 2, 3], device='cuda:0')

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

1
2
y = x**2
print(y)

输出：

1
tensor([1, 4, 9], device='cuda:0')

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

1
z = y + x.cpu()

会报错:

1
RuntimeError: Expected object of type torch.cuda.LongTensor but found type torch.LongTensor for argument #3 'other'

模型的GPU计算

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

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

输出：

1
device(type='cpu')

可见模型在CPU上，将其转换到GPU上:

1
2
net.cuda()
print(list(net.parameters())[0].device)

输出：

1
device(type='cuda', index=0)

同样的，需要保证模型输入的Tensor和模型都在同一设备上，否则会报错。

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

输出：

1
2
tensor([[-0.5800],
        [-0.2995]], device='cuda:0', grad_fn=<ThAddmmBackward>)