查看 CUDA 和驱动。
!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 可以用。
from mxnet import np, npx
from mxnet.gluon import nn
npx.set_np()
npx.num_gpus()
2
使用不同的设备。
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)
(gpu(0), cpu(0), [gpu(0), gpu(1)])
在第一个 GPU 上创建数据。
x = np.ones((2, 3), ctx=try_gpu())
print(x.context)
x
gpu(0)
array([[1., 1., 1.], [1., 1., 1.]], ctx=gpu(0))
在第二个 GPU 上创建数据。
y = np.random.uniform(size=(2, 3), ctx=try_gpu(1))
y
array([[0.67478997, 0.07540122, 0.9956977 ], [0.09488854, 0.415456 , 0.11231736]], ctx=gpu(1))
在不同的设备间复制数据。
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)
一个操作子的输入必须在同一个设备上,计算也将会在那个设备上。
y + z
array([[1.6747899, 1.0754012, 1.9956977], [1.0948886, 1.415456 , 1.1123173]], ctx=gpu(1))
在第一个 GPU 上初始化模型参数。
net = nn.Sequential()
net.add(nn.Dense(1))
net.initialize(ctx=try_gpu())
当输入数据也在 GPU 上时,Gluon 会在那个 GPU 上计算结果。
net(x)
array([[0.04995865], [0.04995865]], ctx=gpu(0))
确认模型参数和结果都在同一个 GPU 上。
net[0].weight.data()
array([[0.0068339 , 0.01299825, 0.0301265 ]], ctx=gpu(0))