MXNET：深度學習計算-GPU

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