從上到下分別為各種主語言的嵌入,編程接口(矩陣運算,符號表達式,分布式通訊),兩種編程模式的統一系統實現,以及各硬件的支持。接下一章我們將介紹編程接口,然后下一章介紹系統實現。之后我們給出一些實驗對比結果,以及討論MXNet的未來。
MXNet使用多值輸出的符號表達式來聲明計算圖。符號是由操作子構建而來。一個操作子可以是一個簡單的矩陣運算“+”,也可以是一個復雜的神經網絡里面的層,例如卷積層。一個操作子可以有多個輸入變量和多個輸出變量,還可以有內部狀態變量。一個變量既可以是自由的,我們可以之后對其賦值;也可以是某個操作子的輸出。例如下面的代碼中我們使用Julia來定義一個多層感知機,它由一個代表輸入數據的自由變量,和多個神經網絡層串聯而成。
在執行一個符號表達式前,我們需要對所有的自由變量進行賦值。上例中,我們需要給定數據,和各個層里隱式定義的輸入,例如全連接層的權重和偏值。我們同時要申明所需要的輸出,例如softmax的輸出。
除了執行獲得softmax輸出外(通常也叫forward),符號表達式也支持自動求導來獲取各權重和偏值對應的梯度(也稱之為backward)。此外,我們還可以提前估計計算時需要的內存,符號表達式的可視化,讀入和輸出等。
這里梯度由Symbol計算而得。Symbol的輸出結果均表示成NDArray,我們可以通過NDArray提供的張量計算來更新權重。此外,我們還利用了主語言的for循環來進行迭代,學習率eta也是在主語言中進行修改。
上面的混合實現跟使用純符號表達式實現的性能相差無二,然后后者在表達控制邏輯時會更加復雜。其原因是NDArray的執行會和Symbol類似的構建一個計算圖,并與其他運算一同交由后臺引擎執行。對于運算-=由于我們只是將其結果交給另一個Symbol的forward作為輸入,因此我們不需要立即得到結果。當上面的for循環結束時,我們只是將數個 Symbol和NDarray對應的計算圖提交給了后臺引擎。當我們最終需要結果的時候,例如將weight復制到主語言中或者保存到磁盤時,程序才會被阻塞直到所有計算完成。
- push: 將key-value對從一個設備push進存儲
- pull:將某個key上的值從存儲中pull出來此外,KVStore還接受自定義的更新函數來控制收到的值如何寫入到存儲中。最后KVStore提供數種包含最終一致性模型和順序一致性模型在內的數據一致性模型。 </p>
在下面例子中,我們將前面的梯度下降算法改成分布式梯度下降。
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會對計算圖進行優化,以及為所有變量提前申請空間。
計算圖優化
計算圖優化已經在數據庫等領域被研究多年,我們目前只探索了數個簡單的方法。
- 注意到我們提前申明了哪些輸出變量是需要的,這樣我們只需要計算這些輸出需要的操作。例如,在預測時我們不需要計算梯度,所以整個backforward圖都可以忽略。而在特征抽取中,我們可能只需要某些中間層的輸出,從而可以忽略掉后面的計算。
- 我們可以合并某些操作。例如 ab+1*只需要一個blas或者cuda函數即可,而不需要將其表示成兩個操作。
- 我們實現了一些“大”操作,例如一個卷積層就只需要一個操作子。這樣我們可以大大減小計算圖的大小,并且方便手動的對這個操作進行優化。 </p>
內存申請
內存通常是一個重要的瓶頸,尤其是對GPU和智能設備而言。而神經網絡計算時通常需要大量的臨時空間,例如每個層的輸入和輸出變量。對每個變量都申請一段獨立的空間會帶來高額的內存開銷。幸運的是,我們可以從計算圖推斷出所有變量的生存期,就是這個變量從創建到最后被使用的時間段,從而可以對兩個不交叉的變量重復使用同一內存空間。這個問題在諸多領域,例如編譯器的寄存器分配上,有過研究。然而最優的分配算法需要 O(n2) 時間復雜度,這里n是圖中變量的個數。
MXNet提供了兩個啟發式的策略,每個策略都是線性的復雜度。
- inplace。在這個策略里,我們模擬圖的遍歷過程,并為每個變量維護一個還有多少其他變量需要它的計數。當我們發現某個變量的計數變成0時,我們便回收其內存空間。
- co-share。我們允許兩個變量使用同一段內存空間。這么做當然會使得這兩個變量不能同時在寫這段空間。所以我們只考慮對不能并行的變量進行co- share。每一次我們考慮圖中的一條路(path),路上所有變量都有依賴關系所以不能被并行,然后我們對其進行內存分配并將它們從圖中刪掉。 </p>
引擎
在MXNet中,所有的任務,包括張量計算,symbol執行,數據通訊,都會交由引擎來執行。首先,所有的資源單元,例如NDArray,隨機數生成器,和臨時空間,都會在引擎處注冊一個唯一的標簽。然后每個提交給引擎的任務都會標明它所需要的資源標簽。引擎則會跟蹤每個資源,如果某個任務所需要的資源到到位了,例如產生這個資源的上一個任務已經完成了,那么引擎會則調度和執行這個任務。
通常一個MXNet運行實例會使用多個硬件資源,包括CPU,GPU,PCIe通道,網絡,和磁盤,所以引擎會使用多線程來調度,既任何兩個沒有資源依賴沖突的任務都可能會被并行執行,以求最大化資源利用。
與通常的數據流引擎不同的是,MXNet的引擎允許一個任務修改現有的資源。為了保證調度正確性,提交任務時需要分開標明哪些資源是只讀,哪些資源會被修改。這個附加的寫依賴可以帶來很多便利。例如我們可以方便實現在numpy以及其他張量庫中常見的數組修改操作,同時也使得內存分配時更加容易,比如操作子可以修改其內部狀態變量而不需要每次都重來內存。再次,假如我們要用同一個種子生成兩個隨機數,那么我們可以標注這兩個操作會同時修改種子來使得引擎不會并行執行,從而使得代碼的結果可以很好的被重復。
數據通訊
KVStore的實現是基于參數服務器。但它跟前面的工作有兩個顯著的區別。
- 我們通過引擎來管理數據一致性,這使得參數服務器的實現變得相當簡單,同時使得KVStore的運算可以無縫的與其他結合在一起。
- 我們使用一個兩層的通訊結構,原理如下圖所示。第一層的服務器管理單機內部的多個設備之間的通訊。第二層服務器則管理機器之間通過網絡的通訊。第一層的服務器在與第二層通訊前可能合并設備之間的數據來降低網絡帶寬消費。同時考慮到機器內和外通訊帶寬和延時的不同性,我們可以對其使用不同的一致性模型。例如第一層我們用強的一致性模型,而第二層我們則使用弱的一致性模型來減少同步開銷。 </p>
可移植性
輕量和可移植性是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能為大家做深度學習相關研究和應用帶來便利。也希望能與更多的開發者一起學習和進步。
擴展閱讀
- 此文大部分內容已經發表在NIPS LearningSys 2016上,paper link
- 本文只是對MXNet各個部件做了初步的介紹,更多文檔參見 MXNet/doc
- 本文實驗代碼均在 MXNet/example
來自:https://github.com/dmlc/mxnet/issues/797
本文由用戶
jopen 自行上傳分享,僅供網友學習交流。所有權歸原作者,若您的權利被侵害,請聯系管理員。
轉載本站原創文章,請注明出處,并保留原始鏈接、圖片水印。
本站是一個以用戶分享為主的開源技術平臺,歡迎各類分享!