GPU 计算

目前为止我们一直在使用 CPU 计算。对于复杂的神经网络和大规模的数据来说,使用 CPU 来计算可能不够高效。本节中,我们将介绍如何使用单块 Nvidia GPU 来计算。首先,需要确保至少有一块 Nvidia GPU 已经安装好了。然后,下载 CUDA 并按照提示设置好相应的路径。这些准备工作都完成后,下面就可以通过 nvidia-smi 来查看显卡信息了。

In [1]:
!nvidia-smi
Wed Jul 18 08:10:34 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 375.26                 Driver Version: 375.26                    |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|===============================+======================+======================|
|   0  Tesla M60           On   | 0000:00:1D.0     Off |                    0 |
| N/A   44C    P0    45W / 150W |    298MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           On   | 0000:00:1E.0     Off |                    0 |
| N/A   46C    P0    38W / 150W |    289MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
+-----------------------------------------------------------------------------+

接下来,我们需要确认安装了 MXNet 的 GPU 版本。如果装了 MXNet 的 CPU 版本,我们需要先卸载它。例如我们可以使用 pip uninstall mxnet。然后根据 CUDA 的版本安装对应的 MXNet 版本。假设你安装了 CUDA 9.1,那么我们可以通过 pip install --pre mxnet-cu91 来安装支持 CUDA 9.1 的 MXNet 版本。

计算设备

MXNet 使用 context 来指定用来存储和计算的设备,例如可以是 CPU 或者 GPU。默认情况下,MXNet 会将数据创建在主内存,然后利用 CPU 来计算。在 MXNet 中,CPU 和 GPU 可分别由 cpu()gpu() 来表示。需要注意的是,mx.cpu()(或者在括号里填任意整数)表示所有的物理 CPU 和内存。这意味着计算上会尽量使用所有的 CPU 核。但 mx.gpu() 只代表一块显卡和相应的显卡内存。如果有多块 GPU,我们用 mx.gpu(i) 来表示第 \(i\) 块 GPU(\(i\) 从 0 开始)且 mx.gpu(0)mx.gpu() 等价。

In [2]:
import mxnet as mx
from mxnet import nd
from mxnet.gluon import nn

[mx.cpu(), mx.gpu(), mx.gpu(1)]
Out[2]:
[cpu(0), gpu(0), gpu(1)]

NDArray 的 GPU 计算

默认情况下,NDArray 存在 CPU 上。因此,之前我们每次打印 NDArray 的时候都会看到 @cpu(0) 这个标识。

In [3]:
x = nd.array([1,2,3])
x
Out[3]:

[ 1.  2.  3.]
<NDArray 3 @cpu(0)>

我们可以通过 NDArray 的 context 属性来查看其所在的设备。

In [4]:
x.context
Out[4]:
cpu(0)

GPU 上的存储

我们有多种方法将 NDArray 放置在 GPU 上。例如我们可以在创建 NDArray 的时候通过 ctx 指定存储设备。下面我们将 a 创建在 GPU 0 上。注意到在打印 a 时,设备信息变成了 @gpu(0)。创建在 GPU 上时我们会只用 GPU 内存,你可以通过 nvidia-smi 查看 GPU 内存使用情况。通常你需要确保不要创建超过 GPU 内存上限的数据。

In [5]:
a = nd.array([1, 2, 3], ctx=mx.gpu())
a
Out[5]:

[ 1.  2.  3.]
<NDArray 3 @gpu(0)>

假设你至少有两块 GPU,下面代码将会在 GPU 1 上创建随机数组

In [6]:
b = nd.random.uniform(shape=(2, 3), ctx=mx.gpu(1))
b
Out[6]:

[[ 0.59118998  0.313164    0.76352036]
 [ 0.97317863  0.35454726  0.11677533]]
<NDArray 2x3 @gpu(1)>

除了在创建时指定,我们也可以通过 copytoas_in_context 函数在设备之间传输数据。下面我们将 CPU 上的 x 复制到 GPU 0 上。

In [7]:
y = x.copyto(mx.gpu())
y
Out[7]:

[ 1.  2.  3.]
<NDArray 3 @gpu(0)>
In [8]:
z = x.as_in_context(mx.gpu())
z
Out[8]:

[ 1.  2.  3.]
<NDArray 3 @gpu(0)>

需要区分的是,如果源变量和目标变量的 context 一致,as_in_context 使目标变量和源变量共享源变量的内存,

In [9]:
y.as_in_context(mx.gpu()) is y
Out[9]:
True

copyto 总是为目标变量新创建内存。

In [10]:
y.copyto(mx.gpu()) is y
Out[10]:
False

GPU 上的计算

MXNet 的计算会在数据的 context 上执行。为了使用 GPU 计算,我们只需要事先将数据放在 GPU 上面。而计算结果会自动保存在相同的 GPU 上。

In [11]:
(z + 2).exp() * y
Out[11]:

[  20.08553696  109.19629669  445.23950195]
<NDArray 3 @gpu(0)>

注意,MXNet 要求计算的所有输入数据都在同一个 CPU/GPU 上。这个设计的原因是不同 CPU/GPU 之间的数据交互通常比较耗时。因此,MXNet 希望用户确切地指明计算的输入数据都在同一个 CPU/GPU 上。例如,如果将 CPU 上的 x 和 GPU 上的 y 做运算,会出现错误信息。

当我们打印 NDArray 或将 NDArray 转换成 NumPy 格式时,如果数据不在主内存里,MXNet 会自动将其先复制到主内存,从而带来隐形的传输开销。

Gluon 的 GPU 计算

同 NDArray 类似,Gluon 的模型可以在初始化时通过 ctx 指定设备。下面代码将模型参数初始化在 GPU 上。

In [12]:
net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=mx.gpu())

当输入是 GPU 上的 NDArray 时,Gluon 会在相同的 GPU 上计算结果。

In [13]:
net(y)
Out[13]:

[[ 0.0068339 ]
 [ 0.01366779]
 [ 0.02050169]]
<NDArray 3x1 @gpu(0)>

确认一下模型参数存储在相同的 GPU 上。

In [14]:
net[0].weight.data()
Out[14]:

[[ 0.0068339]]
<NDArray 1x1 @gpu(0)>

小结

  • 通过 context,我们可以在不同的 CPU/GPU 上存储数据和计算。

练习

  • 试试大一点的计算任务,例如大矩阵的乘法,看看 CPU 和 GPU 的速度区别。如果是计算量很小的任务呢?
  • GPU 上应如何读写模型参数?

扫码直达 讨论区