作者:劉光聰
鏈接:http://www.lxweimin.com/p/a5574ebcdeab
TensorFlow是什么?
TensorFlow基于數據流圖,用于大規模分布式數值計算的開源框架。節點表示某種抽象的計算,邊表示節點之間相互聯系的張量。
計算圖實例
TensorFlow支持各種異構的平臺,支持多CPU/GPU,服務器,移動設備,具有良好的跨平臺的特性;TensorFlow架構靈活,能夠支持各種網絡模型,具有良好的通用性;此外,TensorFlow架構具有良好的可擴展性,對OP的擴展支持,Kernel特化方面表現出眾。
TensorFlow最初由Google大腦的研究員和工程師開發出來,用于機器學習和神經網絡方面的研究,于2015.10宣布開源,在眾多深度學習框架中脫穎而出,在Github上獲得了最多的Star量。
本文將闡述TensorFlow的系統架構,幫助讀者加深理解TensorFlow的工作機理。
本文假設讀者已經了解TensorFlow的基本編程模型,包括計算圖,OP,Tensor,Session等基本概念。
TensorFlow的系統結構以C API為界,將整個系統分為「前端」和「后端」兩個子系統:
前端系統:提供編程模型,負責構造計算圖;
后端系統:提供運行時環境,負責執行計算圖。
TensorFlow系統架構
如上圖所示,重點關注系統中如下4個基本組件,它們是系統分布式運行機制的核心。
Client
Client是前端系統的主要組成部分,它是一個支持多語言的編程環境。它提供基于計算圖的編程模型,方便用戶構造各種復雜的計算圖,實現各種形式的模型設計。
Client通過Session為橋梁,連接TensorFlow后端的「運行時」,并啟動計算圖的執行過程。
Distributed Master
在分布式的運行時環境中,Distributed Master根據Session.run的Fetching參數,從計算圖中反向遍歷,找到所依賴的「最小子圖」。
然后,Distributed Master負責將該「子圖」再次分裂為多個「子圖片段」,以便在不同的進程和設備上運行這些「子圖片段」。
最后,Distributed Master將這些「子圖片段」派發給Work Service;隨后Work Service啟動「子圖片段」的執行過程。
Worker Service
對于每以個任務,TensorFlow都將啟動一個Worker Service。Worker Service將按照計算圖中節點之間的依賴關系,根據當前的可用的硬件環境(GPU/CPU),調用OP的Kernel實現完成OP的運算(一種典型的多態實現技術)。
另外,Worker Service還要負責將OP運算的結果發送到其他的Work Service;或者接受來自其他Worker Service發送給它的OP運算的結果。
Kernel Implements
Kernel是OP在某種硬件設備的特定實現,它負責執行OP的運算。
組件交互
如上圖所示,假設存在兩個任務:
/job:ps/task:0: 負責模型參數的存儲和更新
/job:worker/task:0: 負責模型的訓練或推理
接下來,我們將進一步抽絲剝繭,逐漸挖掘出TensorFlow計算圖的運行機制。
Client基于TensorFlow的編程接口,構造計算圖。目前,TensorFlow主流支持Python和C++的編程接口,并對其他編程語言接口的支持日益完善。
此時,TensorFlow并未執行任何計算。直至建立Session會話,并以Session為橋梁,建立Client與后端運行時的通道,將Protobuf格式的GraphDef發送至Distributed Master。
也就是說,當Client對OP結果進行求值時,將觸發Distributed Master的計算圖的執行過程。
如下圖所示,Client構建了一個簡單計算圖。它首先將w與x進行矩陣相乘,再與截距b按位相加,最后更新至s。
構造計算圖
在分布式的運行時環境中,Distributed Master根據Session.run的Fetching參數,從計算圖中反向遍歷,找到所依賴的最小子圖。
然后Distributed Master負責將該子圖再次分裂為多個「子圖片段」,以便在不同的進程和設備上運行這些「子圖片段」。
最后,Distributed Master將這些圖片段派發給Work Service。隨后Work Service啟動「本地子圖」的執行過程。
Distributed Master將會緩存「子圖片段」,以便后續執行過程重復使用這些「子圖片段」,避免重復計算。
執行圖計算
如上圖所示,Distributed Master開始執行計算子圖。在執行之前,Distributed Master會實施一系列優化技術,例如「公共表達式消除」,「常量折疊」等。隨后,Distributed Master負責任務集的協同,執行優化后的計算子圖。
子圖片段
子圖片段
如上圖所示,存在一種合理的「子圖片段」劃分算法。Distributed Master將模型參數相關的OP進行分組,并放置在PS任務上。其他OP則劃分為另外一組,放置在Worker任務上執行。
SEND/RECV節點
插入SEND/RECV節點
如上圖所示,如果計算圖的邊被任務節點分割,Distributed Master將負責將該邊進行分裂,在兩個分布式任務之間插入SEND和RECV節點,實現數據的傳遞。
隨后,Distributed Master將「子圖片段」派發給相應的任務中執行,在Worker Service成為「本地子圖」,它負責執行該子圖的上的OP。
對于每個任務,都將存在相應的Worker Service,它主要負責如下3個方面的職責:
處理來自Master的請求;
調度OP的Kernel實現,執行本地子圖;
協同任務之間的數據通信。
執行本地子圖
Worker Service派發OP到本地設備,執行Kernel的特定實現。它將盡最大可能地利用多CPU/GPU的處理能力,并發地執行Kernel實現。
另外,TensorFlow根據設備類型,對于設備間的SEND/RECV節點進行特化實現:
使用cudaMemcpyAsync的API實現本地CPU與GPU設備的數據傳輸;
對于本地的GPU之間則使用端到端的DMA,避免了跨host CPU昂貴的拷貝過程。
對于任務之間的數據傳遞,TensorFlow支持多協議,主要包括:
gRPC over TCP
RDMA over Converged Ethernet
TensorFlow的運行時包含200多個標準的OP,包括數值計算,多維數組操作,控制流,狀態管理等。每一個OP根據設備類型都會存在一個優化了的Kernel實現。在運行時,運行時根據本地設備的類型,為OP選擇特定的Kernel實現,完成該OP的計算。
TensorFlow Core
其中,大多數Kernel基于Eigen::Tensor實現。Eigen::Tensor是一個使用C++模板技術,為多核CPU/GPU生成高效的并發代碼。但是,TensorFlow也可以靈活地直接使用cuDNN實現更高效的Kernel。
此外,TensorFlow實現了矢量化技術,使得在移動設備,及其滿足高吞吐量,以數據為中心的應用需求,實現更高效的推理。
如果對于復合OP的子計算過程很難表示,或執行效率低下,TensorFlow甚至支持更高效的Kernle實現的注冊,其擴展性表現相當優越。
最后,按照TensorFlow的軟件層次,通過一張表格羅列TensorFlow的技術棧,以便更清晰地對上述內容做一個簡單回顧。
相關學習資料移步: