GPU 计算

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

In [1]:
!nvidia-smi
Mon Sep 17 14:11:41 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   58C    P0    48W / 150W |      0MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           On   | 0000:00:1E.0     Off |                    0 |
| N/A   54C    P8    15W / 150W |      0MiB /  7612MiB |      0%      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 中,CPU 和 GPU 可分别由cpu()gpu()来表示。需要注意的是,mx.cpu()(或者在括号里填任意整数)表示所有的物理 CPU 和内存。这意味着 MXNet 的计算会尽量使用所有的 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属性来查看该 NDArray 所在的设备。

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

GPU 上的存储

我们有多种方法将 NDArray 存储在 GPU 上。例如我们可以在创建 NDArray 的时候通过ctx参数指定存储设备。下面我们将 NDArray 变量a创建在gpu(0)上。注意到在打印a时,设备信息变成了@gpu(0)。创建在 GPU 上的 NDArray 只消耗相同 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.59119    0.313164   0.76352036]
 [0.9731786  0.35454726 0.11677533]]
<NDArray 2x3 @gpu(1)>

除了在创建时指定,我们也可以通过copytoas_in_context函数在设备之间传输数据。下面我们将 CPU 上的 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 计算,我们只需要事先将数据存储在 GPU 上。计算结果会自动保存在相同的 GPU 上。

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

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

注意,MXNet 要求计算的所有输入数据都在 CPU 或同一个 GPU 上。这个设计的原因是 CPU 和不同的 GPU 之间的数据交互通常比较耗时。因此,MXNet 希望用户确切地指明计算的输入数据都在 CPU 或同一个 GPU 上。例如,如果将 CPU 上的 NDArray 变量x和 GPU 上的 NDArray 变量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)>

小结

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

练习

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

扫码直达讨论区

参考文献

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