使用GPU来计算

【注意】运行本教程需要GPU。没有GPU的同学可以大致理解下内容,至少是context这个概念,因为之后我们也会用到。但没有GPU不会影响运行之后的大部分教程(好吧,还是有点点,可能运行会稍微慢点)。

前面的教程里我们一直在使用CPU来计算,因为绝大部分的计算设备都有CPU。但CPU的设计目的是处理通用的计算,例如打开浏览器和运行Jupyter,它一般只有少数的一块区域复杂数值计算,例如nd.dot(A, B)。对于复杂的神经网络和大规模的数据来说,单块CPU可能不够给力。

常用的解决办法是要么使用多台机器来协同计算,要么使用数值计算更加强劲的硬件,或者两者一起使用。本教程关注使用单块Nvidia GPU来加速计算,更多的选项例如多GPU和多机器计算则留到后面。

首先需要确保至少有一块Nvidia显卡已经安装好了,然后下载安装显卡驱动和CUDA(推荐下载8.0,CUDA自带了驱动)。完成后应该可以通过nvidia-smi查看显卡信息了。(Windows用户需要设一下PATH:set PATH=C:\Program Files\NVIDIA Corporation\NVSMI;%PATH%)。

In [1]:
!nvidia-smi
Thu Jan 25 18:25:12 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   43C    P0    39W / 150W |    692MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla M60           On   | 0000:00:1E.0     Off |                    0 |
| N/A   51C    P0    38W / 150W |    267MiB /  7612MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID  Type  Process name                               Usage      |
|=============================================================================|
|    0    109799    C   .../miniconda3/envs/gluon_zh_docs/bin/python   426MiB |
|    0    117705    C   .../miniconda3/envs/gluon_zh_docs/bin/python   264MiB |
|    1    117705    C   .../miniconda3/envs/gluon_zh_docs/bin/python   265MiB |
+-----------------------------------------------------------------------------+

接下来要要确认正确安装了的mxnet的GPU版本。具体来说是卸载了mxnetpip uninstall mxnet),然后根据CUDA版本安装mxnet-cu75或者mxnet-cu80(例如pip install --pre mxnet-cu80)。

使用pip来确认下:

In [2]:
import pip
for pkg in ['mxnet', 'mxnet-cu75', 'mxnet-cu80']:
    pip.main(['show', pkg])
Name: mxnet-cu80
Version: 1.0.1b20180125
Summary: MXNet is an ultra-scalable deep learning framework. This version uses CUDA-8.0.
Home-page: https://github.com/apache/incubator-mxnet
Author: UNKNOWN
Author-email: UNKNOWN
License: Apache 2.0
Location: /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages
Requires: numpy, requests, graphviz

Context

MXNet使用Context来指定使用哪个设备来存储和计算。默认会将数据开在主内存,然后利用CPU来计算,这个由mx.cpu()来表示。GPU则由mx.gpu()来表示。注意mx.cpu()表示所有的物理CPU和内存,意味着计算上会尽量使用多有的CPU核。但mx.gpu()只代表一块显卡和其对应的显卡内存。如果有多块GPU,我们用mx.gpu(i)来表示第i块GPU(i从0开始)。

In [3]:
import mxnet as mx
[mx.cpu(), mx.gpu(), mx.gpu(1)]
Out[3]:
[cpu(0), gpu(0), gpu(1)]

NDArray的GPU计算

每个NDArray都有一个context属性来表示它存在哪个设备上,默认会是cpu。这是为什么前面每次我们打印NDArray的时候都会看到@cpu(0)这个标识。

In [4]:
from mxnet import nd
x = nd.array([1,2,3])
x.context
Out[4]:
cpu(0)

GPU上创建内存

我们可以在创建的时候指定创建在哪个设备上(如果GPU不能用或者没有装MXNet GPU版本,这里会有error):

In [5]:
a = nd.array([1,2,3], ctx=mx.gpu())
b = nd.zeros((3,2), ctx=mx.gpu())
c = nd.random.uniform(shape=(2,3), ctx=mx.gpu())
(a,b,c)
Out[5]:
(
 [ 1.  2.  3.]
 <NDArray 3 @gpu(0)>,
 [[ 0.  0.]
  [ 0.  0.]
  [ 0.  0.]]
 <NDArray 3x2 @gpu(0)>,
 [[ 0.66865093  0.17409194  0.38500249]
  [ 0.24678314  0.35134333  0.84042978]]
 <NDArray 2x3 @gpu(0)>)

