mxnet的设备管理

MXNet 使用 context 来指定用来存储和计算的设备，例如可以是 CPU 或者 GPU。默认情况下，MXNet 会将数据创建在主内存，然后利用 CPU 来计算。在 MXNet 中，CPU 和 GPU 可分别由 cpu() 和 gpu() 来表示。

需要注意的是，mx.cpu()（或者在括号里填任意整数）表示所有的物理 CPU 和内存。这意味着计算上会尽量使用所有的 CPU 核。
但 mx.gpu() 只代表一块显卡和相应的显卡内存。如果有多块 GPU，我们用 mx.gpu(i) 来表示第 i 块 GPU（i 从 0 开始）且 mx.gpu(0) 和 mx.gpu() 等价。

NDArray 的 GPU 计算

默认情况下，NDArray 存在 CPU 上

x = nd.array([1,2,3])
x
x.context
# output
[ 1.  2.  3.]
<NDArray 3 @cpu(0)>
cpu(0)

我们有多种方法将 NDArray 放置在 GPU 上。例如我们可以在创建 NDArray 的时候通过 ctx 指定存储设备。

a = nd.array([1, 2, 3], ctx=mx.gpu())
a
b = nd.random.uniform(shape=(2, 3), ctx=mx.gpu(1))
b
# output
[ 1.  2.  3.]
<NDArray 3 @gpu(0)>
[[ 0.59118998  0.313164    0.76352036][ 0.97317863  0.35454726  0.11677533]]
<NDArray 2x3 @gpu(1)>

除了在创建时指定，我们也可以通过 copyto 和 as_in_context 函数在设备之间传输数据。下面我们将 CPU 上的 x 复制到 GPU 0 上。

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

需要区分的是，如果源变量和目标变量的 context 一致，as_in_context 使目标变量和源变量共享源变量的内存;而 copyto 总是为目标变量新创建内存。

GPU 上的计算

MXNet 的计算会在数据的 context 上执行。为了使用 GPU 计算，我们只需要事先将数据放在 GPU 上面。而计算结果会自动保存在相同的 GPU 上。

注意，MXNet 要求计算的所有输入数据都在同一个 CPU/GPU 上。这个设计的原因是不同 CPU/GPU 之间的数据交互通常比较耗时。因此，MXNet 希望用户确切地指明计算的输入数据都在同一个 CPU/GPU 上。例如，如果将 CPU 上的 x 和 GPU 上的 y 做运算，会出现错误信息。

当我们打印 NDArray 或将 NDArray 转换成 NumPy 格式时，如果数据不在主内存里，MXNet 会自动将其先复制到主内存，从而带来隐形的传输开销。

Gluon 的 GPU 计算

同 NDArray 类似，Gluon 的模型可以在初始化时通过 ctx 指定设备。下面代码将模型参数初始化在 GPU 上。

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

当输入是 GPU 上的 NDArray 时，Gluon 会在相同的 GPU 上计算结果。

net(y)
# output
[[ 0.0068339 ][ 0.01366779][ 0.02050169]]
<NDArray 3x1 @gpu(0)>

模型参数存储在相同的 GPU 上。

net[0].weight.data()
[[ 0.0068339]]
<NDArray 1x1 @gpu(0)>