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)>