SparkNET: 用Spark訓練深度神經網絡
SparkNet: Training Deep Network in Spark
這篇論文是 Berkeley 大學 Michael I. Jordan 組的 ICLR2016(under review) 的最新論文,有興趣可以看看原文和源碼: paper , github .
訓練深度神經網絡是一個非常耗時的過程,比如用卷積神經網絡去訓練一個目標識別任務需要好幾天來訓練。因此,充分利用集群的資源,加快訓練速度成了一個非常重要的領域。不過,當前非常熱門的批處理計算架構 (例如:MapReduce 和 Spark)都不是設計用來專門支持異步計算和現有的一些通信密集型的深度學習系統。
SparkNet 是基于Spark的深度神經網絡架構,
- 它提供了便捷的接口能夠去訪問Spark RDDs;
- 同時提供Scala接口去調用caffe;
- 還擁有一個輕量級的tensor 庫;
- 使用了一個簡單的并行機制來實現SGD的并行化,使得SparkNet能夠很好的適應集群的大小并且能夠容忍極高的通信延時;
- 它易于部署,并且不需要對參數進行調整;
- 它還能很好的兼容現有的caffe模型;
下面這張圖是SparkNet的架構:
從上圖可以看出,Master 向每個worker 分發任務之后,各個worker都單獨的使用Caffe(利用GPU)來進行訓練。每個worker完成任務之后,把參數傳回Master。論文用了5個節點的EC2集群,broadcast 和 collect 參數(每個worker幾百M),耗時20秒,而一個minibatch的計算時間是2秒。
Implementation
SparkNet 是建立在 Apache Spark 和Caffe深度學習庫的基礎之上的。SparkNet 用 Java 來訪問Caffe的數據,用Scala來訪問Caffe的參數,用ScalaBuff來使得Caffe網絡在運行時保持動態結構。SparkNet能夠兼容Caffe的一些模型定義文件,并且支持Caffe模型參數的載入。
下面簡單貼一下SparkNet的api和模型定義、模型訓練代碼。
并行化的SGD
為了讓模型能夠在帶寬受限的環境下也能運行得很好,論文提出了一種SGD的并行化機制使得最大幅度減小通信,這也是全文最大了亮點。這個方法也不是只針對SGD,實際上對Caffe的各種優化求解方法都有效。
在將SparkNet的并行化機制之前,先介紹一種Naive的并行機制。
Naive SGD Parallelization
Spark擁有一個master節點和一些worker節點。數據分散在各個worker中的。
在每一次的迭代中,Spark master節點都會通過broadcast(廣播)的方式,把模型參數傳到各個worker節點中。
各個worker節點在自己分到的部分數據,在同一個模型上跑一個minibatch 的SGD。
完成之后,各個worker把訓練的模型參數再發送回master,master將這些參數進行一個平均操作,作為新的(下一次迭代)的模型參數。
這是很多人都會采用的方法,看上去很對,不過它有一些缺陷。
Naive 并行化的缺陷
這個缺陷就是需要消耗太多的通信帶寬,因為每一次minibatch訓練都要broadcast 和 collect 一次,而這個過程特別消耗時間(20秒左右)。
令 N a ( b ) 表示,在batch-size為 b 的情況下,到達準確率 a 所需要的迭代次數。
令 C ( b ) 表示,在batch-size 為 b 的情況下,SGD訓練一個batch的訓練時間(約2秒)。
顯然,使用SGD達到準確率為 a 所需要的時間消耗是:
N a ( b ) C ( b )
假設有 K 個機器,通信(broadcast 和 collect)的時間為 S ,那么Naive 并行 SGD
的時間消耗就是:
N a ( b ) ( C ( b ) / K + S )
SparkNet 的并行化機制
基本上過程和Naive 并行化差不多。唯一的區別在于,各個worker節點每次不再只跑一個迭代,而是在自己分到的一個minibatch數據集上,迭代多次,迭代次數是一個固定值 τ 。
SparkNet的并行機制是分好幾個rounds來跑的。在每一個round中,每個機器都在batch size為 b 的數據集上跑 τ 次迭代。沒一個round結束,再把參數匯總到master進行平均等處理。
我們用 M a ( b , K , τ ) 表示達到準確率 a 所需要的 round 次數。
因此,SparkNet需要的時間消耗就是:
M a ( b , K , τ ) ? ( τ C ( b ) + S )
下面這張圖,很直觀的對比了Naive 并行機制跟 SparkNet 并行機制的區別:
Naive 并行機制:
SparkNet 并行機制:
論文還做了各種對比實驗,包括時間,準確率等。實驗模型采用AlexNet,數據集是ImageNet的子集(100類,每類1000張)。
假設 S = 0 ,那么 τ M a ( b , K , τ ) / N a ( b ) 就是SparkNet的加速倍數。論文通過改變 τ 和 K 得出了下面的表格(使準確率達到20%的耗時情況):
上面的表格還是體現了一些趨勢的:
(1). 看第一行,當 K = 1 ,因為只有一個worker節點,所以異步計算的 τ 這時并沒有起到什么作用,可以看到第一行基本的值基本都是接近1.
(2). 看最右邊這列,當 τ = 1 ,這其實就相當于是Naive 并行機制,只不過,Naive的batch是 b / K ,這里是 b . 這一列基本上是跟 K 成正比。
(3). 注意到每一行的值并不是從左到右一直遞增的。
當 S ! = 0 的時候,naive 跟 SparkNet 的耗時情況又是怎么樣的呢?作者又做了一些實驗。
可以看到,當 S 接近與0的時候(帶寬高),Naive會比SparkNet速度更快,但是,當 S 變大(帶寬受限),SparkNet的性能將超過Naive,并且可以看出,Naive受 S 變化劇烈, 而SparkNet相對平穩。
而作者實驗用EC2環境, S 大概是 20 秒,所以,顯然,SparkNet會比Naive好很多。
論文還做了一些事情,比如:
令 τ = 50 ,分別測試 K = 1、3、5、10時,準確率與時間的關系;
當 K = 5 ,分別測試 τ = 20、50、100、150時,準確率與時間的關系。
總結一下,這篇論文其實沒有太多復雜的創新(除了SGD并行化時候的一點小改進),不過我很期待后續的工作,同時也希望這個SparkNet能夠維護的越來越好。有時間的話,還是很想試試這個SparkNet的。
來自: http://dataunion.org/24000.html