到目前为止,我们一直在使用CPU计算。对复杂的神经网络和大规模的数据来说,使用CPU来计算可能不够高效。在本节中,我们将介绍如何使用单块NVIDIA GPU来计算。首先,需要确保已经安装好了至少一块NVIDIA GPU。然后,下载CUDA并按照提示设置好相应的路径(可参考附录中“使用AWS运行代码”一节)。这些准备工作都完成后,下面就可以通过nvidia-smi
命令来查看显卡信息了。
!nvidia-smi # 对Linux/macOS用户有效
Thu Nov 7 15:19:36 2019 +-----------------------------------------------------------------------------+ | NVIDIA-SMI 418.87.00 Driver Version: 418.87.00 CUDA Version: 10.1 | |-------------------------------+----------------------+----------------------+ | GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC | | Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. | |===============================+======================+======================| | 0 GeForce RTX 208... On | 00000000:18:00.0 On | N/A | | 45% 44C P8 16W / 280W | 31MiB / 10989MiB | 0% Default | +-------------------------------+----------------------+----------------------+ | 1 GeForce RTX 208... On | 00000000:AF:00.0 On | N/A | | 45% 42C P8 5W / 280W | 11MiB / 10989MiB | 0% Default | +-------------------------------+----------------------+----------------------+ +-----------------------------------------------------------------------------+ | Processes: GPU Memory | | GPU PID Type Process name Usage | |=============================================================================| +-----------------------------------------------------------------------------+
接下来,我们需要确认安装了PyTorch的GPU版本。安装方法见“获取和运行本书的代码”一节。运行本节中的程序需要至少2块GPU。
PyTorch可以指定用来存储和计算的设备,如使用内存的CPU或者使用显存的GPU。默认情况下,PyTorch会将数据创建在内存,然后利用CPU来计算。在PyTorch中,torch.device('cpu')
(或者在括号里填任意整数)表示所有的物理CPU和内存。这意味着,MXNet的计算会尽量使用所有的CPU核。但torch.device('cuda')
只代表一块GPU和相应的显存。如果有多块GPU,我们用torch.cuda.device_count()
来获得GPU数量,之后设置os.environ["CUDA_VISIBLE_DEVICES"] = "i"
($i$从0开始)来指定全局使用的默认GPU,或者设置device
属性为cuda:i
来指定当前使用的GPU。
import torch
from torch import nn
torch.device('cpu'), torch.device('cuda')
(device(type='cpu'), device(type='cuda'))
Tensor
的GPU计算¶在默认情况下,Tensor
存在内存上。
x = torch.Tensor([1, 2, 3])
x
tensor([1., 2., 3.])
我们可以通过Tensor
的device
属性来查看该Tensor
所在的设备。
x.device
device(type='cpu')
我们有多种方法将Tensor
存储在显存上。例如,我们可以在创建Tensor
的时候通过device
参数指定存储设备。下面我们将Tensor
变量a
创建在cuda
上。注意,在打印a
时,设备信息变成了cuda:0
。创建在显存上的Tensor
只消耗同一块显卡的显存。我们可以通过nvidia-smi
命令查看显存的使用情况。通常,我们需要确保不创建超过显存上限的数据。
a = torch.Tensor([1, 2, 3]).to('cuda')
a
tensor([1., 2., 3.], device='cuda:0')
假设至少有2块GPU,下面代码将会在cuda:1
上创建随机数组。
B = torch.rand(2, 3, device='cuda:1')
B
tensor([[0.6489, 0.3837, 0.8226], [0.1939, 0.6765, 0.7554]], device='cuda:1')
除了在创建时指定,我们也可以通过to
函数在设备之间传输数据。下面我们将内存上的Tensor
变量x
复制到cuda
上。
y = x.to('cuda')
y
tensor([1., 2., 3.], device='cuda:0')
z = x.to('cuda')
z
tensor([1., 2., 3.], device='cuda:0')
需要区分的是,如果源变量和目标变量的device
一致,to
函数使目标变量和源变量共享源变量的内存或显存。
y.to('cuda') is y
True
PyTorch的计算会在数据的device
属性所指定的设备上执行。为了使用GPU计算,我们只需要事先将数据存储在显存上。计算结果会自动保存在同一块显卡的显存上。
(z + 2).exp() * y
tensor([ 20.0855, 109.1963, 445.2395], device='cuda:0')
注意,PyTorch要求计算的所有输入数据都在内存或同一块显卡的显存上。这样设计的原因是CPU和不同的GPU之间的数据交互通常比较耗时。因此,PyTorch希望用户确切地指明计算的输入数据都在内存或同一块显卡的显存上。例如,如果将内存上的Tensor
变量x
和显存上的Tensor
变量y
做运算,会出现错误信息。当我们打印Tensor
或将Tensor
转换成NumPy格式时,如果数据不在内存里,PyTorch会将它先复制到内存,从而造成额外的传输开销。
同Tensor
类似,PyTorch的模型可以在初始化时通过to
函数转移到指定设备。下面的代码将模型参数初始化在显存上。
net = nn.Sequential(
nn.Linear(3, 1)
)
net.to('cuda')
Sequential( (0): Linear(in_features=3, out_features=1, bias=True) )
当输入是显存上的Tensor
时,PyTorch会在同一块显卡的显存上计算结果。
net(y)
tensor([0.8643], device='cuda:0', grad_fn=<AddBackward0>)
下面我们确认一下模型参数存储在同一块显卡的显存上。
net[0].weight.data
tensor([[ 0.4655, 0.4900, -0.1772]], device='cuda:0')