GPUs

查看 CUDA 和驱动。

In [1]:
!nvidia-smi
Wed Sep  4 11:20:44 2019       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 418.67       Driver Version: 418.67       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  Tesla V100-SXM2...  Off  | 00000000:00:1B.0 Off |                    0 |
| N/A   58C    P0    58W / 300W |      0MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   1  Tesla V100-SXM2...  Off  | 00000000:00:1C.0 Off |                    0 |
| N/A   52C    P0    58W / 300W |    322MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   2  Tesla V100-SXM2...  Off  | 00000000:00:1D.0 Off |                    0 |
| N/A   52C    P0    61W / 300W |      0MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
|   3  Tesla V100-SXM2...  Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   51C    P0    43W / 300W |     11MiB / 16130MiB |      0%      Default |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    1     66377      C   /home/ubuntu/miniconda3/bin/python           311MiB |
+-----------------------------------------------------------------------------+

查看有多少 GPU 可以用。

In [2]:
from mxnet import np, npx
from mxnet.gluon import nn
npx.set_np()

npx.num_gpus()
Out[2]:
2

使用不同的设备。

In [3]:
print(npx.cpu(), npx.gpu(), npx.gpu(1))

def try_gpu(i=0):
    return npx.gpu(i) if npx.num_gpus() >= i + 1 else npx.cpu()

def try_all_gpus():
    ctxes = [npx.gpu(i) for i in range(npx.num_gpus())]
    return ctxes if ctxes else [npx.cpu()]

try_gpu(), try_gpu(3), try_all_gpus()
cpu(0) gpu(0) gpu(1)
Out[3]:
(gpu(0), cpu(0), [gpu(0), gpu(1)])

在第一个 GPU 上创建数据。

In [4]:
x = np.ones((2, 3), ctx=try_gpu())
print(x.context)
x
gpu(0)
Out[4]:
array([[1., 1., 1.],
       [1., 1., 1.]], ctx=gpu(0))

在第二个 GPU 上创建数据。

In [5]:
y = np.random.uniform(size=(2, 3), ctx=try_gpu(1))
y
Out[5]:
array([[0.67478997, 0.07540122, 0.9956977 ],
       [0.09488854, 0.415456  , 0.11231736]], ctx=gpu(1))

在不同的设备间复制数据。

In [6]:
z = x.copyto(try_gpu(1))
print(x)
print(z)
[[1. 1. 1.]
 [1. 1. 1.]] @gpu(0)
[[1. 1. 1.]
 [1. 1. 1.]] @gpu(1)

一个操作子的输入必须在同一个设备上,计算也将会在那个设备上。

In [7]:
y + z
Out[7]:
array([[1.6747899, 1.0754012, 1.9956977],
       [1.0948886, 1.415456 , 1.1123173]], ctx=gpu(1))

在第一个 GPU 上初始化模型参数。

In [8]:
net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=try_gpu())

当输入数据也在 GPU 上时,Gluon 会在那个 GPU 上计算结果。

In [9]:
net(x)
Out[9]:
array([[0.04995865],
       [0.04995865]], ctx=gpu(0))

确认模型参数和结果都在同一个 GPU 上。

In [10]:
net[0].weight.data()
Out[10]:
array([[0.0068339 , 0.01299825, 0.0301265 ]], ctx=gpu(0))