动手学PyTorch | (18) GPU计算

到⽬前为止，我们⼀直在使⽤CPU计算。对复杂的神经网络和⼤规模的数据来说，使用CPU来计算可能不够⾼效。在本节中，我们将介绍如何使用单块NVIDIA GPU来计算。所以需要确保已经安装好了PyTorch GPU版本（又可用的GPU资源）。

可以在命令行使用nvidia-smi查看显卡信息：

nvidia-smi #Linux/Mac OS

1. 计算设备

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

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

import torch
from torch import nn

torch.cuda.is_available() # cuda是否可用

查看GPU数量：

torch.cuda.device_count() # gpu数量

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


torch.cuda.current_device() # 当前设备索引, 从0开始

根据索引号查看GPU的名字：

torch.cuda.get_device_name(0) # 返回gpu名字

2. Tensor的GPU计算

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

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

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

x = x.cuda(0)
x

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

x.device

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


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
x = torch.tensor([1, 2, 3], device=device)
# or
x = torch.tensor([1, 2, 3]).to(device)
x

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

y = x**2
y

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

z = y + x.cpu()

3. 模型的GPU计算

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

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

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

net.cuda() #net.cuda(0)
list(net.parameters())[0].device

同样的，我么需要保证输⼊到模型的Tensor和模型在同一个设备上，否则会报错：

x = torch.rand(2,3).cuda() #默认.cuda(0)
net(x)

4. 小结

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

2）PyTorch要求计算的所有输⼊数据都在内存或同一块显卡的显存上。

来源：CSDN

作者：CoreJT

链接：https://blog.csdn.net/sdu_hao/article/details/103173064

标签

cpu参数

tensor