乾貨分享：深度學習框架技術剖析

GIF

摘要：深度學習框架正在快速演化，各大公司都推出了自己的框架，TensorFlow, PyTorch, Caffe2, MXNet, PaddlePaddle，大大推動了深度學習的發展，同時也讓用戶有目不暇接無所適從之感。我們認為，深度學習框架用戶有必要去了解深度學習框架的一些基本原理，這有助於我們用好「框架」這個工具，也有助於根據自身需要去選擇合適的框架。

作為框架的開發者，我們發現，雖然框架多種多樣，但框架核心技術正呈現收斂的態勢，經過幾年的發展，在深度學習框架開發者眼裡出現一些「共識」，或所謂「最佳實踐」，幾乎所有框架都去擁抱了這樣技術選型，在架構和技術選擇上趨同。另一方面，也有一些技術在框架開發者眼裡屬於舉棋不定或無計可施的狀態。這次報告會對已經收斂的技術（「最佳實踐」）做一個梳理，讀者會發現，開發一個深度學習框架沒有那麼難。本報告也會簡要討論目前框架未解決的難題，讀者也會發現，開發一個超越已有技術的框架有很難的問題。最後，我們會從框架開發者的視角去對主流深度學習框架做一句話點評，供用戶在做技術選型時參考。

註：首先介紹深度學習框架的背景，然後介紹深度學習框架開發中已經收斂的技術和仍未解決的問題，其次點評主流深度學習框架，最後對2018年的深度學習框架技術發展做出展望。

註：在進入正文前，讓我們首先聲明一些前提，如果這些前提不成立，那麼「深度學習框架」就沒那麼重要。本文僅對第四點做一些闡述，用軟體實現深度學習演算法加速，可分微觀和宏觀兩個層次。微觀層次主要聚焦在單個設備或晶元上的代碼優化，設備廠商通常會工作在這個層次，他們會提供高性能的庫，譬如x86 或arm CPU上MKL, OpenBlas，Nvidia GPU 上的CuBlas, Cudnn等，在大部分稠密計算場景都能貼近設備的理論性能，優化空間不大（當然在終端設備等低功耗場景有很多發揮空間）。宏觀層次，主要是多設備和多計算節點層面的優化，這要靠分散式框架的支持，是計算力推向更高水平的關鍵，在性能上仍有巨大優化空間。最後，深度學習軟體要解決編程不夠快（程序員的效率）和程序運行不夠快（計算機的效率）兩個痛點，「兩個痛點的描述」出自尼克著《人工智慧簡史》里的一句話。

註：神經網路由若干層次構成，每個層次通常都可以表示成對矩陣的處理，這種稠密計算形式特別適合使用高度並行的硬體加速（GPU, TPU等）。

註：限於硬體製造工藝水平，單個設備（GPU, TPU) 不可能無限大，而工業級應用對計算力的渴求是無止境的，此時就必須用高速互聯的多個設備協同來完成大規模計算。上圖展示了GPU集群和TPU集群，在這種配置里，通常是CPU和GPU （或TPU) 一塊兒工作，CPU 負責任務的調度和管理，而GPU 負責實現稠密計算，這就是經常說的異構計算（Heterogenous computing）。

註：「硬體越快，軟體越難」這個觀點分享過多次。簡要說：自上而下看，深度學習模型訓練通常使用隨機梯度下降（SGD)演算法，是更接近流式計算的一種負載：每處理一小片數據，就引起系統內部狀態的變化；自下而上看，深度學習廣泛採用異構計算技術，GPU 此類的設備吞吐率非常高，是CPU的10倍以上，意味著同樣大小的計算任務，GPU可以更快完成。從小粒度和快設備兩方面看，深度學習訓練中的計算任務粒度非常小，通常是數十毫秒到百毫秒級別。但是，設備互聯帶寬並沒有實質改進，譬如同機內部PCIe或多機互聯使用的高速乙太網或Infiniband的傳輸帶寬要低於GPU內部數據帶寬一兩個數量級。以上因素給分散式軟體框架帶來巨大壓力，如果處理不好，就會造成設備利用率低，整體系統性能差的後果。打個比方，雖然高鐵要比普通的列車開起來快很多，但如果讓高鐵每個車站都挺兩分鐘，那麼高鐵並不會比普通列車快多少。

註：軟體層和硬體層都是屬於「計算力」範疇，軟體層扮演了傳統意義上操作系統（OS，如Windows, Linux），或者互聯網時代瀏覽器，或者移動互聯網時代Android, IOS，或者大數據時代Hadoop的角色，是上層應用的入口。同時軟體生態又定義了底層硬體的角色，會影響底層硬體的發展和命運。

註：我們首先介紹深度學習框架中已經收斂的技術，理解了這些原理，每個人應該能開發出一個自己的深度學習框架。

註：在進入技術細節之前，讓我們先來理解兩個很重要的概念：控制流（Control flow） 和數據流（Data flow)，這倆概念事關後續一些關鍵的技術選擇。以a = x + y; b = a * a; c = 4 - a; 這樣一段簡單的程序為例，有兩種編程模式，一種是以C語言為代表的命令式編程（imperative programming），語句的排列順序隱式的刻畫了程序的執行順序（左圖中虛線箭頭表示執行順序），有哪些語句可以並行執行，並不太明確，如果要在多個線程中執行這幾條語句，為了防止出現多個線程間的讀寫衝突，可能要使用鎖 （lock）等技術來保護某一個變數（某一段內存）防止出現data race。另一種編程模式是以Lisp為代表的函數式編程（functional programming），程序用一系列表達式來刻畫，程序的執行不是按表達式的聲明順序來執行，而是從表達式中挖掘出各個表達式之間的數據依賴關係，把數據依賴關係用一個有向無環圖來表示，圖顯式刻畫了哪些表達式必須在另一些表達式之前求值，或者哪些表達式之間不存在依賴關係，可以並行執行。在並行和並發越來越多的世界，functional programming 和數據流程序正在越來越受重視。

註：數據流模型一般表示成有向無環圖（Directed acyclic graph, DAG）。譬如上一頁的a = x + y; b = a * a; c = 4 - a; 三個表達式可以表示成這樣一張圖，圓圈表示數據，方塊表示運算元。運算元之間靠數據的生產和消費關係產生依賴。數據流模型的優勢主要包括兩方面：（1） 表示上的好處，顯式描述了程序中的所有並行機會；（2）實現上的好處，數據流的執行引擎可以用很簡單的方法支持並發和並行，而在控制流程序中對並發和並行的支持就要複雜的多。

註：比較早的框架Caffe 通過Layer 這種抽象，運算和數據是放在一起的。TensorFlow 出現後，有向無環圖中兩個最基本的元素，操作符（運算）和張量（數據）是分開表示的，這種抽象模式也被其它框架所採用。具體地，Operator 一般是運算的描述，Kernel 是運算的具體實現，在實現上還要考慮操作符粒度的問題，理論上如果支持了最基本的加減乘除就可以通過圖計算自動支持更加複雜的運算（如一些神經網路層次的計算），但粒度太細就對編譯器要求特別高，當前編譯器生成的代碼不一定能超過工程師手工優化的代碼，因此在多數框架里都是直接支持一些粗粒度的操作符，譬如卷積操作符，矩陣乘操作符等（值得注意的是TensorFlow XLA, TVM 在細粒度操作符描述的圖優化方面有一些好的實踐）。對於張量計算的支持，業界也有很多技巧，譬如一般使用C++模板元編程來實現，藉助於編譯時計算來提高運行時的效率，TensorFlow 等框架一般基於Eigen庫，MXNet 使用自己開發的Mshadow。

