pytorch使用GPU计算

在之前的blog中早已经讲过如何配置pytorch的GPU加速环境

查看GPU加速是否可用：

import torch
from torch import nnprint(torch.cuda.is_available()) # true 查看GPU是否可用print(torch.cuda.device_count()) #GPU数量， 1torch.cuda.current_device() #当前GPU的索引， 0torch.cuda.get_device_name(0) #输出GPU名称

Tensor的GPU计算：

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

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

x.device #通过Tensor的device属性判断该Tensor所在的设备

我们也可以在创建Tensor时就指定其设备属性。

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)

如果对在GPU设备上的数据进行计算，那么计算结果也会在该设备上：

y = x**2
print(y.device) # cuda:0

但是，位于不同设备上的数据是不可以直接进行计算的，这包括在GPU上的数据不可以与在CPU上的数据直接进行计算，在不同的GPU设备上的数据也不可以直接进行计算：

z = y + x.cpu()

报错：

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

与Tensor类似，pytorch模型也可以通过.cuda()转移到设备上，而我们可以通过查看模型参数的device属性查看模型所在的设备：

module = nn.Linear(3, 1)
list(module.parameters())[0].device # device(type='cpu')

说明模型的所在的设备为cpu,将模型转移到GPU上：

module.cuda(0)
list(module.parameters())[0].device  #device=(type='cuda', index=0)

同理，我们在使用模型的时候，需要保证模型的输入的Tensor和模型在同一个设备上，否则会报错