MXNet設計和實現簡介

時間 2019-11-21

標籤 mxnet 設計實現簡介简体版

原文原文鏈接

Docs »javascript
MXNet設計和實現簡介html
Edit on GitHubjava

The following is an overview of MXNet in Chinese. For english readers, please refer to our NIPS learnsys paperpython

MXNet設計和實現簡介

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

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

這兩種方式各有利弊，總結以下算法

| |淺嵌入，命令式編程 | 深嵌入，聲明式編程 | | — |— | — |數據庫

| 如何執行 a=b+1 | 須要b已經被賦值。當即執行加法，將結果保存在a中。 | 返回對應的計算圖(computation graph)，咱們能夠以後對b進行賦值，而後再執行加法運算 |編程

| 優勢 | 語義上容易理解，靈活，能夠精確控制行爲。一般能夠無縫地和主語言交互，方便地利用主語言的各種算法，工具包，debug和性能調試器。 | 在真正開始計算的時候已經拿到了整個計算圖，因此咱們能夠作一系列優化來提高性能。實現輔助函數也容易，例如對任何計算圖都提供forward和 backward函數，對計算圖進行可視化，將圖保存到硬盤和從硬盤讀取。|數組

| 缺點 | 實現統一的輔助函數和提供總體優化都很困難。 | 不少主語言的特性都用不上。某些在主語言中實現簡單，但在這裏卻常常麻煩，例如if-else語句。debug也不容易，例如監視一個複雜的計算圖中的某個節點的中間結果並不簡單。|

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

下表咱們比較MXNet和其餘流行的深度學習系統

| | 主語言 | 從語言 | 硬件 | 分佈式 | 命令式 | 聲明式 |

| — | — | — | — | :—: | :—: | :—:|

| Caffe | C++ | Python/Matlab | CPU/GPU | x | x | v |

| Torch | Lua | - | CPU/GPU/FPGA | x | v | x |

| Theano | Python | - | CPU/GPU | x | v | x |

| TensorFlow | C++ | Python | CPU/GPU/Mobile | v | x | v |

| MXNet | C++ | Python/R/Julia/Go | CPU/GPU/Mobile | v | v | v |

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

MXNet的系統架構以下圖所示：

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

編程接口

`Symbol`：聲明式的符號表達式

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

using MXNetmlp = @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 < n; ++i) {
  net.forward();
  net.backward();
  net.weight -= eta * net.grad}

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

上面的混合實現跟使用純符號表達式實現的性能相差無二，而後後者在表達控制邏輯時會更加複雜。其緣由是NDArray的執行會和Symbol相似的構建一個計算圖，並與其餘運算一同交由後臺引擎執行。對於運算-=由於咱們只是將其結果交給另外一個Symbol的forward做爲輸入，所以咱們不須要當即獲得結果。當上面的for循環結束時，咱們只是將數個 Symbol和NDarray對應的計算圖提交給了後臺引擎。當咱們最終須要結果的時候，例如將weight複製到主語言中或者保存到磁盤時，程序纔會被阻塞直到全部計算完成。

`KVStore`：多設備間的數據交互

MXNet提供一個分佈式的key-value存儲來進行數據交換。它主要有兩個函數，

push：將key-value對從一個設備push進存儲
pull：將某個key上的值從存儲中pull出來此外，KVStore還接受自定義的更新函數來控制收到的值如何寫入到存儲中。最後KVStore提供數種包含最終一致性模型和順序一致性模型在內的數據一致性模型。

在下面例子中，咱們將前面的梯度降低算法改爲分佈式梯度降低。

KVStore kv("dist_async");kv.set_updater([](NDArray w, NDArray g) {
    w -= eta * g;
  });for (int i = 0; i < max_iter; ++i) {
   kv.pull(net.weight);
   net.forward();
   net.backward();
   kv.push(net.grad);}

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

這裏push和pull跟NDArray同樣使用了延後計算的技術。它們只是將對應的操做提交給後臺引擎，而引擎則調度實際的數據交互。因此上述的實現跟咱們使用純符號實現的性能相差無幾。

