754      阿裏雲  技術社區[雲棲]

【筆談】Github上新鮮出爐的深度學習係統mxnet的中文介紹

MXNet設計和實現簡介

神經網絡本質上是一種語言，我們通過它來表達對應用問題的理解。例如我們用卷積層來表達空間相關性，RNN來表達時間連續性。根據問題的複雜性和信息如何從輸入到輸出一步步提取，我們將不同大小的層按一定原則連接起來。近年來隨著數據的激增和計算能力的大幅提升，神經網絡也變得越來越深和大。例如最近幾次imagnet競賽的冠軍都使用有數十至百層的網絡。對於這一類神經網絡我們通常稱之為深度學習。從應用的角度而言，對深度學習最重要的是如何方便的表述神經網絡，以及如何快速訓練得到模型。

對於一個優秀的深度學習係統，或者更廣來說優秀的科學計算係統，最重要的編程接口的設計。他們都采用將一個領域特定語言(domain specific language)嵌入到一個主語言中。例如numpy將矩陣運算嵌入到python中。這類嵌入一般分為兩種，其中一種嵌入的較淺，其中每個語句都按原來的意思執行，且通常采用命令式編程(imperative programming)，其中numpy和Torch就是屬於這種。而另一種則用一種深的嵌入方式，提供一整套針對具體應用的迷你語言。這一種通常使用聲明式語言(declarative programing)，既用戶隻需要聲明要做什麼，而具體執行則由係統完成。這類係統包括Caffe，theano和剛公布的TensorFlow。

這兩種方式各有利弊，總結如下

目前現有的係統大部分都采用上兩種編程模式的一種。與它們不同的是，MXNet嚐試將兩種模式無縫的結合起來。在命令式編程上MXNet提供張量運算，而聲明式編程中MXNet支持符號表達式。用戶可以自由的混合它們來快速實現自己的想法。例如我們可以用聲明式編程來描述神經網絡，並利用係統提供的自動求導來訓練模型。另一方便，模型的迭代訓練和更新模型法則中可能涉及大量的控製邏輯，因此我們可以用命令式編程來實現。同時我們用它來進行方便的調式和與主語言交互數據。

下表我們比較MXNet和其他流行的深度學習係統

（注，TensforFlow暫時沒有公開其分布式實現）

MXNet的係統架構如下圖所示：

從上到下分別為各種主語言的嵌入，編程接口（矩陣運算，符號表達式，分布式通訊），兩種編程模式的統一係統實現，以及各硬件的支持。接下一章我們將介紹編程接口，然後下一章介紹係統實現。之後我們給出一些實驗對比結果，以及討論MXNet的未來。

編程接口

Symbol： 聲明式的符號表達式

MXNet使用多值輸出的符號表達式來聲明計算圖。符號是由操作子構建而來。一個操作子可以是一個簡單的矩陣運算“+”，也可以是一個複雜的神經網絡裏麵的層，例如卷積層。一個操作子可以有多個輸入變量和多個輸出變量，還可以有內部狀態變量。一個變量既可以是自由的，我們可以之後對其賦值；也可以是某個操作子的輸出。例如下麵的代碼中我們使用Julia來定義一個多層感知機，它由一個代表輸入數據的自由變量，和多個神經網絡層串聯而成。

using MXNet
mlp = @mx.chain mx.Variable(:data) =>
  mx.FullyConnected(num_hidden=64) =>
  mx.Activation(act_type=:relu)    =>
  mx.FullyConnected(num_hidden=10) =>
  mx.Softmax()

在執行一個符號表達式前，我們需要對所有的自由變量進行賦值。上例中，我們需要給定數據，和各個層裏隱式定義的輸入，例如全連接層的權重和偏值。我們同時要申明所需要的輸出，例如softmax的輸出。

除了執行獲得softmax輸出外（通常也叫forward），符號表達式也支持自動求導來獲取各權重和偏值對應的梯度（也稱之為backward）。此外，我們還可以提前估計計算時需要的內存，符號表達式的可視化，讀入和輸出等。

NDArray：命令式的張量計算

MXNet提供命令式的張量計算來橋接主語言的和符號表達式。下麵代碼中，我們在GPU上計算矩陣和常量的乘法，並使用numpy來打印結果

>>> import MXNet as mx
>>> a = mx.nd.ones((2, 3),
... mx.gpu())
>>> print (a * 2).asnumpy()
[[ 2.  2.  2.]
 [ 2.  2.  2.]]