註：autograd 已經成為深度學習框架的標配。有了autograd，用戶寫程序時，只需要描述前向計算是怎麼做的，後向計算過程由系統自己推導完成。autograd通過導數的鏈式法則實現，逆拓撲序搭建反向計算圖。需要注意兩點：（1）後向計算過程可能會依賴於前向計算產生的中間數據，所以前向計算的中間數據可能要一直保持到對應的後向計算完成才能釋放，否則就需要在後向計算時重新進行前向計算。（2）如果前向計算過程有多個操作符消費了同一個數據，後向計算圖時就需要把這幾個操作符對應的梯度運算元上傳過來的誤差信號進行累加。上面的示意圖來自陳天奇在華盛頓大學一門《Deep learning systems》課程的課件，感興趣的讀者可以去課程網站獲取更多資料。

註：給定用戶輸入的DAG （稱之為邏輯圖，logical graph）, 系統一般會利用編譯器技術對圖進行優化重寫，上圖羅列的一些優化技巧就不一一展開解釋了。經過優化最終送到執行引擎執行的圖叫物理圖（physical graph），物理圖可能和輸入的邏輯圖已經截然不同了。在TensorFlow, PyTorch, MXNet, Caffe2 中都可以看到這些常見的優化辦法。

註：執行引擎是深度學習引擎的核心，基本原理是按拓撲序去執行運算元/操作符，以上圖為例，剛開始，乘法和減法運算無法執行，因為它們依賴的一個數據a還沒有生成，引擎首先執行輸入數據已經就緒的操作符，即加法，當加法完成後，執行引擎會從DAG 中刪掉已經執行的節點，然後發現乘法和減法的執行條件已經滿足了，再去執行乘法和減法。事實上，當前所有大數據處理引擎的內核都是用這個原理實現的。在深度學習框架里，需要注意調度器是在CPU上執行的，而操作符的真實運算是在GPU上完成的，為了高效協調CPU和GPU之間的工作，在具體實現上有一些技巧。感興趣的讀者可以觀摩TensorFlow, MXNet, Caffe2 的執行引擎，也許你會發現有更好的實現辦法。

註：從執行效率考慮，深度學習框架底層一般基於C++開發，從易用性角度出發，也同時提供Python前端便於數據科學家使用。上圖來自李飛飛教授在斯坦福的cs231n課程課件，展示了Numpy，TensorFlow和PyTorch 對同一個簡單神經網路訓練過程的實現。最左側是Numpy 代碼，它的第一個特色是imperative programming，是即時求值（eager evaluation），運行完b = a + z 這條語句，b的結果就出來了；第二個特色是沒有autograd，所以用戶不僅要寫前向計算的代碼，還要把後向梯度下降的代碼寫出來，而TensorFlow和PyTorch都支持了autograd，在代碼中只需要寫前向計算的過程，而後向計算過程是系統自動構建的。TensorFlow 和PyTorch的區別則是，前者是lazy evaluation，後者是eager evaluation。在TensorFlow 中，a = x + y; b = a + z 只是一些表達式，構建了一個數據流圖，在執行sess.run 時刻才會被真正執行，而且執行順序不一定和表達式聲明順序一致。在PyTorch中，和Numpy原理類似，每條語句都是立刻執行，而且按照語句的排列順序執行。看上去，PyTorch的代碼的確更加簡潔，後文會進一步討論延遲求值和即時求值的問題。

註：深度學習框架不只要解決易用性和高效性，還要方便部署運維。當前主流的深度學習框架都是基於檢查點機制實現容錯，Fail fast and warm start。深度學習框架還需要和上下游的開源工具鏈有機銜接，譬如分散式數據存儲和數據預處理依靠Hadoop或者Spark。部署運維，現在比較推崇基於Docker和Kubernetes相結合的方案。用戶有時需要在多個框架之間切換，隨著ONNX標準的推出，也大大便利了各種框架間的遷移，譬如使用PyTorch 描述或訓練的模型可以按ONNX規範導出，並被Caffe2框架使用。除了解決訓練問題，深度學習框架還便於上線部署，為此TensorFlow推出了單獨的serving模塊。