尝试将内存开到另外一块GPU上。如果不存在会报错。当然,如果你有大于10块GPU,那么下面代码会顺利执行。

In [6]:
import sys

try:
    nd.array([1,2,3], ctx=mx.gpu(10))
except mx.MXNetError as err:
    sys.stderr.write(str(err))
[18:25:15] src/storage/storage.cc:63: Check failed: e == cudaSuccess || e == cudaErrorCudartUnloading CUDA: invalid device ordinal

Stack trace returned 10 entries:
[bt] (0) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2a6788) [0x7f1d8432f788]
[bt] (1) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2a6b98) [0x7f1d8432fb98]
[bt] (2) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x28cc7d6) [0x7f1d869557d6]
[bt] (3) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x28d14eb) [0x7f1d8695a4eb]
[bt] (4) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x23dd031) [0x7f1d86466031]
[bt] (5) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(MXNDArrayCreateEx+0x145) [0x7f1d86466805]
[bt] (6) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/../../libffi.so.6(ffi_call_unix64+0x4c) [0x7f1dc52eee6c]
[bt] (7) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/../../libffi.so.6(ffi_call+0x165) [0x7f1dc52edfd5]
[bt] (8) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(_ctypes_callproc+0x2ce) [0x7f1dc5502dee]
[bt] (9) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(+0x12825) [0x7f1dc5503825]

我们可以通过copytoas_in_context来在设备直接传输数据。

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

这两个函数的主要区别是,如果源和目标的context一致,as_in_context不复制,而copyto总是会新建内存:

In [8]:
yy = y.as_in_context(mx.gpu())
zz = z.copyto(mx.gpu())
(yy is y, zz is z)
Out[8]:
(True, False)

GPU上的计算

计算会在数据的context上执行。所以为了使用GPU,我们只需要事先将数据放在上面就行了。结果会自动保存在对应的设备上:

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

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

注意所有计算要求输入数据在同一个设备上。不一致的时候系统不进行自动复制。这个设计的目的是因为设备之间的数据交互通常比较昂贵,我们希望用户确切的知道数据放在哪里,而不是隐藏这个细节。下面代码尝试将CPU上x和GPU上的y做运算。

In [10]:
try:
    x + y
except mx.MXNetError as err:
    sys.stderr.write(str(err))
[18:25:15] src/imperative/./imperative_utils.h:55: Check failed: inputs[i]->ctx().dev_mask() == ctx.dev_mask() (2 vs. 1) Operator broadcast_add require all inputs live on the same context. But the first argument is on cpu(0) while the 2-th argument is on gpu(0)

Stack trace returned 10 entries:
[bt] (0) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2a6788) [0x7f1d8432f788]
[bt] (1) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x2a6b98) [0x7f1d8432fb98]
[bt] (2) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x24af99c) [0x7f1d8653899c]
[bt] (3) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x24be50e) [0x7f1d8654750e]
[bt] (4) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(+0x240038b) [0x7f1d8648938b]
[bt] (5) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/site-packages/mxnet/libmxnet.so(MXImperativeInvokeEx+0x63) [0x7f1d864898f3]
[bt] (6) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/../../libffi.so.6(ffi_call_unix64+0x4c) [0x7f1dc52eee6c]
[bt] (7) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/../../libffi.so.6(ffi_call+0x165) [0x7f1dc52edfd5]
[bt] (8) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(_ctypes_callproc+0x2ce) [0x7f1dc5502dee]
[bt] (9) /var/lib/jenkins/miniconda3/envs/gluon_zh_docs/lib/python3.6/lib-dynload/_ctypes.cpython-36m-x86_64-linux-gnu.so(+0x12825) [0x7f1dc5503825]

默认会复制回CPU的操作

如果某个操作需要将NDArray里面的内容转出来,例如打印或变成numpy格式,如果需要的话系统都会自动将数据copy到主内存。

In [11]:
print(y)
print(y.asnumpy())
print(y.sum().asscalar())

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

Gluon的GPU计算

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

In [12]:
from mxnet import gluon
net = gluon.nn.Sequential()
net.add(gluon.nn.Dense(1))

net.initialize(ctx=mx.gpu())

输入GPU上的数据,会在GPU上计算结果

In [13]:
data = nd.random.uniform(shape=[3,2], ctx=mx.gpu())
net(data)
Out[13]:

[[ 0.00271713]
 [ 0.01119637]
 [ 0.0097794 ]]
<NDArray 3x1 @gpu(0)>

确认下权重:

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

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

总结

通过context我们可以很容易在不同的设备上计算。

练习

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

吐槽和讨论欢迎点这里