首先在python中安装对应的pytorch以及cuda版本
- 查看cuda是否安装成功,gpu是否可用
import torch
from torch import nn
print(torch.duda.is_available()) # true 查看GPU是否可用
print(torch.cuda.device_count()) #GPU数量, 2
torch.cuda.current_device() #当前GPU的索引, 0
torch.cuda.get_device_name(0) #输出GPU名称
- Tensor(张量)中使用GPU加速,将cpu上的计算转换到GPU上计算
x = torch.tensor([1, 2, 3;4,5,6])
x = x.cuda(0)
- 报错
- RuntimeError: Expected object of type torch.cuda.LongTensor but found type torch.LongTensor for argument #3 ‘other’
- 这个说明模型在cpu上,cpu上的模型不能再GPU上计算,GPU上的模型也不能在CPU上计算
- 解决办法把张量转换到cpu上 self.nerwork.cpu()或者转换到gpu上 self.cuda()
如何在torch中使用GPU
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 判断是否有GPU
model.to(device)
mytensor = my_tensor.to(device)#调用GPU
多GPU下使用
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 判断是否有GPU
if torch.cuda.device_count() > 1:
model = nn.DataParallel(model)
model.to(device)
mytensor = my_tensor.to(device)
多GPU模型在单GPU下加载失败的问题
多GPU训练的模型保存和读取的时候用
model = nn.DataParallel(model,device_ids=device_ids)
单GPUdevice_ids=[0]
,多GPU就更改这个为device_ids=[0,1]