mxnet的设备管理
MXNet 使用 context 来指定用来存储和计算的设备,例如能够是 CPU 或者 GPU。默认状况下,MXNet 会将数据建立在主内存,而后利用 CPU 来计算。在 MXNet 中,CPU 和 GPU 可分别由 cpu() 和 gpu() 来表示。dom
须要注意的是,mx.cpu()(或者在括号里填任意整数)表示全部的物理 CPU 和内存。这意味着计算上会尽可能使用全部的 CPU 核。 但 mx.gpu() 只表明一块显卡和相应的显卡内存。若是有多块 GPU,咱们用 mx.gpu(i) 来表示第 i 块 GPU(i 从 0 开始)且 mx.gpu(0) 和 mx.gpu() 等价。函数
NDArray 的 GPU 计算
默认状况下,NDArray 存在 CPU 上spa
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 上。code
y = x.copyto(mx.gpu()) z = x.as_in_context(mx.gpu())
须要区分的是,若是源变量和目标变量的 context 一致,as_in_context 使目标变量和源变量共享源变量的内存;而 copyto 老是为目标变量新建立内存。orm
GPU 上的计算内存
MXNet 的计算会在数据的 context 上执行。为了使用 GPU 计算,咱们只须要事先将数据放在 GPU 上面。而计算结果会自动保存在相同的 GPU 上。it
注意,**MXNet 要求计算的全部输入数据都在同一个 CPU/GPU 上。**这个设计的缘由是不一样 CPU/GPU 之间的数据交互一般比较耗时。所以,MXNet 但愿用户确切地指明计算的输入数据都在同一个 CPU/GPU 上。例如,若是将 CPU 上的 x 和 GPU 上的 y 作运算,会出现错误信息。form
当咱们打印 NDArray 或将 NDArray 转换成 NumPy 格式时,若是数据不在主内存里,MXNet 会自动将其先复制到主内存,从而带来隐形的传输开销。class
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)>