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

轉載于:https://www.cnblogs.com/houkai/p/9522182.html

總結

以上是生活随笔為你收集整理的MXNET：深度学习计算-GPU的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。