開源書籍
關于TensorFlow的內核詳細原理與運行機制,請查閱我的開源技術書籍:TensorFlow內核剖析。
TensorFlow是什么?
TensorFlow
基于數據流圖,用于大規模分布式數值計算的開源框架。節點表示某種抽象的計算,邊表示節點之間相互聯系的張量。
TensorFlow
支持各種異構的平臺,支持多CPU/GPU
,服務器,移動設備,具有良好的跨平臺的特性;TensorFlow
架構靈活,能夠支持各種網絡模型,具有良好的通用性;此外,TensorFlow
架構具有良好的可擴展性,對OP
的擴展支持,Kernel
特化方面表現出眾。
TensorFlow
最初由Google
大腦的研究員和工程師開發出來,用于機器學習和神經網絡方面的研究,于2015.10
宣布開源,在眾多深度學習框架中脫穎而出,在Github
上獲得了最多的Star
量。
本文將闡述TensorFlow
的系統架構,幫助讀者加深理解TensorFlow
的工作機理。
本文假設讀者已經了解
TensorFlow
的基本編程模型,包括計算圖,OP
,Tensor
,Session
等基本概念。
系統概述
TensorFlow
的系統結構以C API
為界,將整個系統分為「前端」和「后端」兩個子系統:
- 前端系統:提供編程模型,負責構造計算圖;
- 后端系統:提供運行時環境,負責執行計算圖。
如上圖所示,重點關注系統中如下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
在分布式的運行時環境中,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節點
如上圖所示,如果計算圖的邊被任務節點分割,Distributed Master
將負責將該邊進行分裂,在兩個分布式任務之間插入SEND
和RECV
節點,實現數據的傳遞。
隨后,Distributed Master
將「子圖片段」派發給相應的任務中執行,在Worker Service
成為「本地子圖」,它負責執行該子圖的上的OP
。
Worker Service
對于每個任務,都將存在相應的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
Kernel Implements
TensorFlow
的運行時包含200
多個標準的OP
,包括數值計算,多維數組操作,控制流,狀態管理等。每一個OP
根據設備類型都會存在一個優化了的Kernel
實現。在運行時,運行時根據本地設備的類型,為OP
選擇特定的Kernel
實現,完成該OP
的計算。
其中,大多數Kernel
基于Eigen::Tensor
實現。Eigen::Tensor
是一個使用C++
模板技術,為多核CPU/GPU
生成高效的并發代碼。但是,TensorFlow
也可以靈活地直接使用cuDNN
實現更高效的Kernel
。
此外,TensorFlow
實現了矢量化技術,使得在移動設備,及其滿足高吞吐量,以數據為中心的應用需求,實現更高效的推理。
如果對于復合OP
的子計算過程很難表示,或執行效率低下,TensorFlow
甚至支持更高效的Kernle
實現的注冊,其擴展性表現相當優越。
技術棧
最后,按照TensorFlow
的軟件層次,通過一張表格羅列TensorFlow
的技術棧,以便更清晰地對上述內容做一個簡單回顧。
開源技術書
https://github.com/horance-liu/tensorflow-internals