註：下面我們探討一些當前框架開發者還舉棋不定或一籌莫展的技術問題。

註：Define-and-run 和 Define-by-run 近期關注度比較高，PyTorch 靠Define-by-run 這個特性吸引了很多開發者。這個技術問題還有其它等價的描述，譬如define-and-run，基本上和lazy evaluation, 或 declarative programming, data flow 是一回事，通常意味著效率高。define-by-run 則基本上和eager evaluation， imperative programming 或control flow 是一回事，通常意味著靈活性。最近，很多框架在積極的推進支持define-by-run的運行模式，譬如TensorFlow 增加了eager evaluation，MXNet 推出了gluon 介面，PaddlePaddle Fluid 也是一種imperative programming的用法。那這兩種技術選擇到底是怎麼回事呢？ 我們認為：（1）Imperative programming 只不過是大部分程序員更加熟悉的編程方式，實現一個imperative programming的深度學習框架要比實現一個declarative programming的框架簡單（最簡只須實現autograd，複雜點要支持JIT）。傳統的lazy evaluation框架增加imperative programming介面可以做到和PyTorch 完全一樣的用戶體驗，只不過要費些周章。（2）只要declarative programming 解決了調試困難等問題，就是對用戶更友好的一種編程模式，用戶只要在寫程序時描述 what，而不需要關心 how，底層細節對用戶透明，這是現代變編程語言的發展趨勢。（3）並行和並發代表著未來的趨勢，那麼數據流（聲明式編程，函數式編程）代表著未來，data flow 模型在任務描述和系統執行引擎的簡潔性上都有天然優勢。

註：並行計算可以擴大處理任務的規模，也可以加速計算。具體到深度學習框架，總體情況是：數據並行已經得到解決，無論是哪個框架都能把適合數據並行的任務做到接近理想加速比，譬如計算機視覺領域各種CNN模型；數據並行和流水線並行不支持或支持的不好，導致在某些分散式訓練場景，硬體利用率過低，訓練周期過長。在集合通信（Collective communication operation）上有基於參數伺服器的，MXNet, PaddlePaddle，TensorFlow，也有基於MPI（或類MPI）的，譬如Caffe2。TensorFlow 在宏觀架構上還區分了Client, Master, Worker節點，在重構版的PaddlePaddle也使用了類似的架構。

註：現有框架都能良好支持數據並行。原本限制數據並行的一個問題是隨機梯度下降演算法依賴的mini-batch不能太大，太大的mini-batch 演算法不收斂，這就導致即使有大量的並行設備，也發揮不出來威力。近期有一系列成果攻克了這個問題，譬如Facebook 推出的一個小時訓練ImageNet，以及尤洋做的一系列工作，可以把mini-batch 推廣到32K, 保證演算法仍然收斂，這就能充分發揮數據並行的優勢。

註：模型並行本身實現複雜度不是特別高，主要困難在於有的場景適合數據並行，有的場景適合數據並行，有的場景同時需要模型並行和數據並行，這就需要根據實際情況正確的對數據重新組織（分裂，合併）和路由（把數據正確的發送到目的地，scatter或broadcast）。再有就是，當數據路由比較複雜時，就很難高效的支持，就需要流水線並行的支持。

註：當神經網路模型或中間隱狀態體積很大時，譬如超過一張顯卡顯存的容量，除了使用模型並行，還可以使用流水線並行。流水線並行有點像接力比賽，上圖展示了一個簡單的例子，第一個GPU 做完第一層的計算後，把結果傳遞給第二塊GPU，第二塊GPU 完成中間四層的計算之後，把結果傳遞給第三塊GPU完成計算。通常訓練深度學習模型時，存在多個階段，譬如把數據從磁碟載入到主存，把數據從主存搬運到GPU, GPU 完成一個階段的計算之後，可能還需要把數據通過網路傳送到另一台機器。在這樣多階段任務中，流水線並行對於系統性能至關重要。可以注意到，大部分框架在IO階段會做流水線並行，而在多個計算階段，或者計算與通信階段之間都沒有相應支持。基於現有框架去支持模型並行，流水線並行還非常困難。