另一方麵，NDArray可以無縫和符號表達式進行對接。假設我們使用Symbol定義了一個神經網絡，那麼我們可以如下實現一個梯度下降算法

for (int i = 0; i < max_iter; ++i) {
  network.forward();
  network.backward();
  network.weight -= eta * network.gradient
}

這裏梯度由Symbol計算而得。Symbol的輸出結果均表示成NDArray，我們可以通過NDArray提供的張量計算來更新權重。此外，我們還利用了主語言的for循環來進行迭代，學習率eta也是在主語言中進行修改。

上麵的混合實現跟使用純符號表達式實現的性能相差無二，然後後者在表達控製邏輯時會更加複雜。其原因是NDArray的執行會和Symbol類似的構建一個計算圖，並與其他運算一同交由後台引擎執行。對於運算-=由於我們隻是將其結果交給另一個Symbol的forward作為輸入，因此我們不需要立即得到結果。當上麵的for循環結束時，我們隻是將數個Symbol和NDarray對應的計算圖提交給了後台引擎。當我們最終需要結果的時候，例如將weight複製到主語言中或者保存到磁盤時，程序才會被阻塞直到所有計算完成。

KVStore：多設備間的數據交互

MXNet提供一個分布式的key-value存儲來進行數據交換。它主要有兩個函數，
1. push： 將key-value對從一個設備push進存儲
2. pull：將某個key上的值從存儲中pull出來此外，KVStore還接受自定義的更新函數來控製收到的值如何寫入到存儲中。最後KVStore提供數種包含最終一致性模型和順序一致性模型在內的數據一致性模型。

在下麵例子中，我們將前麵的梯度下降算法改成分布式梯度下降。

KVStore kvstore("dist_async");
kvstore.set_updater([](NDArray weight, NDArray gradient) {
    weight -= eta * gradient;
  });
for (int i = 0; i < max_iter; ++i) {
   kvstore.pull(network.weight);
   network.forward();
   network.backward();
   kvstore.push(network.gradient);
}

在這裏我們先使用最終一致性模型創建一個kvstore，然後將更新函數注冊進去。在每輪迭代前，每個計算節點先將最新的權重pull回來，之後將計算的得到的梯度push出去。kvstore將會利用更新函數來使用收到的梯度更新其所存儲的權重。

這裏push和pull跟NDArray一樣使用了延後計算的技術。它們隻是將對應的操作提交給後台引擎，而引擎則調度實際的數據交互。所以上述的實現跟我們使用純符號實現的性能相差無幾。

讀入數據模塊

數據讀取在整體係統性能上占重要地位。MXNet提供工具能將任意大小的樣本壓縮打包成單個或者數個文件來加速順序和隨機讀取。

通常數據存在本地磁盤或者遠端的分布式文件係統上（例如HDFS或者Amazon S3)，每次我們隻需要將當前需要的數據讀進內存。MXNet提供迭代器可以按塊讀取不同格式的文件。迭代器使用多線程來解碼數據，並使用多線程預讀取來隱藏文件讀取的開銷。

訓練模塊

MXNet實現了常用的優化算法來訓練模型。用戶隻需要提供數據數據迭代器和神經網絡的Symbol便可。此外，用戶可以提供額外的KVStore來進行分布式的訓練。例如下麵代碼使用分布式異步SGD來訓練一個模型，其中每個計算節點使用兩快GPU。

import MXNet as mx
model = mx.model.FeedForward(
    ctx                = [mx.gpu(0), mx.gpu(1)],
    symbol             = network,
    num_epoch          = 100,
    learning_rate      = 0.01,
    momentum           = 0.9,
    wd                 = 0.00001,
    initializer        = mx.init.Xavier(factor_type="in", magnitude=2.34))
model.fit(
    X                  = train_iter,
    eval_data          = val_iter,
    kvstore            = mx.kvstore.create('dist_async'),
    epoch_end_callback = mx.callback.do_checkpoint('model_'))

係統實現

計算圖

一個已經賦值的符號表達式可以表示成一個計算圖。下圖是之前定義的多層感知機的部分計算圖，包含forward和backward。

其中圓表示變量，方框表示操作子，箭頭表示數據依賴關係。在執行之前，MXNet會對計算圖進行優化，以及為所有變量提前申請空間。

計算圖優化

