pytorch使用GPU计算
在之前的blog中早已经讲过如何配置pytorch的GPU加速环境
查看GPU加速是否可用:
import torchfrom torchimport nnprint(torch.cuda.is_available())# true 查看GPU是否可用print(torch.cuda.device_count())#GPU数量, 1
torch.cuda.current_device()#当前GPU的索引, 0
torch.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 ii 块GPU及相应的显存( i ii从0开始)且cuda(0)和cuda()等价。
x.device#通过Tensor的device属性判断该Tensor所在的设备
我们也可以在创建Tensor时就指定其设备属性。
device= torch.device('cuda'if torch.cuda.is_availableelse'cpu')
x= torch.tensor([1,2,3], device=device)#or# x = torch.tensor([1,2,3]).to(device)
如果对在GPU设备上的数据进行计算,那么计算结果也会在该设备上:
y= x**2print(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和模型在同一个设备上,否则会报错