mxnet framework簡介

最近一直都在看mxnet的源碼，用了近兩個月時間，終於研究的七七八八。

下面簡單總結下mxnet的系統級架構。

(一) Mxnet系統結構

整個系統架構如圖所示，上面兩層是用戶層(黃色)，包含幾大模塊：kvstore模塊，IO模塊，NDArry模塊，Executor模塊和Symbol模塊。每個模塊都提供了可供用戶操作的API。

最底下一層是系統層(藍色)，主要是mxnet內部交互使用，包括StorageAllocator模塊，Engine模塊，ResourceManager模塊和Operator模塊。

想了解各個模塊的功能，可以通過一個數據並行的訓練過程簡單看懂。

(二)單機多卡的數據並行訓練過程

Pre-process的過程

(1) 用戶調用Symbol模塊的API來定義整個深度神經網路。

a) 這裡用戶需要定義好每層網路做的事情，以及每層網路的輸入

b) 具體每層網路的實際計算在Operator模塊中定義。該模塊具體實現對用戶不可見。

(2) 用戶調用IO模塊來構造訓練數據集的迭代器。

a) 這裡的IO模塊主要是用來把用戶存好的訓練數據封裝成mxnet易於優化的迭代器結構

(3) 用戶調用IO模塊來構造校驗數據集的迭代器

(4) 用戶定義訓練參數

a) 調用kvstore模塊來定義全局的kvstore。單機多卡環境下使用device類型的kvstore。多機環境下用dist類型。

b) 定義訓練需要的所有設備(GPU) devs。一般情況下每個設備Load一套模型

c) 其他參數，包括learning rate, epoch size, data batch size等

(5) 用戶調用Executor模塊構造執行器executorgroup。

a) 每個executor綁定一個設備，executor之間通過kvstore交互信息。

(6) 系統初始化每個executor。Executor的初始化涉及到mxnet與多個系統級模塊的交互。其步驟如下：

a) 系統調用Symbol模塊去把用戶定義的Symbol轉化成mxnet使用的計算圖結構。每個節點可以簡單的看成operator

i. 構造計算圖的第一步是簡單的把用戶定義的Symbol連成圖結構

ii. 補充圖結構的鏡像節點，用來做反向傳播

iii. 對於前後兩個節點在不同設備上的情況，會添加新的節點用來做數據拷貝

iv. 深度優先法遍歷計算圖，得到其拓撲排序

b) 系統調用StorageAllocator模塊為計算圖分配內存空間

i. 這裡會對計算圖先進行著色，不可並行的路徑用相同顏色。然後著色相同的節點可以用相同的內存塊保存數據，從而節省內存空間。

ii. 部分節點的輸入和輸出還會使用同一塊內存塊來節省空間。

iii. 所有的內存塊都是以NDArray的形式保存在系統內，這裡都會調用NDArray模塊

c) 系統調用ResourceManager模塊為某些節點分配臨時資源

d) 將某些節點合併成一個更大的節點，以節省操作數量(例如在CNN中將conv+ReLU+ BN打包成一個操作)。

(7) 系統初始化kvstore

a) 主要是給每個key綁定value的形狀以及分配緩衝區。

Process過程

在每個epoch內

(1) 對每個data batch

a) 系統調用executor group的forward函數完成前向傳播

i. executor group的forward函數會依次調用每個executor的forward函數完成前向傳播計算

ii. 每個executor的forward函數實際上是把executor內部的計算圖按照拓撲排序次序依次push到engine模塊啟動的engine中。

iii. Engine會利用線程池技術，利用一個調度演算法，來非同步執行每個計算圖節點的Operator，從而達到多executor的並行計算。

b) 系統調用executor group的backward函數完成反向傳播，更新模型的參數值

i. 與前向傳播過程類似

c) 調用kvstore同步多塊卡之間模型的參數值

i. 每個executor在一個設備上執行，因為數據不同，計算出的模型參數值也不同。這時就需要通過kvstore同步

ii. kvstore先做一個reduce，將不同executor計算出的模型參數merge到一起，再做一個broadcast將該值發送給所有executor以達到同步目的

iii. 進一步來說kvstore里還會傳入optimizer，來做參數更新。optimizer就是整個網路使用的優化演算法(例如SGD)，kvstore內部的緩衝區會儲存當前的所有param的值（即訓練好的weight），然後調用push(grad)，把不同卡上計算的梯度先merge到一起，然後會計算param+= optimizer(grad)來更新param。例如梯度下降法，optimizer（grad）函數實際就是計算的梯度方向，這裡optimizer(grad)=learning rate * grad。其他優化演算法就會用其他方法計算下降方向然後更新param。最後kvstore調用pull(weight)把緩衝區里存的param值broadcast到每個executor里的weight里，從而達到同步目的。

d) 執行batch end的回調函數。一般是統計計算速度等

(2) 執行epoch end的回調函數。一般是做checkpoint存儲當前計算好的模型到文件

(3) 用校驗數據檢驗當前模型的精度

預測過程完全類似，只不過只需要構造一個executor，不需要kvstore做多executor間的通信。