計算圖優化已經在數據庫等領域被研究多年，我們目前隻探索了數個簡單的方法。
1. 注意到我們提前申明了哪些輸出變量是需要的，這樣我們隻需要計算這些輸出需要的操作。例如，在預測時我們不需要計算梯度，所以整個backforward圖都可以忽略。而在特征抽取中，我們可能隻需要某些中間層的輸出，從而可以忽略掉後麵的計算。
2. 我們可以合並某些操作。例如 ab+1*隻需要一個blas或者cuda函數即可，而不需要將其表示成兩個操作。
3. 我們實現了一些“大”操作，例如一個卷積層就隻需要一個操作子。這樣我們可以大大減小計算圖的大小，並且方便手動的對這個操作進行優化。

內存申請

內存通常是一個重要的瓶頸，尤其是對GPU和智能設備而言。而神經網絡計算時通常需要大量的臨時空間，例如每個層的輸入和輸出變量。對每個變量都申請一段獨立的空間會帶來高額的內存開銷。幸運的是，我們可以從計算圖推斷出所有變量的生存期，就是這個變量從創建到最後被使用的時間段，從而可以對兩個不交叉的變量重複使用同一內存空間。這個問題在諸多領域，例如編譯器的寄存器分配上，有過研究。然而最優的分配算法需要 O(n2) 時間複雜度，這裏n是圖中變量的個數。

MXNet提供了兩個啟發式的策略，每個策略都是線性的複雜度。
1. inplace。在這個策略裏，我們模擬圖的遍曆過程，並為每個變量維護一個還有多少其他變量需要它的計數。當我們發現某個變量的計數變成0時，我們便回收其內存空間。
2. co-share。我們允許兩個變量使用同一段內存空間。這麼做當然會使得這兩個變量不能同時在寫這段空間。所以我們隻考慮對不能並行的變量進行co-share。每一次我們考慮圖中的一條路（path），路上所有變量都有依賴關係所以不能被並行，然後我們對其進行內存分配並將它們從圖中刪掉。

引擎

在MXNet中，所有的任務，包括張量計算，symbol執行，數據通訊，都會交由引擎來執行。首先，所有的資源單元，例如NDArray，隨機數生成器，和臨時空間，都會在引擎處注冊一個唯一的標簽。然後每個提交給引擎的任務都會標明它所需要的資源標簽。引擎則會跟蹤每個資源，如果某個任務所需要的資源到到位了，例如產生這個資源的上一個任務已經完成了，那麼引擎會則調度和執行這個任務。

通常一個MXNet運行實例會使用多個硬件資源，包括CPU，GPU，PCIe通道，網絡，和磁盤，所以引擎會使用多線程來調度，既任何兩個沒有資源依賴衝突的任務都可能會被並行執行，以求最大化資源利用。

與通常的數據流引擎不同的是，MXNet的引擎允許一個任務修改現有的資源。為了保證調度正確性，提交任務時需要分開標明哪些資源是隻讀，哪些資源會被修改。這個附加的寫依賴可以帶來很多便利。例如我們可以方便實現在numpy以及其他張量庫中常見的數組修改操作，同時也使得內存分配時更加容易，比如操作子可以修改其內部狀態變量而不需要每次都重來內存。再次，假如我們要用同一個種子生成兩個隨機數，那麼我們可以標注這兩個操作會同時修改種子來使得引擎不會並行執行，從而使得代碼的結果可以很好的被重複。

數據通訊

KVStore的實現是基於參數服務器。但它跟前麵的工作有兩個顯著的區別。
1. 我們通過引擎來管理數據一致性，這使得參數服務器的實現變得相當簡單，同時使得KVStore的運算可以無縫的與其他結合在一起。
2. 我們使用一個兩層的通訊結構，原理如下圖所示。第一層的服務器管理單機內部的多個設備之間的通訊。第二層服務器則管理機器之間通過網絡的通訊。第一層的服務器在與第二層通訊前可能合並設備之間的數據來降低網絡帶寬消費。同時考慮到機器內和外通訊帶寬和延時的不同性，我們可以對其使用不同的一致性模型。例如第一層我們用強的一致性模型，而第二層我們則使用弱的一致性模型來減少同步開銷。

可移植性

輕量和可移植性是MXNet的一個重要目標。MXNet核心使用C++實現，並提供C風格的頭文件。因此方便係統移植，也使得其很容易被其他支持C FFI (forigen language interface )的語言調用。此外，我們也提供一個腳本將MXNet核心功能的代碼連同所有依賴打包成一個單一的隻有數萬行的C++源文件，使得其在一些受限的平台，例如智能設備，方便編譯和使用。

實驗結果

這裏我們提供一些早期的實驗結果。

