MXnet：一個靈活、高效的異構分佈式系統機器學習庫。
摘要
MXnet是一個多語言機器學習（ML）庫，用於簡化ML算法的開發，特別是對於深度神經網絡。它嵌入在宿主語言中，將聲明性符號表達式與命令式張量計算混合在一起。它提供自動微分來推導梯度。MXnet具有計算和內存效率，可以在各種異構系統上運行，從移動設備到分佈式GPU集羣。
本文描述了MXNET的API設計和系統實現，並解釋瞭如何統一處理符號表達式和張量運算的嵌入。我們的初步實驗表明，在使用多個GPU機器的大規模深度神經網絡應用中，有着很好的結果。
1 引言
機器學習（ML）算法的規模和複雜性越來越大。幾乎所有最近的ImageNet挑戰[12]的獲獎者都使用具有很深層次的神經網絡，需要數十億次浮點運算來處理一個樣本。結構和計算複雜度的提高對ML系統的設計和實現提出了有趣的挑戰。大多數ML系統將域SPECI C語言（DSL）嵌入到宿主語言（例如Python、Lua、C++）中。可能的編程範例包括命令式的，在命令式的，用戶指定確切的“如何”計算需要執行，以及聲明式的，在聲明式的，用戶指定的重點。
關於“要做什麼”。命令式編程的例子包括numpy和matlab，而caffe、cxxnet等包在層定義上進行抽象並隱藏實際實現的內部工作。二者之間的分界線有時可能是模糊的，像theano和最近的張量流這樣的框架也可以看作兩者的混合物，它們聲明瞭一個計算圖，但圖中的計算必須是特定的。
與編程範式問題相關的是如何進行計算。執行可以是具體的，結果立即在同一線程上返回，或者異步或延遲執行，在將語句發佈到可用設備之前，首先將其收集並轉換爲數據流圖作爲中間表示。這兩個執行模型對如何發現固有的並行性有不同的含義。具體執行是限制性的（例如並行矩陣乘法），而異步/延遲執行則會自動識別數據流圖實例範圍內的所有並行性。
編程範式和執行模型的結合產生了一個很大的設計空間，其中一些空間比其他空間更有趣（也更有效）。事實上，我們的團隊和社區其他人一樣，已經集體完成了一些任務。例如，Minerva[14]將命令式編程與異步執行結合起來。雖然Theano採用聲明性方法，但它支持更全局的圖形感知優化。在嘌呤2[10]中採用了類似的學科。相反，cxnet採用聲明式編程（過度張量抽象）和具體的執行，類似於caffe[7]。表1給出了更多的例子。
| - | 命令式程序 | 聲明性程序 |
|---|---|---|
| 執行a = b+1 | 計算結果並將其存儲在a上，就像b的類型一樣。 | 返回計算圖；將數據綁定到b並稍後進行計算。 |
| 優勢 | 概念簡單，通常與宿主語言的內置數據結構、函數、調試器和第三方庫無縫工作。 | 在執行前獲取整個計算圖，有利於優化性能和內存利用率。也便於實現加載、保存和可視化等功能。 |
表1：比較域特定語言的命令式和聲明式。
表2：與其他流行的開放源碼ML庫比較
我們的合併新的努力導致了MXnet（或“混合網”），打算混合不同方法的優勢。聲明式編程在全局計算圖上提供了清晰的邊界，發現了更多的優化機會，而命令式編程提供了更多的靈活性。在深入學習的環境中，聲明式編程在神經網絡配置中指定計算結構時很有用，而命令式編程對於參數更新和交互式調試更爲自然。我們還努力將其嵌入到多個宿主語言中，包括C++、Python、R、Go和Julia。
儘管支持多種語言和不同編程範式的組合，我們還是能夠將執行融合到同一個後端引擎。引擎跟蹤計算圖和命令式操作之間的數據依賴性，並有效地聯合調度它們。我們積極減少內存佔用，儘可能執行就地更新和內存空間重用。最後，我們設計了一個緊湊的通信API，使得MXnet程序可以在多臺機器上運行，而不會發生任何變化。
與其他開放源碼ML系統相比，MXnet爲Torch7[3]、Theano[1]、Chainer[5]和Caffe[7]提供了超集編程接口，並支持更多系統，如GPU集羣。MXnet與TensorFlow[11]很接近，但可以另外嵌入命令式張量操作。MXNet是輕量級的，例如預測代碼T到單個50K行C++源代碼中，沒有其他依賴性，並且有更多的語言支持。更詳細的比較如表2所示。
2 編程接口
2.1 符號：聲明性符號表達式
MXNet 使用多輸出符號表達式，符號，聲明計算圖。符號由運算符組成，例如簡單的矩陣運算（例如“+”）或複雜的神經網絡層（例如卷積層）。一個運算符可以接受多個輸入變量，生成多個輸出變量，並具有內部狀態變量。變量可以是自由的，稍後我們可以用值綁定它，也可以是另一個符號的輸出。圖2顯示了通過鏈接一個變量（顯示輸入數據）和多個層操作符來構造多層感知符號。
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()
圖2:Julia中的符號表達式構造。
>>> import mxnet as mx
>>> a = mx.nd.ones((2, 3), mx.gpu())
>>> print (a * 2).asnumpy()
[[ 2. 2. 2.]
[ 2. 2. 2.]]
圖3:python中的ndarray接口
爲了評估一個符號，我們需要用數據綁定自由變量並聲明所需的輸出。除了評估（“forward”），符號還支持自動符號區分（“backward”）。還爲符號提供了其他功能，如加載、保存、內存估計和可視化。
2.2 NDArray：命令張量計算
MXnet提供了帶強制張量計算的ndarray，以彌補聲明性符號表達式和宿主語言之間的差距。圖3顯示了一個在GPU上執行矩陣常量乘法，然後按numpy.ndarray打印結果的示例。
ndarray抽象與符號聲明的執行無縫工作，我們可以將前者的命令張量計算與後者混合。例如，給出一個符號神經網絡和權重更新函數，例如w = w - ng，然後我們可以通過
while(1) { net.foward_backward(); net.w -= eta * net.g };
上述方法與使用單一但非常複雜的符號表達式的實現同樣有效。原因是MXnet使用了延遲的ndarray計算，後端引擎可以正確地解決兩者之間的數據依賴關係。
2.3 KVStore：設備上的數據同步
KVStore是一個分佈式的鍵值存儲，用於在多個設備上進行數據同步。它支持兩個原語：將一個鍵-值對從設備推到存儲，並從存儲中提取鍵的值。此外，用戶定義的更新程序可以指定如何合併推送的值。最後，通過一致性模型[8]控制模型發散。目前，我們支持順序和最終的一致性。
下面的例子通過數據並行實現了分佈式梯度下降。
while(1){ kv.pull(net.w); net.foward_backward(); kv.push(net.g); }
當權重更新函數註冊到kvstore時，每個工作人員重複地從該存儲中提取最新的權重，然後推出本地計算的梯度。
與單個聲明性程序相比，上面的混合實現具有相同的性能，因爲實際的數據推送和拉取是由延遲評估執行的，後者和其他的一樣由後端引擎調度。
2.4其他模塊
MXnet 附帶工具，可以將任意大小的示例打包到單個緊湊文件中，以方便順序搜索和隨機搜索。還提供了數據迭代器。數據預取和預處理是多線程的，減少了由於可能的遠程文件存儲讀取和/或圖像解碼和轉換而產生的開銷。
訓練模塊實現了常用的優化算法，如隨機梯度下降。它在給定的符號模塊和數據迭代器上訓練模型，如果提供了額外的kvstore，則可以選擇分佈式。
3 實施
3.1 計算圖
將一個綁定的符號表達式表示爲一個計算圖進行計算。圖4顯示了圖2中MLP符號向前和向後圖形的一部分。在評估之前，MXNET轉換圖形以優化效率並將內存分配給內部變量。
圖4：向前和向後的計算圖。
圖形優化： 我們將探討以下簡單的優化。我們首先注意到，只需要獲得綁定期間指定輸出所需的子圖。例如，在預測中，只需要前向圖，而從內部層提取特徵時，可以跳過最後一層。其次，可以將運算符分組爲單個運算符。例如，a x b +1 被單個BLAS或GPU調用替換。最後，我們手工實現了優化的“big”操作，如神經網絡中的一個層。
內存分配： 請注意，每個變量的生命週期，即創建和最後一次之間的時間段，將用於計算圖。因此，我們可以爲不相交的變量重用內存。然而，理想的分配策略需要 O(n2) 時間複雜性，其中n是變量數。
我們提出了兩種具有線性時間複雜性的啓發式算法。第一個調用inplace，模擬遍歷圖的過程，並保留到目前爲止尚未使用的依賴節點的引用計數器。如果計數器達到零，內存將被回收。第二個名爲co-share，它允許兩個節點共享一段內存，前提是它們不能並行運行。探索共同共享會施加一個額外的依賴約束。特別是，每次調度時，在圖中的掛起路徑中，我們都會確定最長的路徑並執行所需的內存分配。
3.2 依賴引擎
在MXnet中，每個源單元，包括ndarray、隨機數生成器和時間空間，都用一個唯一的標記註冊到引擎中。然後，任何操作（如矩陣操作或數據通信）都會被推送到引擎中，並指定所需的資源標記。如果解決了依賴關係，引擎會連續調度推送的操作以供執行。由於通常存在多個計算資源，如CPU、GPU和內存/PCIe總線，因此該引擎使用多個線程來調度操作，以獲得更好的資源利用率和並行化。
與大多數數據流引擎不同[14]，我們的引擎將突變操作作爲一個現有的資源單元進行跟蹤。也就是說，我們的支持一個操作除了讀取之外還將寫入的標記的特定化。這使得可以像numpy和其他張量庫那樣調度數組突變。它還可以通過將參數更新表示爲參數數組的變化，更容易地實現參數的內存重用。它還使一些特殊操作的調度變得更容易。例如，當生成具有相同隨機種子的兩個隨機數時，我們可以通知引擎它們將寫入種子，這樣它們就不應該並行執行。這有助於再現性。
3.3 數據通信
我們基於參數服務器[8，9，4]實現了kvstore（圖5）。它與以前的工作在兩個方面有所不同：第一，我們使用引擎來調度kvstore操作和管理數據一致性。該策略不僅使數據同步與計算無縫結合，而且大大簡化了實現過程。第二，我們採用兩級結構，一級服務器管理單臺機器內設備之間的數據同步，二級服務器管理機器間的同步。來自一級服務器的出站數據可以聚合，從而減少帶寬需求；機器內和機器間的同步可以使用不同的一致性模型（例如，內部是連續的，最終是內部的）。
圖5：通信。
4 評價
原始性能： 我們在流行的“convnet基準”上最能將mxnet與torch7、caffe和tensorflow進行比較[2]。所有這些系統都是用CUDA7.5和CUDNN3編譯的，除了TensorFlow，它只支持CUDA7.0和CUDNN2。我們對所有網絡使用批量大小32，並在單個Nvidia GTX 980卡上運行實驗。結果如圖6所示。正如預期的那樣，MXnet的性能與Torch7和Caffe相似，因爲大多數計算都花費在CUDA/CUDNN內核上。TensorFlow總是慢2倍，這可能是因爲它使用了較低的CUDN版本。
圖6：將MXnet與其他MXnet在單個前向後退性能上進行比較。
圖7:mxnet在各種分配策略下的內部內存使用情況，只針對批大小爲64的正向（左）和正向（右）。
內存使用情況： 圖7顯示了除輸出外的內部變量的內存使用情況。可以看出，“inplace”和“co-share”都可以有效地減少內存佔用。在模型訓練過程中，將它們組合在一起可使所有網絡減少2倍，並進一步提高到4倍以進行模型預測。例如，即使是最昂貴的VGG網絡，培訓也需要不到16MB的額外空間。
可擴展性： 我們在AmazonEC2G2.8x實例上運行了這個實驗，每個實例都附帶了四個Nvidia GK104 GPU和10G以太網。我們在由130萬張圖像和1000個類組成的ILSVRC12數據集[13]上使用批標準化[6]對Googlenet進行培訓。我們將學習率固定爲：05，動量固定爲：9，重量衰減爲0:05，並在一批中爲每個GPU提供36個圖像。
收斂結果如圖8所示。可以看出，與單臺機器相比，分佈式訓練在開始時收斂較慢，但在10次數據傳遞之後表現更好。在一臺機器和10臺機器上，數據傳遞的平均成本分別爲14K和1.4K秒。因此，本實驗揭示了超線性加速。
圖8：1臺和10臺機器上的ILSVRC12數據集上的Googlenet進度。
5 結論
MXnet是一個機器學習庫，將符號表達式與張量計算結合起來，以最大限度地提高效率和靈活性。它是輕量級的，嵌入多種主機語言，並且可以在分佈式設置中運行。實驗結果令人鼓舞。雖然我們繼續探索新的設計選擇，但我們相信它已經有利於相關的研究團體。這些代碼可在 http://dmlc.io 上找到。
文獻
- [1] Fr´ed´eric Bastien, Pascal Lamblin, Razvan Pascanu, James Bergstra, Ian Goodfellow, Arnaud
Bergeron, Nicolas Bouchard, David Warde-Farley, and Yoshua Bengio. Theano: new features
and speed improvements. arXiv preprint arXiv:1211.5590, 2012. - [2] Soumith Chintala. Easy benchmarking of all public open-source implementations of convnets,
- [3] Ronan Collobert, Koray Kavukcuoglu, and Cl´ ement Farabet. Torch7: A matlab-like envi-
ronment for machine learning. In BigLearn, NIPS Workshop, number EPFL-CONF-192376,
- [4] J. Dean, G. Corrado, R. Monga, K. Chen, M. Devin, Q. Le, M. Mao, M. Ranzato, A. Senior, P. Tucker, K. Yang, and A. Ng. Large scale distributed deep networks. In Neural Information Processing Systems, 2012.
- [5] Chainer Developers. Chainer: A powerful, flexible, and intuitive framework of neural networks, 2015. http://chainer.org/.
- [6] Sergey Ioffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. arXiv preprint arXiv:1502.03167, 2015.
- [7] Yangqing Jia, Evan Shelhamer, Jeff Donahue, Sergey Karayev, Jonathan Long, Ross Girshick, Sergio Guadarrama, and Trevor Darrell. Caffe: Convolutional architecture for fast feature embedding. In Proceedings of the ACM International Conference on Multimedia, pages 675–678. ACM, 2014.
- [8] M. Li, D. G. Andersen, J. Park, A. J. Smola, A. Amhed, V. Josifovski, J. Long, E. Shekita, and B. Y. Su. Scaling distributed machine learning with the parameter server. In OSDI, 2014.
- [9] M. Li, D. G. Andersen, A. J. Smola, and K. Yu. Communication efficient distributed machine learning with the parameter server. In Neural Information Processing Systems, 2014.
- [10] Min Lin, Shuo Li, Xuan Luo, and Shuicheng Yan. Purine: A bi-graph based deep learning framework. arXiv preprint arXiv:1412.6249, 2014.
- [11] Abadi Martn, Ashish Agarwal, Paul Barham, Eugene Brevdo, Zhifeng Chen, Craig Citro, Greg Corrado, Andy Davis, Jeffrey Dean, Matthieu Devin, Sanjay Ghemawat, Ian Goodfellow, Andrew Harp, Geoffrey Irving, Michael Isard, Yangqing Jia, Rafal Jozefowicz, Lukasz Kaiser, Manjunath Kudlur, Josh Levenberg, Dan Mane, Rajat Monga, Sherry Moore, Derek Murray, Chris Olah,Mike Schuster, Jonathon Shlens, Benoit Steiner, Ilya Sutskever, Kunal Talwar, Paul Tucker, Vincent Vanhoucke, Vijay Vasudevan, Fernanda Viegas, Oriol Vinyals, Pete Warden, Martin Wattenberg, Martin Wicke, Yuan Yu, and Xiaoqiang Zheng. Tensorflow: Large-scale machine learning on heterogeneous systems. 2015.
- [12] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei. ImageNet Large Scale Visual Recognition Challenge. International Journal of Computer Vision (IJCV), 115(3):211–252, 2015.
- [13] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, et al. Imagenet large scale visual recognition challenge. International Journal of Computer Vision, pages 1–42, 2014.
- [14] Minjie Wang, Tianjun Xiao, Jianpeng Li, Jiaxing Zhang, Chuntao Hong, and Zheng Zhang. Minerva: A scalable and highly efficient training platform for deep learning, 2014.