註：下面分享一些我們對各種框架的理解和判斷，如果觀點有偏差，敬請理解，歡迎批評指正。

註：以上框架用戶比較多，開發團隊技術實力雄厚。深度學習框架的先驅Theano 已停止更新，它的autograd機制已經被這些現代框架所吸收；我們沒有研究微軟開發的CNTK；Intel 收購的Nervana 在開發一個新的框架NGraph，也值得關注，不過它目前主要聚焦在單設備優化；DMLC的NVVM和TVM 放在MXNet內；有一個來自日本研究人員的框架Chainer也比較有特色，Python 前端非常清爽。

註：TensorFlow 是系統完整度最高的，支持training和inference （serving），支持圖像常用的CNN，也支持自然語言和語音常用的RNN/LSTM， 還有移動端的TensorFlow Lite，支持lazy execution也支持eager execution，社區生態最強大，Google在深度學習演算法和應用方向的研究也是冠絕天下（參見Jeff Dean 寫的 Google Brain 2017 年度回顧 https://zhuanlan.zhihu.com/p/32905123 ）深度學習領域很多新的研究成果都是以TensorFlow代碼發布的。但TensorFlow的性能也是廣受詬病，我們不大清楚TensorFlow在Google內部是不是性能卓越，在外部用戶那裡經常能聽到用戶抱怨TF 慢，在學界和工業界很多Benchmark里，大家也喜歡拿TensorFlow做baseline，譬如CMU Eric Xing教授團隊發表的Poseidon 論文，基於TF 做了一系列優化之後，性能提高非常顯著；Uber 團隊改造了TensorFlow的分散式實現後（把PS換成MPI), CNN 數據並行場景還能提高一倍的性能 （見Uber Horovod https://github.com/uber/horovod) 。從框架開發者的角度看，我們以為TensorFlow要解決性能問題，須有壯士斷腕的決心，推翻一系列原有設計和實現。最後，TensorFlow畢竟是功能最完整的框架，如果要訓練大規模RNN/LSTM，目前只能選擇它，儘管要忍受一下很長的訓練周期。

註：Facebook AI Lab出品的PyTorch 是深度學習框架的一匹黑馬，靠Eager evaluation 博得了大批深度學習研究人員的青睞。基於Imperative programming 的理念和基於Python的語言構造（控制流）來搭建神經網路，靈活性高。NLP 領域常有一些動態圖的需求，PyTorch是首選。我們認為在單機場景下，易用性和靈活性是最重要的用戶需求，其它框架為了滿足這樣的需求，必須在原本為分散式設計的技術架構上做一些妥協，難以和PyTorch的最簡內核競爭。PyTorch 為了克服Eager evaluation的一些問題，也在通過JIT來享受一些Lazy evaluation的優點，同時也在向分散式場景進軍。如前文所述，在大規模分散式應用場景下，用戶程序只能是Lazy evaluation風格，數據流的執行引擎在高並發高並行的場景有天然的優勢，PyTorch 現在的設計實現距離這個目標還比較遙遠。

註：MXNet 開發團隊實力雄厚，現在是Apache孵化項目，有Amazon官方支持加持。MXNet 的特點是包含很多正對Geek品味的實現技巧， 很值得喜歡鑽研前沿技術的人去觀摩。但不好的地方是，給人一種比較「雜」的感覺，讓開發者感到困惑，譬如矩陣庫有兩套實現Mshadow和NDArray。MXNet 在計算機視覺領域總是能緊跟前沿應用，一些新的模型出來社區總是第一時間支持。MXNet 有一些關聯項目，譬如NNVM和TVM，目前來看TVM 更獨特，NNVM里實現的一些圖優化技術在其它框架里也是有對應實現的，而TVM 和TensorFlow XLA 應該是處於一個層次：聚焦於單設備程序性能優化。基於TVM, Halide, TensorFlow XLA ，用戶可以在單設備上使用聲明式編程，編譯器自動生成高效的後端代碼。