與其他係統相比

我們首先使用一個流行卷積網絡測試方案來對比MXNet與Torch，Caffe和TensorFlow在過去幾屆imagenet競賽冠軍網絡上的性能。每個係統使用同樣的CUDA 7.0和CUDNN 3，但TensorFlow使用其隻支持的CUDA 6.5 和CUDNN 2。我們使用單塊GTX 980並報告單個forward和backward的耗時。

可以看出MXNet，Torch和Caffe三者在性能上不相上下。這個符合預期，因為在單卡上我們評測的幾個網絡的絕大部分運算都由CUDA和CUDNN完成。TensorFlow比其他三者都慢2倍以上，這可能由於是低版本的CUDNN和項目剛開源的緣故。

內存的使用

接下來我們考察不同的內存分配算法對內存占用的影響。下圖分別表示使用batch=128時，在做預測時和做訓練時的不同算法在內部變量（除去模型，最初輸入和最終輸出）上的內存開銷。

可以看出，inplace和co-share兩者都可以極大的降低內存使用。將兩者合起來可以在訓練時減少2倍內存使用，在預測時則可以減小4倍內存使用。特別的，即使是最複雜的vggnet，對單張圖片進行預測時，MXNet隻需要16MB額外內存。

Scalability

最後我們報告在分布式訓練下的性能。我們使用imagenet 1k數據（120萬224x224x3圖片，1000類），並用googlenet加上batch normalization來訓練。我們使用Amazon EC2 g2.8x，單機和多機均使用同樣的參數，下圖表示了使用單機和10台g2.8x時的收斂情況。

從訓練精度來看，單機的收斂比多機快，這個符合預期，因為多機時有效的batch大小比單機要大，在處理同樣多的數據上收斂通常會慢。但有意思的是兩者在測試精度上非常相似。

單機下每遍曆一次數據需要1萬4千秒，而在十台機器上，每次隻需要1千4百秒。如果考慮運行時間對比測試精度，10台機器帶來了10倍的提升。

過去，現狀，和未來

大半年前我們拉來數個優秀的C++機器學習係統的開發人員成立了DMLC，本意是更方便共享各自項目的代碼，並給用戶提供一致的體驗。當時我們有兩個深度學習的項目，一個是CXXNet，其通過配置來定義和訓練神經網絡。另一個是Minerva，提供類似numpy一樣的張量計算接口。前者在圖片分類等使用卷積網絡上很方便，而後者更靈活。那時候我們想要能不能一個兩者功能都具備的係統，於是這樣就有了MXNet。其名字來自Minerva的M和CXXNet的XNet。其中Symbol的想法來自CXXNet，而NDArray的想法來自Minerva。我們也常把MXNet叫“mix net”。

MXNet是DMLC第一個結合了所有的成員努力的項目，也同時吸引了很多核心成員的加入。MXNet目的是做一個有意思的係統，能夠讓大家用著方便的係統，一個輕量的和可以快速測試係統和算法想法的係統。對於未來，我們主要關注下麵四個方向：

1. 支持更多的硬件，我們目前在積極考慮支持AMD GPU，高通GPU，Intel Phi，FPGA，和更多智能設備。相信MXNet的輕量和內存節省可以在這些上大有作為。
2. 更加完善的操作子。目前不論是Symbol還是NDArray支持的操作還是有限，我們希望能夠盡快的擴充他們。
3. 更多編程語言。除了C++，目前MXNet對Python，R和Julia的支持比較完善。但我們希望還能有很多的語言，例如javascript。

4. 更的應用。我們之前花了很多精力在圖片分類上，下麵我們會考慮很多的應用。例如上周我們試了下如何利用一張圖片的風格和一張圖片的內容合成一張新圖片。下圖是利用我辦公室窗景和梵高的starry night來合成圖片

   

   

   接下來我們希望能夠在更多應用，例如語音，翻譯，問答，上有所產出。

我們忠心希望MXNet能為大家做深度學習相關研究和應用帶來便利。也希望能與更多的開發者一起學習和進步。

擴展閱讀

1. 此文大部分內容已經發表在NIPS LearningSys 2016上，paper link
2. 本文隻是對MXNet各個部件做了初步的介紹，更多文檔參見 MXNet/doc
3. 本文實驗代碼均在 MXNet/example

最後更新：2017-04-01 13:37:10

  上一篇： 平台化三部曲之二模塊化開發 - Google Guice 平台模塊化開發的果汁

  下一篇： android仿淘寶刮刮卡功能實現