使用GPU前,我们首先需要查看GPU信息,这可以通过如下命令实现:
watch nvidia-smi
默认情况下Pytorch将数据创建在内存,然后利用CPU进行计算,所以我们我们需要手动设定GPU信息。接下来介绍几个相关指令
查看GPU是否可用:torch.cuda.is_available()
查看GPU数量:torch.cuda.device_count()
查看当前GPU索引号(从0开始):torch.cuda.current_device()
根据索引号查看GPU名字:torch.cuda.get_device_name(index_number:int)
默认情况下,Tensor会被存储在内存上。因此,我们打印Tensor时候看不到GPU相关标识:
x=torch.tensor([1,2,3])
print(x)
输出:
tensor([1,2,3])
使用.cuda()
可以内存中的Tensor转换到GPU上。如果有多块GPU,可以使用.cuda(i)
来表示第i块GPU所对应的显存(从0开始),注意cuda(0)
和cuda()
等价:
x=x.cuda(0)
print(x)
输出:
tensor([1, 2, 3], device='cuda:0')
我们可以通过Tensor的device属性来查看该Tensor所在的设备:
print(x.device)
输出:
cuda:0
我们可以在创建的时候就指定设备:
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
x=torch.tensor([1,2,3],device=device)
# 或者通过如下的方式指定
x=torch.tensor([1,2,3]).to(device)
print(x)
输出:
tensor([1, 2, 3], device='cuda:0')
如果对GPU上的数据进行运算,那么结果还是存放在GPU上:
y=x**2
print(y)
输出:
tensor([1, 4, 9], device='cuda:0')
注意存储在CPU上的数据无法直接与存放在GPU上的数据进行运算,会报错。
和Tensor类似,Pytorch模型也可以通过.cuda
转换到GPU上,我们可以通过检查模型参数的device
属性来查看存放模型的设备:
net=nn.Linear(3,1)
print(list(net.parameters())[0].device)
输出:
cpu
此时可见模型在CPU上,我们将其转换到GPU上:
net.cuda()
print(list(net.parameters())[0].device)
输出:
cuda:0
注意我们需要保证模型输入的Tensor和模型都在同一台设备上,否则会报错。