MXNet入坑（二）- Engine 「依賴引擎」與它的核心演算法

@Mingxuan Liu（知乎用戶）劉大佬參考了MXNet的官方文檔，寫的深入淺出。經同意，原文直接搬過來了。

Dependency Engine for Deep Learningmxnet.incubator.apache.org

前言：我們總是希望深度學習的庫能夠更快的運行並能處理大的數據集。自然而然的我們想到了並行計算。首先我們碰到的問題是：在並行計算中，我們怎麼同步數據？設計一個運行時的「依賴引擎」是解決這個問題的常用方法。下面我們就 MXNet 設計「依賴引擎」的目的丶碰到的問題丶解決方法（演算法思想）進行具體的描述。

設計「依賴引擎」的目的

設計一個操作（operation）獨立的庫，以非同步並行的方式優化用戶寫的串列程序。

難點

1.數據流依賴

數據流依賴描述了一個計算結果怎麼影響到其它的計算。這是每一個「依賴引擎「必須解決的問題。下面，我們就一個簡單的例子來具體描述它。

我們來看這樣一個串列程序：

① A = 2

② B = A + 1

③ C = A + 2

④ D = B * C

在這個串列程序中，我們可以直觀的看到：

操作②和操作③都直接地依賴於操作①，操作④直接地依賴於操作②和③

換而言之，操作①的計算結果直接地影響了操作②和③，②和③的計算結果直接地影響了操作④。通過分析，我們很容易知道，我們可以以下面的方式優化這個串列程序：

• 第一步，執行操作①
• 第二步，並行的執行操作②和③
• 第三步，執行操作④

這就是所謂的依賴調度，在演算法部分我們會更詳細的描述它。

2.內存的回收

在串列程序中，我們可以很容易的知道，當變數越過它的作用域時，我們就可以回收它佔用的內存。但是在並行計算時就比較難以判斷了。MXNet 中使用回收內存操作和其它操作之間的依賴關係來確定什麼時候執行它。

3.隨機數的生成

我們知道，計算機中的隨機數並不是真正意義的隨機數，是偽隨機數。一個偽隨機數生成器並不是線程安全的，因為生成新隨機數時可能會改變它的內部狀態。MXNet 使用序列化隨機數操作的方式來解決這個問題。

核心演算法（依賴跟蹤演算法）

為了以通俗易通的方式描述「依賴跟蹤演算法「，先提出以下三個概念：

變數狀態：MXNet 中，變數有三種狀態：可讀可寫，可讀不可寫，不可讀不可寫。為了保證並行狀態下數據的一致性，當有一個操作在改變（寫操作）變數時，該變數就不能被其它操作寫或讀；當一個變數正準備被某些操作讀取時，後來的操作就不能改變這個變數；當一個變數處於可寫狀態，由於沒有操作在改變它，不必考慮操作會影響數據的一致性，所以它必然是可讀的。

操作序列：MXNet 用戶寫下一個串列程序，調用「依賴引擎」的介面把程序的每一個操作按順序壓入操作序列中，並告知「依賴引擎」每一個操作依賴的每一個變數及依賴的每一個變數的狀態。

變數的未決隊列：變數當前的狀態使依賴於該變數的操作處於不可執行狀態，這樣的操作叫做未決操作。例如操作 B = A + 2 依賴於 A 處於可讀狀態，B 處於可寫狀態。當 A 處於不可讀狀態或 B 處於不可讀狀態都將導致這個操作處於未決狀態。每一個變數都有一個未決隊列，這個隊列存放著依賴於該變數的未決操作。

下面通過分析 MXNet 中 Dependency Engine 設計文檔中的例圖來更具體的理解其中的演算法思想：

前情提要

*Push Sequence 對應於上面描述的操作序列

*包含變數名的圓角矩形表示變數的當前狀態

串列程序：

A = 2

B = 2

B = A + B

C = A + 2

A.del()

我們注意到，橙色矩形表示操作（如：A=2{1},B=A+B{2}）。每一個操作後面都有一個被大括弧包圍的數字，我們以 count 來命名它。

Count 的意義：

Push Sequence：表示該操作依賴的變數個數。下面將表示為 pCount。例如:操作 A = 2 依賴與一個變數 A 的狀態，故大括弧後面為 1：A = 2{1};操作 B = A + B 依賴於兩個變數 A 和 B 的狀態，故大括弧後面為 2：B = A + B{2}。Variable Pending Queue：表示執行該操作依賴的為準備好的變數的個數。下面將表示為 vCount。例如：第三行中操作 B = A + B 依賴於 A 處於可讀狀態丶 B 處於可寫狀態，但此時 A 處於不可讀狀態，B 處於不可寫狀態，故大括弧後面為 2：B = A + B{2}；第四行中，操作 A = 2 執行完畢，這時操作 B = A + B 只依賴於操作 B = 2 的完成，故大括弧裡面為 1：B = A + B{1}。

