MxNet學習： 深度學習的依賴引擎
我們總是希望深度學習庫執行得更快，能夠擴充套件到更大的資料集。一種自然的方法是，看看我們是否可以從向這個問題投入更多硬體中獲益，就像同時使用多個gpu。
然後，庫設計人員會問:如何跨裝置平行計算?更重要的是，我們如何在引入多執行緒時同步計算?執行時依賴引擎是這些問題的通用解決方案。
在本文中，我們研究了使用執行時依賴排程來加速深度學習的方法。我們的目的是解釋執行時依賴排程如何既能加速又能簡化多裝置深度學習。我們還探討了一種通用依賴引擎的潛在設計，這種引擎既可以是獨立於庫的，也可以是獨立於操作的。
本文的大部分討論都是從MXNet依賴引擎中獲得靈感的。我們討論的依賴跟蹤演算法最初是由Yutian Li和Mingjie Wang開發的。
1 依賴排程
雖然大多數使用者希望利用平行計算，但我們大多數人更熟悉序列程式。因此，一個自然的問題是:我們如何編寫序列程式並構建一個庫以非同步方式自動並行我們的程式?
例如，在下面的程式碼中，我們可以按任意順序執行B = A + 1和C = A + 2，或者並行執行:
A = 2
B = A + 1
C = A + 2
D = B * C
但是，手工編寫序列程式碼是相當困難的，因為最後一個操作D = B * C需要等待前面的兩個操作都完成之後才能開始。下面的依賴關係圖/資料流圖說明了這一點。
依賴引擎是一個庫，它接受一系列操作，並根據依賴模式(可能是並行的)對它們進行排程。因此在本例中，依賴項庫可以並行執行B = a + 1和C = a + 2，並在這些操作完成後執行D = B * C。
2 依賴排程中的問題
依賴引擎減輕了編寫併發程式的負擔。然而，隨著操作的並行化，新的依賴跟蹤問題出現了。在本節中，我們將討論這些問題。
2.1 資料流依賴
資料流依賴關係描述如何在其他計算中使用一個計算的結果。每個依賴引擎都必須解決資料流依賴問題。
因為我們在前一節中討論了這個問題，所以這裡包含了相同的圖。擁有資料流跟蹤引擎的庫包括Minerva和Purine2。
2.2 記憶體複用
什麼時候應該回收分配給陣列的記憶體?在序列處理中，這很容易確定。我們只是在變數超出作用域後回收記憶體。然而，如下圖所示，這在並行處理中有點困難。
在這個例子中，因為這兩個計算都需要使用來自A的值，所以在完成之前我們不能回收記憶體。引擎必須根據依賴關係調度記憶體回收操作，並確保在B = A + 1和C = A + 2完成之後執行這些操作。
2.3 隨機數生成
隨機數生成器是機器學習中常用的一種生成器，它給依賴引擎帶來了有趣的挑戰。考慮下面的例子:
在這個例子中，我們按順序生成隨機數。雖然這兩個隨機數生成似乎可以並行化，但通常不是這樣。偽隨機數生成器(PRNG)不是執行緒安全的，因為在生成新數字時，它可能導致某些內部狀態發生突變。即使PRNG是執行緒安全的，序列化數字生成也是可取的，這樣我們就可以得到可重複的隨機數。
2.4 案例研究:一個多gpu神經網路的依賴引擎
在上一節中，我們討論了在設計依賴引擎時可能遇到的問題。在考慮如何設計一個通用引擎來解決這些問題之前，讓我們先考慮一個依賴引擎如何在神經網路的多gpu訓練中提供幫助。下面的虛擬碼Python程式演示了在雙層神經網路上的一個批次的訓練過程。
# Example of one iteration Two GPU neural Net
data = next_batch()
data[gpu0].copyfrom(data[0:50])
data[gpu1].copyfrom(data[50:100])
# forward, backprop on GPU 0
fc1[gpu0] = FullcForward(data[gpu0], fc1_weight[gpu0])
fc2[gpu0] = FullcForward(fc1[gpu0], fc2_weight[gpu0])
fc2_ograd[gpu0] = LossGrad(fc2[gpu0], label[0:50])
fc1_ograd[gpu0], fc2_wgrad[gpu0] =
FullcBackward(fc2_ograd[gpu0] , fc2_weight[gpu0])
_, fc1_wgrad[gpu0] = FullcBackward(fc1_ograd[gpu0] , fc1_weight[gpu0])
# forward, backprop on GPU 1
fc1[gpu1] = FullcForward(data[gpu1], fc1_weight[gpu1])
fc2[gpu1] = FullcForward(fc1[gpu1], fc2_weight[gpu1])
fc2_ograd[gpu1] = LossGrad(fc2[gpu1], label[50:100])
fc1_ograd[gpu1], fc2_wgrad[gpu1] =
FullcBackward(fc2_ograd[gpu1] , fc2_weight[gpu1])
_, fc1_wgrad[gpu1] = FullcBackward(fc1_ograd[gpu1] , fc1_weight[gpu1])
# aggregate gradient and update
fc1_wgrad[cpu] = fc1_wgrad[gpu0] + fc1_wgrad[gpu1]
fc2_wgrad[cpu] = fc2_wgrad[gpu0] + fc2_wgrad[gpu1]
fc1_weight[cpu] -= lr * fc1_wgrad[cpu]
fc2_weight[cpu] -= lr * fc2_wgrad[cpu]
fc1_weight[cpu].copyto(fc1_weight[gpu0] , fc1_weight[gpu1])
fc2_weight[cpu].copyto(fc2_weight[gpu0] , fc2_weight[gpu1])
在這個程式中，資料0到50複製到GPU 0，資料50到100複製到GPU 1。計算出的梯度在CPU中聚合，然後CPU執行一個簡單的SGD更新，並將更新後的權重複制回每個GPU。這是一種以序列方式編寫並行程式的常見方法。下圖顯示瞭如何並行化它:
說明：
- 一旦我們得到一個層的梯度，可以立即將其複製到CPU；
- 權重一經更新即可複製回GPU；
- 在前向計算中，有一個對上一輪迭代的依賴項：fc1_weight[cpu].copyto(fc1_weight[gpu0] , fc1_weight[gpu1])；
- 從上一次向後傳遞到第k層到下一次向前呼叫第k層之間有一個計算延遲，在此延遲期間，我們可以將k層的權值同步與其他計算並行進行同步。
這種優化方法被CXXNet等多gpu深度學習庫所使用。重點是使權重同步(通訊)與計算重疊。但是，這並不容易，因為複製操作需要在層的後向傳遞完成後立即觸發，然後再觸發reduce、update等。
依賴引擎可以排程這些操作並執行多執行緒和依賴跟蹤。
3 設計一個通用的依賴引擎
我們希望您確信依賴引擎對於將深度學習程式擴充套件到多個裝置是有用的。現在讓我們討論如何為依賴引擎設計和實現通用介面。這種解決方案並不是依賴引擎唯一可能的設計。這是一個我們認為在大多數情況下有用的例子。
我們的目標是建立一個通用的輕量級依賴引擎。理想情況下，我們希望引擎能夠方便地插入到現有的深度學習程式碼中，並且能夠在稍加修改的情況下擴充套件到多臺機器。要做到這一點，我們只需要關注依賴跟蹤，而不是假設使用者可以或不能做什麼。
這裡是引擎目標的總結：
- 引擎不應該知道它執行什麼操作，以便使用者可以執行他們定義的任何操作；
- 不應該限制引擎可以排程的物件型別；
我們應該能夠排程對GPU和CPU記憶體的依賴 ；
我們應該能夠跟蹤隨機數生成器的依賴關係，等等； - 引擎不應該分配資源。它應該只跟蹤依賴項。使用者可以分配自己的記憶體、PRNG等。
下面的Python程式碼片段提供了一個引擎介面，可以幫助我們實現目標。注意，真正的實現將更接近於metal，通常是在c++中。
class DepEngine(object):
def new_variable():
"""Return a new variable tag
Returns
-------
vtag : Variable Tag
The token of the engine to represent dependencies.
"""
pass
def push(exec_func, read_vars, mutate_vars):
"""Push the operation to the engine.
Parameters
----------
exec_func : callable
The real operation to be performed.
read_vars : list of Variable Tags
The list of variables this operation will read from.
mutate_vars : list of Variable Tags
The list of variables this operation will mutate.
"""
pass
因為我們不能假設我們正在排程什麼物件，所以我們要求使用者分配一個與每個物件相關聯的虛擬標記來表示我們需要排程什麼。因此，在一開始，使用者可以分配變數標記，並將其附加到我們希望排程的每個物件上。
然後使用者呼叫push來告訴引擎要執行的函式。使用者還需要指定操作的依賴關係，使用read_vars和write_vars:
- read_vars是隻讀變數的標記，不需要更改其內部狀態；
- mutate_vars是內部狀態可改變的物件的變數標記。
上面的圖顯示瞭如何將操作B = A + 1推送到依賴引擎。B.data和A.data是分配好的空間。請注意，引擎只知道變數標記。任何執行函式都可以被處理。這個介面對於我們想要排程的操作和資源是通用的。
為了好玩，讓我們通過以下程式碼片段來了解引擎內部是如何處理標記的：
B = A + 1
C = A + 2
A = C * 2
D = A + 3
第一行讀取變數A和可變變數B。第二行讀取變數A和可變變數C，以此類推。
引擎為每個變數維護一個佇列，如下面的動畫所示。綠色塊表示讀取操作，而紅色塊表示修改操作。
在構建此佇列時，引擎會看到A佇列開頭的前兩個綠色塊實際上可以並行執行，因為它們都是讀取操作，不會相互衝突。下圖說明了這一點。
所有這些排程的一個很酷的地方是，它不只侷限於數值計算。因為計劃的所有事情都只是一個標記，引擎可以計劃所有事情!
下圖給出了我們在前面幾節中提到的程式的完整推送序列。
3.1 將現有程式碼移植到依賴引擎
因為通用介面不控制記憶體分配和要執行的操作，所以依賴引擎可以將大多數現有程式碼排程為兩個步驟:
- 分配與記憶體blob、PRNG等資源相關的變數標記；
- 呼叫push，將執行函式作為要執行的原始程式碼，並將相應資源的變數標記正確地放入read_vars和mutate_vars中。
3.2 實現通用依賴引擎
我們已經描述了通用引擎介面，以及如何使用它來排程各種操作。在本節中，我們將對如何實現這樣的引擎進行高階討論。
總體思路如下:
- 使用佇列跟蹤每個變數標記上所有掛起的依賴項；
- 在每個操作上使用計數器來跟蹤還需要完成多少依賴項；
- 操作完成後，更新佇列和依賴項計數器的狀態，以排程新的操作。
下圖演示了排程演算法，可能會讓您更好地瞭解引擎中正在發生的事情。
下面，我們將展示另一個涉及隨機數生成器的示例。
如您所見，該演算法的目的是更新操作的掛起佇列，並在操作完成時進行正確的狀態轉換。應該更加小心地確保狀態轉換以執行緒安全的方式完成。
3.3 獨立的依賴跟蹤和執行策略
如果您正在仔細閱讀，您可能已經注意到前面的部分只顯示了決定何時可以執行操作的演算法。我們沒有展示如何實際執行一個操作。實際上，可以有許多不同的策略。例如，我們可以使用一個全域性執行緒池來執行所有操作，或者使用一個特定的執行緒在每個裝置上執行操作。
這種執行策略通常獨立於依賴跟蹤，可以分為獨立模組或基本依賴跟蹤模組的虛擬介面。開發一個對所有操作和排程都公平的優雅的執行時策略本身就是一個有趣的系統問題。
4 討論
我們在本文中討論的設計並不是依賴跟蹤問題的唯一解決方案。這只是我們如何解決這個問題的一個例子。當然，其中一些設計選擇是有爭議的。我們將在本節中討論其中一些。
4.1 動態 VS. 靜態
本主題中討論的依賴項引擎介面在某種程度上是動態的，因為使用者可以逐個推進操作，而不是宣告整個依賴項圖(靜態)。就資料結構而言，動態排程可能比靜態宣告需要更多的開銷。但是，它也支援更多的靈活性，例如支援命令式程式的自動並行性，或者命令式程式和符號程式的混合。您還可以向介面新增一些預先宣告的操作，以支援資料結構重用。
4.2 可變的 VS. 不可變
本文提供的通用引擎介面支援顯式的可變排程。在典型的資料流引擎中，資料通常是不可變的。使用不可變資料有很多好處。例如，不可變資料通常更適合並行化，並且有助於在分散式設定中更好地容錯(通過重新計算)。
然而，不可變性帶來了幾個挑戰：
- 在處理隨機數和刪除時，很難安排資源競爭問題；
- 引擎通常需要管理資源(記憶體、隨機數)以避免衝突。很難處理使用者分配的空間，等等；
- 預先分配的靜態記憶體不可用，這也是因為通常的模式是寫入預先分配的層空間，如果資料是不可變的，則不支援這種模式。
允許可變排程可以緩解這些問題。
4.3 MxNet
MxNet是按照本文的思想實現了通用的依賴引擎。
5 參考
文章翻譯自：https://mxnet.incubator.apache.org/architecture/note_engine.html