【MXNet】（十八）：使用GPU计算

MXNet可以指定用来计算和存储的设备。默认情况下，MXNet会将数据创建在内存，然后利用CPU来计算。在MXNet中，mx.cpu()（或者在括号里填任意整数）表示所有的物理CPU和内存。这意味着，MXNet的计算会尽量使用所有的CPU核。但mx.gpu()只代表一块GPU和相应的显存。如果有多块GPU，我们用mx.gpu(i)来表示第ii块GPU及相应的显存（ii从0开始）且mx.gpu(0)和mx.gpu()等价。

import mxnet as mx
from mxnet import nd
from mxnet.gluon import nn

mx.cpu(), mx.gpu(), mx.gpu(0)

在默认情况下，NDArray存在内存上。

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

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

x.context

可以在创建NDArray的时候通过ctx参数指定存储设备。

a = nd.array([1, 2, 3], ctx=mx.gpu())
a

B = nd.random.uniform(shape=(2, 3), ctx=mx.gpu(0))
B

也可以通过copyto函数和as_in_context函数在设备之间传输数据。

y = x.copyto(mx.gpu())
y

z = x.as_in_context(mx.gpu())
z

如果源变量和目标变量的context一致，as_in_context函数使目标变量和源变量共享源变量的内存或显存。

y.as_in_context(mx.gpu()) is y

而copyto函数总是为目标变量开新的内存或显存。

y.copyto(mx.gpu()) is y

MXNet的计算会在数据的context属性所指定的设备上执行。为了使用GPU计算，我们只需要事先将数据存储在显存上。计算结果会自动保存在同一块显卡的显存上。MXNet要求计算的所有输入数据都在内存或同一块显卡的显存上。

(z + 2).exp() * y

Gluon的模型可以在初始化时通过ctx参数指定设备。

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

看一下模型参数是不是存储在同一块显卡的显存上。

net[0].weight.data()