1. 程式人生 > >Caffe、TensorFlow、MXnet三個開源庫對比

Caffe、TensorFlow、MXnet三個開源庫對比

from http://chenrudan.github.io/blog/2015/11/18/comparethreeopenlib.html#0-tsina-1-2654-397232819ff9a47a7b7e80a40613cfe1

最近Google開源了他們內部使用的深度學習框架TensorFlow[1],結合之前開源的MXNet[2]和Caffe[3],對三個開源庫做了一些討論,其中只有Caffe比較仔細的看過原始碼,其他的兩個庫僅閱讀官方文件和一些研究者的評論部落格有感,本文首先對三個庫有個整體的比較,再針對一些三者設計的不同資料結構、計算方式、gpu的選擇方式等方面做了比較詳細的討論。表格1是三者的一些基本情況的記錄和比較。其中示例指的是官方給出的example是否易讀易理解,因為TensorFlow直接安裝python包,所以一開始沒有去下原始碼,從文件中找example不如另外兩個下原始碼直接。實際上TensorFlow更加像一套獨立的python介面,它不止能夠完成CNN/RNN的功能,還見到過有人用它做Kmeans聚類。這個表主觀因素比較明顯,僅供參考。

庫名稱 開發語言 支援介面 安裝難度(ubuntu) 文件風格 示例 支援模型 上手難易
Caffe c++/cuda c++/python/matlab *** * *** CNN **
MXNet c++/cuda python/R/Julia ** *** ** CNN/RNN *
TensorFlow c++/cuda/python c++/python * ** * CNN/RNN/… ***
  • 安裝難度: *(簡單) --> ***(複雜)
  • 文件風格: *(一般) --> ***(好看、全面)
  • 示例: *(給的少) --> ***(給的多、全)
  • 上手難易: *(易) --> ***(難)

1.基本資料結構

庫名稱 資料結構名稱 設計方式
Caffe Blob 儲存的資料可以看成N維的c陣列,有(n,k,h,w)四個維數,一個blob裡面有兩塊資料空間儲存前向和後向求導資料
MXNet NDArray 提供cpu/gpu的矩陣和向量計算,能夠自動並行
TensorFlow tensor 相當於N維的array或者list,維數可變,資料型別一旦定義不能改變

caffe的資料儲存類blob,當把資料可以看成是一個N維的c陣列,它們的儲存空間連續。例如儲存圖片是4維(num, channel, height, width),變數(n,k,h,w)在陣列中儲存位置為((n*K+k)*H+h)*W+w。blob有以下三個特徵[4]:

  • 兩塊資料,一個是原始data,一個是求導值diff
  • 兩種記憶體分配方式,一種是分配在cpu上,一種是分配在gpu上,通過字首cpu、gpu來區分
  • 兩種訪問方式,一種是不能改變資料,一種能改變資料

Caffe最讓我覺得精妙的地方在於一個blob儲存前向和後向的資料。雖然就程式碼本身而言,前向資料是因為輸入資料不同而改變,後向求導是因為求導不同而改變,根據SRP原則,在一個類裡面因為兩個原因而改變了資料這種是不合適的設計。但是從邏輯層面,前向資料的改變引起了反向求導的不同,它們實際上是一起在改變,本身應該是一個整體。所以我很喜歡這個設計,雖然基本上其他框架中都是將兩個資料給分離出來,caffe2也不知是否保留。

MXNet的NDArray類似numpy.ndarray,也支援把資料分配在gpu或者cpu上進行運算。但是與numpy和caffe不同的是,當在操作NDArray,它能自動的將需要執行的資料分配到多臺gpu和cpu上進行計算,從而完成高速並行。在呼叫者的眼中程式碼可能只是一個單執行緒的,資料只是分配到了一塊記憶體中,但是背後執行的過程實際上是並行的。將指令(加減等)放入中間引擎,然後引擎來評估哪些資料有依賴關係,哪些能並行處理。定義好資料之後將它繫結到網路中就能處理它了。

TensorFlow的tensor,它相當於N維的array或者list,與MXNet類似,都是採用了以python呼叫的形式展現出來。某個定義好的tensor的資料型別是不變的,但是維數可以動態改變。用tensor rank和TensorShape來表示它的維數(例如rank為2可以看成矩陣,rank為1可以看成向量)。tensor是個比較中規中矩的型別。唯一特別的地方在於在TensorFlow構成的網路中,tensor是唯一能夠傳遞的型別,而類似於array、list這種不能當成輸入。

值得一提的是cuda-convnet採用的資料結構是NVMatrix,NV表示資料分配在gpu上,即將所有變數都當成矩陣來處理,它只有兩維,它算是最早用cuda實現的深度學習框架,而上面三種框架都採用了多維可變維的思想,這種可變維在用矩陣做卷積運算的時候是很有效的。

2.網路實現方式

Caffe是典型的功能(過程)計算方式,它首先按照每一個大功能(視覺化、損失函式、非線性激勵、資料層)將功能分類並針對部分功能實現相應的父類,再將具體的功能實現成子類,或者直接繼承Layer類,從而形成了XXXLayer的形式。然後將不同的layer組合起來就成了net。

1

圖1 caffe的網路結構

MXNet是符號計算和過程計算混合[5],它設計了Symbol大類,提供了很多符號運算的介面,每個symbol定義了對資料進行怎樣的處理,symbol只是定義處理的方式,這步還並未真正的執行運算。其中一個需要注意的是symbol裡面有Variable,它作為承載資料的符號,定義了需要傳遞什麼樣的資料給某個Variable,並在後續的操作中將資料繫結到Variable上。下面的程式碼是一個使用示例,它實現了將激勵函式連線到前面定義好的net後面,並給出了這一個symbol的名字和激勵函式型別,從而構造出net。下圖左邊部分是定義symbol的合集,中間將資料繫結到Variable上之後變成了右邊真正的執行流程圖。

