GPU 计算¶

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

In [1]:

!nvidia-smi

Fri Dec 28 20:17:41 2018
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 396.37                 Driver Version: 396.37                    |
|-------------------------------+----------------------+----------------------+
| 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           Off  | 00000000:00:1D.0 Off |                    0 |
| N/A   47C    P0    46W / 150W |      0MiB /  7618MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   53C    P0    44W / 150W |      0MiB /  7618MiB |     99%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

接下来，我们需要确认安装了 MXNet 的 GPU 版本。如果装了 MXNet 的 CPU 版本，我们需要先卸载它，例如使用pip uninstall mxnet命令，然后根据 CUDA 的版本安装相应的 MXNet 版本。假设你安装了 CUDA 9.0，可以通过pip install mxnet-cu90来安装支持 CUDA 9.0 的 MXNet 版本。运行本节中的程序需要至少两块 GPU。

计算设备¶

MXNet 可以指定用来存储和计算的设备，例如使用内存的 CPU 或者使用显存的 GPU。默认情况下，MXNet 会将数据创建在内存，然后利用 CPU 来计算。在 MXNet 中，mx.cpu()（或者在括号里填任意整数）表示所有的物理 CPU 和内存。这意味着 MXNet 的计算会尽量使用所有的 CPU 核。但mx.gpu()只代表一块 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 存在内存上。因此，之前我们每次打印 NDArray 的时候都会看到@cpu(0)这个标识。

In [3]:

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

Out[3]:

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

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

In [4]:

x.context

Out[4]:

cpu(0)

GPU 上的存储¶

我们有多种方法将 NDArray 存储在显存上。例如我们可以在创建 NDArray 的时候通过ctx参数指定存储设备。下面我们将 NDArray 变量a创建在gpu(0)上。注意到在打印a时，设备信息变成了@gpu(0)。创建在显存上的 NDArray 只消耗同一块显卡的显存。我们可以通过nvidia-smi命令查看显存的使用情况。通常，我们需要确保不创建超过显存上限的数据。

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.59119    0.313164   0.76352036]
 [0.9731786  0.35454726 0.11677533]]
<NDArray 2x3 @gpu(1)>

除了在创建时指定，我们也可以通过copyto和as_in_context函数在设备之间传输数据。下面我们将内存上的 NDArray 变量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 计算，我们只需要事先将数据存储在显存上。计算结果会自动保存在同一块显卡的显存上。

In [11]:

(z + 2).exp() * y

Out[11]:

[ 20.085537 109.1963   445.2395  ]
<NDArray 3 @gpu(0)>

注意，MXNet 要求计算的所有输入数据都在内存或同一块显卡的显存上。这样设计的原因是 CPU 和不同的 GPU 之间的数据交互通常比较耗时。因此，MXNet 希望用户确切地指明计算的输入数据都在内存或同一块显卡的显存上。例如，如果将内存上的 NDArray 变量x和显存上的 NDArray 变量y做运算，会出现错误信息。当我们打印 NDArray 或将 NDArray 转换成 NumPy 格式时，如果数据不在内存，MXNet 会将它先复制到内存，从而造成额外的传输开销。

Gluon 的 GPU 计算¶

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

In [12]:

net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=mx.gpu())

当输入是显存上的 NDArray 时，Gluon 会在同一块显卡的显存上计算结果。

In [13]:

net(y)

Out[13]:

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

下面我们确认一下模型参数存储在同一块显卡的显存上。

In [14]:

net[0].weight.data()

Out[14]:

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

小结¶

MXNet 可以指定用来存储和计算的设备，如使用内存的 CPU 或者使用显存的 GPU。默认情况下，MXNet 会将数据创建在内存，然后利用 CPU 来计算。
MXNet 要求计算的所有输入数据都在内存或同一块显卡的显存上。

练习¶

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

扫码直达讨论区 ¶

参考文献¶

[1] CUDA 下载地址。 https://developer.nvidia.com/cuda-downloads