讀入數據模塊

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

一般數據存在本地磁盤或者遠端的分佈式文件系統上（例如HDFS或者Amazon S3)，每次咱們只須要將當前須要的數據讀進內存。MXNet提供迭代器能夠按塊讀取不一樣格式的文件。迭代器使用多線程來解碼數據，並使用多線程預讀取來隱藏文件讀取的開銷。

訓練模塊

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

import MXNet as mxmodel = 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會對計算圖進行優化，以及爲全部變量提早申請空間。

計算圖優化

計算圖優化已經在數據庫等領域被研究多年，咱們目前只探索了數個簡單的方法。

注意到咱們提早申明瞭哪些輸出變量是須要的，這樣咱們只須要計算這些輸出須要的操做。例如，在預測時咱們不須要計算梯度，因此整個backforward圖均可以忽略。而在特徵抽取中，咱們可能只須要某些中間層的輸出，從而能夠忽略掉後面的計算。
咱們能夠合併某些操做。例如 ab+1*只須要一個blas或者cuda函數便可，而不須要將其表示成兩個操做。
咱們實現了一些「大」操做，例如一個卷積層就只須要一個操做子。這樣咱們能夠大大減少計算圖的大小，而且方便手動的對這個操做進行優化。

內存申請

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

MXNet提供了兩個啓發式的策略，每一個策略都是線性的複雜度。

inplace。在這個策略裏，咱們模擬圖的遍歷過程，併爲每一個變量維護一個還有多少其餘變量須要它的計數。當咱們發現某個變量的計數變成0時，咱們便回收其內存空間。
co-share。咱們容許兩個變量使用同一段內存空間。這麼作固然會使得這兩個變量不能同時在寫這段空間。因此咱們只考慮對不能並行的變量進行 co-share。每一次咱們考慮圖中的一條路（path），路上全部變量都有依賴關係因此不能被並行，而後咱們對其進行內存分配並將它們從圖中刪掉。

引擎

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

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

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

數據通信

KVStore的實現是基於參數服務器。但它跟前面的工做有兩個顯著的區別。

咱們經過引擎來管理數據一致性，這使得參數服務器的實現變得至關簡單，同時使得KVStore的運算能夠無縫的與其餘結合在一塊兒。
咱們使用一個兩層的通信結構，原理以下圖所示。第一層的服務器管理單機內部的多個設備之間的通信。第二層服務器則管理機器之間經過網絡的通信。第一層的服務器在與第二層通信前可能合併設備之間的數據來下降網絡帶寬消費。同時考慮到機器內和外通信帶寬和延時的不一樣性，咱們能夠對其使用不一樣的一致性模型。例如第一層咱們用強的一致性模型，而第二層咱們則使用弱的一致性模型來減小同步開銷。

可移植性

輕量和可移植性是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的目的是作一個有意思的系統，可以讓你們用着方便的系統，一個輕量的和能夠快速測試系統和算法想法的系統。對於將來，咱們主要關注下面四個方向：

支持更多的硬件，咱們目前在積極考慮支持AMD GPU，高通GPU，Intel Phi，FPGA，和更多智能設備。相信MXNet的輕量和內存節省能夠在這些上大有做爲。
更加完善的操做子。目前不管是Symbol仍是NDArray支持的操做仍是有限，咱們但願可以儘快的擴充他們。
更多編程語言。除了C++，目前MXNet對Python，R和Julia的支持比較完善。但咱們但願還能有不少的語言，例如javascript。
更多的應用。咱們以前花了不少精力在圖片分類上，下面咱們會考慮不少的應用。例如上週咱們試了下如何利用一張圖片的風格和一張圖片的內容合成一張新圖片。下圖是利用我辦公室窗景和梵高的starry night來合成圖片

接下來咱們但願可以在更多應用，例如語音、翻譯、問答上有所產出。

咱們忠心但願MXNet能爲你們作深度學習相關研究和應用帶來便利。也但願能與更多的開發者一塊兒學習和進步。