註：Caffe 的用戶還非常多，第二代 Caffe2和第一代已經迥然不同了，但繼承了一些簡潔的品質在裡面。框架的抽象非常簡潔，不厚重，Op/Kernel和引擎主要在C++層實現，而複雜的圖拓撲結構在Python層面處理。Caffe2 借鑒了TensorFlow 對Op/Kernel的抽象，沒有再使用之前Layer那種把數據和操作放在一起的設計。同樣是Op/Kernel抽象，Caffe2 也不是簡單的模仿， 代碼看上去更加舒服。Caffe2 目前支持數據並行，曾創造了一個小時訓練ImageNet的記錄，對Caffe 有感情的用戶可以嘗試。據了解，Caffe2 在Facebook內部承擔了「工業級應用」和「部署」的重任，也開發了很好的移動端支持，這應該是Caffe2 的特色。Caffe2 還有一個很有特色的技術，gloo網路庫是定製的MPI 實現，有「去中心化」集群通信的味道，也便於支持流水線。

註：PaddlePaddle 最大的優勢是百度內部在廣泛使用，經受了實戰檢驗。第一代PaddlePaddle 是比較接近 Caffe，分散式並行也依賴於Parameter server。最近一年，Paddle團隊在開展非常激進的重構。以我們的理解，重構版PaddlePaddle 借鑒了很多TensorFlow的設計，所以Paddle 能否解決TensorFlow 面臨的問題呢？ 重構後的PaddlePaddle 主推命令式編程介面，正像我們評價PyTorch時所說的，命令式編程介面固然親民，但數據流表示在大規模分散式運行場景有天然的優勢（表示能力和引擎實現複雜度方面），要很好的支持大規模分散式場景，深度學習框架最終要把「控制流」代碼翻譯成「數據流」代碼在後台運行。

註：總體來看，深度學習框架開發的絕大部分技術秘密已經公開了，開發一個簡單的深度學習框架沒有那麼難。但從另一方面看，要開發一個易用性和高效性都很卓越的框架是非常困難的，即使是世界上最優秀的開發團隊也感到困難重重，克服這些問題需要艱苦卓絕的創新和堅持。

註：展望未來一年：（1） 我們認為在計算機視覺領域也將迎來模型更大的場景，譬如把Hinton的Capsule net 從Cifar 推向ImageNet 規模，模型大小就遠超當前各種常見的CNN, 這種情況必須把模型分裂到多個設備上去完成，也就是所謂的模型並行。而且學界很關心神經網路結構學習，元學習，也有必要去探索CNN之外的架構，在人腦視皮層儘管存在CNN這種神經元分層組織和局部感受野的神經元，但並沒有所謂weight sharing（神經元功能柱有類似的特異選擇性，但不是嚴格一樣），這種神經元之間的連接規模非常龐大，如果去掉這個約束會發生什麼？如果深度學習框架不能支持模型並行，那麼這種設想就只能在Cifar, MNIST這種數據集合上去探索，並不能在ImageNet甚至更大規模的數據上去探索，而有些規律是在大規模數據上才會「湧現」出來。（2）未來，通用的深度學習框架會支持模型並行，而且用戶使用模型並行易如探囊取物。（3）深度學習向更多場景滲透，自然而然。（4）一旦一項技術被驗證，各種深度學習框架都會去擁抱它，支持它，正如今天很多框架在提供imperative programming介面一樣，同質化是比較嚴重的，未來需要一些新的空氣，新的思路去解決那些懸而未決的問題。（5）在大數據和人工智慧時代，數據積累到了臨界點，工業界除了有數據存儲，數據篩選這些需求，將普遍需要一個大腦（Brain），作為數據驅動業務的引擎，深度學習框架會像Hadoop 一樣經歷一個從「舊時王謝堂前燕，飛入尋常百姓家」的過程。

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！