圖中每一行的執行過程

第一行：A，B 都是可讀可寫。

第二行：第一行中，操作序列中的操作 A = 2 依賴的變數 A 處於可讀可寫狀態，滿足該操作，故執行 A = 2，並把 A 的狀態變為不可讀不可寫。同理，操作B = 2 執行，變數 B 的狀態變為不可讀不可寫。

第三行：第二行中，操作序列中的操作 B = A + B 的執行依於變數 A 的可讀狀態，由於此時 A 處於不可讀狀態，故把操作 B = A + B 壓入 A 的未決隊列並且操作 B = A + B 的 vCount 加 1,此時操作 B = A + B 的 vCount 為 0。操作 B = A +B 的執行又依賴於變數 B 的可寫狀態，由於此時 B 處於不可寫狀態，故把操作 B= A + B 壓入 B 的未決隊列並且 vCount 加 1，此時操作 B = A + B 的 vCount 為2。操作 C = A + 2 同理。

第四行：此時，操作 A = 2 執行完畢，A 的狀態變為可讀不可寫。此時遍歷A 的未決隊列，發現隊列中的兩個操作 B = A + B, C = A + 2 對於 A 的依賴都被滿足了，於是將兩個操作移除 A 的未決隊列並且各自的 vCount 減 1。此時，操作 B = A + B 的 vCount 為 1，操作 C = A + 2 的 vCount 為 0，故執行 C = A +2。

第五行：操作序列中的 A.del()依賴於 A 的可寫狀態，由於此時 A 的狀態為不可寫，故把操作 A.del()壓入 A 的未決隊列，並且操作 A.del()的 vCount 加 1。此時 A.del()的 vCount 為 1。

第六行：B = 2 執行完畢，把 B 的狀態改為可讀可寫，遍歷 B 的未決隊列，發現操作 B = A + B 對於 B 的依賴滿足，故執行 B = A + B 並把 B 的狀態標記為不可讀不可寫。操作 C = A +2 執行完畢，A 的狀態不改變，故不遍歷 A 的未決隊列，C 的狀態變為可讀可寫。

第七行：B = A + B 執行完畢，A 和 B 的狀態都變為可讀可寫。遍歷 A 的未決隊列，發現 A 的狀態滿足操作 A.del()，故執行 A.del()，A 的狀態變為不可讀不可寫。

第八行：A.del()執行完畢，A 變為可讀可寫。

ps.

以上即為該串列程序在「依賴引擎「中的執行流程。細心的人會發現，以上對執行流程的描述過程中，對變數狀態的改變的地方描述得模糊不清。比如描述第四行執行過程中，操作 A = 2 執行完畢，A 的狀態為何變為可讀不可寫不變為可讀可寫；又如第六行，B = 2 執行完畢，為何 B 的狀態變為可讀可寫不變為其它狀態。

參考 MXNet 文檔中關於內存管理部分可知，MXNet 用到了類似於引用計數的方式來管理內存。知道這點，就可以解釋變數的狀態改變了。對於第四行的執行過程，操作 A = 2 執行完畢，A 的狀態由不可寫變為可寫（可寫必然可讀，在上述的「變數狀態「中已經闡述原因）。但由於操作 B = A + B, C = A + 2 引用了變數 A，此時變數 A 的引用計數為 2，故 A 的狀態不能變為可寫，又為了能讓對A 有讀需求的操作（B = A + B, C = A + 2）得到執行，故 A 的狀態變為可讀不可寫。同理，第六行，B = 2 執行完畢，且此時沒有其它操作引用了 B，為了讓

關於隨機數

隨機數生成器在生成一個新的隨機數的時候會改變內部狀態，所以把生成隨機數的操作當作對隨機數生成器的寫操作即可，這裡不再過多的闡述，具體可參考下面這個例子。

關於並行執行操作

以上只描述了如何通過依賴關係來調度操作，並沒有說到如何並行的執行這些操作。其實在 MXNet 中維護了一個線程池，當同一時刻有數個操作的依賴得到滿足，那麼「依賴引擎「把它們分發到不同（或指定的）的且處於空閑狀態的線程上去執行，這樣就達到了以非同步並行的方式優化用戶程序的目的。

總結

用戶首先把操作的順序和依賴關係告訴「依賴引擎「，引擎通過判斷依賴是否滿足來調度操作的執行，並把同一時刻滿足依賴的所有操作儘可能的並行執行。並通過執行完操作修改變數狀態的方式來確定新的依賴狀態，從而達到跟蹤依賴的目的。

*以上只是通過 MXNet 設計文檔得到的結論，關於具體實現細節還是得通過閱讀MXNet 源碼的方式進一步了解。