net = mx.symbol.Activation(data=net, name='relu1', act_type="relu")

1

圖2 MXNet的網路結構

TensorFlow選擇的是符號計算方式,它的程式分為計算構造階段和執行階段,構造階段是構造出computation graph,computation graph就是包含一系列符號操作Operation和Tensor資料物件的流程圖,跟mxnet的symbol類似,它定義好了如何進行計算(加減乘除等)、資料通過不同計算的順序(也就是flow,資料在符號操作之間流動的感覺)。但是暫時並不讀取輸入來計算獲得輸出,而是由後面的執行階段啟動session的run來執行已經定義好的graph。這樣的方式跟mxnet很相似,應該都是借鑑了theano的想法。其中TensorFlow還引入了Variable型別,它不像mxnet的Variable屬於symbol(tf的operation類似mxnet的symbol),而是一個單獨的型別,主要作用是儲存網路權重引數,從而能夠在執行過程中動態改變。tf將每一個操作抽象成了一個符號Operation,它能夠讀取0個或者多個Tensor物件作為輸入(輸出),操作內容包括基本的數學運算、支援reduce、segment(對tensor中部分進行運算。例如tensor長度為10,可以同時計算前5個,中間2個,後面三個的和)、對image的resize、pad、crop、filpping、transposing等。tf沒有像mxnet那樣給出很好的圖形解釋或者例項(可能因為我沒找到。。),按照自己的理解畫了一部分流程圖。有點疑惑的是,為什麼要設計Variable,tf給出的一個alexnet的example原始碼中,輸入資料和權重都設定成了Variable,每一層的輸出並未直接定義,按照tf的說法,只有tensor型別能夠在網路中傳遞,輸出的型別應該是tensor,但是由於輸入和權重改變了,輸出應該也在隨著改變,既然如此,為何不只設計一個tensor,讓tensor也能動態改變。

1

圖3 TensorFlow的computation graph

就設計而言,TensorFlow相對於其他兩個更像是一種通用的機器學習框架,而不是隻針對cnn或rnn,但就現在的效能而言,tf的速度比很多開源框架都要差一點[6]。

3.分散式訓練

Caffe和TensorFlow沒有給出分散式的版本,MXNet提供了多機分散式,因而前兩者只有如何控制使用多gpu。Caffe通過直接在執行指令後面加上-gpu 0,1來表示呼叫兩個gpu0和1,只實現了資料並行,也就是在不同的gpu上執行相同網路和不同資料,caffe會例項化多個solver和net讓每次處理的batch_size加倍。TensorFlow則能夠自己定義某個操作執行在哪個gpu上,通過呼叫with tf.device(‘/gpu:2’)表示接下來的操作要在gpu2上處理,它也是資料並行。MXNet通過執行指令碼時指定多機節點個數來確定在幾臺主機上執行,也是資料並行。MXNet的多gpu分配和它們之間資料同步是通過MXNet的資料同步控制KVStore來完成的。

KVStore的使用首先要建立一個kv空間,這個空間用來在不同gpu不同主機間分享資料,最基本的操作是push和pull,push是把資料放入這個空間,pull是從這個空間取資料。這個空間內儲存的是key-value([int, NDArray]),在push/pull的時候來指定到哪個key。下面的程式碼將不同的裝置上分配的b[i]通過key3在kv空間累加再輸出到a,從而完成了對多gpu的處理。這個是個非常棒的設計,提供了很大的自由度,並且為開發者減少了控制底層資料傳輸的麻煩。

gpus = [mx.gpu(i) for i in range(4)]	
b = [mx.nd.ones(shape, gpu) for gpu in gpus]
kv.push(3, b)
kv.pull(3, out = a)

之前有看過一篇論文,如何將卷積網路放在多gpu上訓練,論文中有兩種方法,一種是常用的資料並行,另一種是模型並行。模型並行指的是將一個完整的網路切分成不同塊放在不同gpu上執行,每個gpu可能只處理某一張圖的四分之一。採用模型並行很大程度上是因為視訊記憶體不夠放不下整個網路的資料,而現在gpu的功能效能提高,一個gpu已經能夠很好的解決視訊記憶體不夠的問題,再加上模型並行會有額外的通訊開銷,因此開源框架採用了資料並行,用來提高並行度。

4.小結

上面針對三個框架的不同方面進行了一些分析與比較,可以看出TensorFlow和MXNet有一些相似的地方,都是想做成更加通用的深度學習框架,貌似caffe2也會採用符號計算[5],說明以後的框架會更加的偏向通用性和高效,個人最喜歡的是caffe,也仿造它和cuda-convnet的結構寫過卷積網路,如果是想提高程式設計能力可以多看看這兩個框架的原始碼。而MXNet給人的感覺是非常用心,更加註重高效,文件也非常的詳細,不僅上手很容易,運用也非常的靈活。TensorFlow則是功能很齊全,能夠搭建的網路更豐富而不是像caffe僅僅侷限在CNN。總之框架都是各有千秋,如何選擇也僅憑個人的喜好,然而google這個大殺器一出現引起的關注度還是最大的,雖然現在單機效能還不夠好,但是看著長長的開發人員名單,也只能說大牛多就是任性。

參考: