Apache MXNet - 分散式訓練

本章介紹 Apache MXNet 中的分散式訓練。讓我們首先了解 MXNet 中的計算模式。

計算模式

MXNet 是一款多語言機器學習庫，為使用者提供了以下兩種計算模式：

命令式模式

這種計算模式提供類似 NumPy API 的介面。例如，在 MXNet 中，使用以下命令式程式碼在 CPU 和 GPU 上構建一個全零張量：

import mxnet as mx
tensor_cpu = mx.nd.zeros((100,), ctx=mx.cpu())
tensor_gpu= mx.nd.zeros((100,), ctx=mx.gpu(0))

如上程式碼所示，MXNet 指定了張量儲存的位置，可以在 CPU 或 GPU 裝置上。在上面的例子中，位置為 0。MXNet 達到了驚人的裝置利用率，因為所有計算都是延遲執行的，而不是立即執行。

符號式模式

雖然命令式模式非常有用，但這種模式的一個缺點是其僵化性，即所有計算都需要預先知道，並且需要預定義的資料結構。

另一方面，符號式模式提供類似 TensorFlow 的計算圖。它透過允許 MXNet 使用符號或變數而不是固定/預定義的資料結構來克服命令式 API 的缺點。之後，這些符號可以解釋為一組操作，如下所示：

import mxnet as mx
x = mx.sym.Variable(“X”)
y = mx.sym.Variable(“Y”)
z = (x+y)
m = z/100

並行型別的

Apache MXNet 支援分散式訓練。它使我們能夠利用多臺機器進行更快、更有效的訓練。

以下是將神經網路訓練的工作負載分佈到多個裝置（CPU 或 GPU 裝置）的兩種方法：

資料並行

在這種並行性中，每個裝置儲存模型的完整副本，並使用資料集的不同部分。裝置還會集體更新共享模型。我們可以將所有裝置放在一臺機器上，也可以跨多臺機器。

模型並行

這是另一種並行性，當模型太大而無法放入裝置記憶體時非常有用。在模型並行中，不同的裝置被分配學習模型的不同部分的任務。這裡需要注意的重要一點是，目前 Apache MXNet 只支援單機模型並行。

分散式訓練的工作原理

以下概念是理解 Apache MXNet 中分散式訓練工作原理的關鍵：

程序型別

程序之間相互通訊以完成模型的訓練。Apache MXNet 有以下三個程序：

工作節點

工作節點的任務是在一批訓練樣本上執行訓練。工作節點在處理每一批之前會從伺服器拉取權重。工作節點在處理完一批後會將梯度傳送到伺服器。

伺服器

MXNet 可以有多個伺服器來儲存模型的引數並與工作節點通訊。

排程器

排程器的作用是設定叢集，包括等待每個節點啟動的訊息以及節點正在監聽的埠。設定集群后，排程器讓所有程序都知道叢集中的其他每個節點。這是因為程序可以相互通訊。只有一個排程器。

KVStore

KVStore 代表**鍵值**儲存。它是用於多裝置訓練的關鍵元件。它很重要，因為在單機或跨多臺機器的裝置之間進行引數通訊是透過一個或多個具有引數 KVStore 的伺服器傳輸的。讓我們透過以下幾點來了解 KVStore 的工作原理：

KVStore 中的每個值都由一個**鍵**和一個**值**表示。
網路中的每個引數陣列都分配一個**鍵**，並且該引數陣列的權重由**值**引用。
之後，工作節點在處理完一批後**推送**梯度。它們還在處理新批次之前**拉取**更新的權重。

KVStore 伺服器的概念僅在分散式訓練期間存在，並且其分散式模式是透過使用包含單詞**dist** 的字串引數呼叫**mxnet.kvstore.create** 函式來啟用的：

kv = mxnet.kvstore.create(‘dist_sync’)

金鑰的分佈

並非所有伺服器都儲存所有引數陣列或鍵，而是將它們分佈在不同的伺服器上。KVStore 透明地處理這種跨不同伺服器的鍵的分佈，並且隨機決定哪個伺服器儲存特定鍵。

如上所述，KVStore 確保每當拉取鍵時，其請求都會發送到擁有相應值的伺服器。如果某個鍵的值很大怎麼辦？在這種情況下，它可能被共享到不同的伺服器上。

分割訓練資料

作為使用者，我們希望每臺機器都能處理資料集的不同部分，尤其是在資料並行模式下執行分散式訓練時。我們知道，為了分割資料迭代器提供的一批樣本以便在單個工作節點上進行資料並行訓練，我們可以使用**mxnet.gluon.utils.split_and_load**，然後將批次的每一部分載入到將進一步處理它的裝置上。

另一方面，對於分散式訓練，首先我們需要將資料集分成**n**個不同的部分，以便每個工作節點獲得不同的部分。一旦獲得，每個工作節點可以使用**split_and_load**再次將資料集的那一部分劃分到單機上的不同裝置上。所有這些都是透過資料迭代器完成的。**mxnet.io.MNISTIterator** 和**mxnet.io.ImageRecordIter** 是 MXNet 中支援此功能的兩個這樣的迭代器。

權重更新

為了更新權重，KVStore 支援以下兩種模式：

第一種方法聚合梯度並使用這些梯度更新權重。
在第二種方法中，伺服器只聚合梯度。

如果使用 Gluon，可以透過傳遞**update_on_kvstore**變數來在這兩種方法之間進行選擇。讓我們透過建立**trainer**物件來了解它：

trainer = gluon.Trainer(net.collect_params(), optimizer='sgd',
   optimizer_params={'learning_rate': opt.lr,
      'wd': opt.wd,
      'momentum': opt.momentum,
      'multi_precision': True},
      kvstore=kv,
   update_on_kvstore=True)

分散式訓練模式

如果 KVStore 建立字串包含單詞 dist，則表示已啟用分散式訓練。以下是透過使用不同型別的 KVStore 可以啟用的不同分散式訓練模式：

dist_sync

顧名思義，它表示同步分散式訓練。在此模式下，所有工作節點在每批的開始都使用相同的同步模型引數集。

此模式的缺點是，在每批之後，伺服器必須等待接收來自每個工作節點的梯度，然後才能更新模型引數。這意味著如果一個工作節點崩潰，它將阻止所有工作節點的進度。

dist_async

顧名思義，它表示非同步分散式訓練。在此模式下，伺服器接收來自一個工作節點的梯度並立即更新其儲存。伺服器使用更新後的儲存來響應任何進一步的拉取請求。

與**dist_sync 模式**相比，它的優勢在於，完成處理一批的工作節點可以從伺服器拉取當前引數並開始下一批。即使其他工作節點尚未完成處理之前的批次，工作節點也可以這樣做。它也比 dist_sync 模式快，因為它可以在沒有任何同步成本的情況下花費更多輪次來收斂。

dist_sync_device

此模式與**dist_sync**模式相同。唯一的區別是，當每個節點上使用多個 GPU 時，**dist_sync_device** 在 GPU 上聚合梯度並更新權重，而**dist_sync** 在 CPU 記憶體上聚合梯度並更新權重。

它減少了 GPU 和 CPU 之間的昂貴通訊。這就是為什麼它比**dist_sync**更快的原因。缺點是它會增加 GPU 的記憶體使用量。

dist_async_device

此模式的工作方式與**dist_sync_device**模式相同，但在非同步模式下。

列印頁面