開源書籍
關(guān)于TensorFlow的內(nèi)核詳細原理與運行機制茅逮,請查閱我的開源技術(shù)書籍:TensorFlow內(nèi)核剖析。
TensorFlow是什么竹握?
TensorFlow
基于數(shù)據(jù)流圖鄙陡,用于大規(guī)模分布式數(shù)值計算的開源框架诺苹。節(jié)點表示某種抽象的計算,邊表示節(jié)點之間相互聯(lián)系的張量婶博。
TensorFlow
支持各種異構(gòu)的平臺瓮具,支持多CPU/GPU
,服務(wù)器,移動設(shè)備呜呐,具有良好的跨平臺的特性想际;TensorFlow
架構(gòu)靈活,能夠支持各種網(wǎng)絡(luò)模型兑巾,具有良好的通用性;此外忠荞,TensorFlow
架構(gòu)具有良好的可擴展性蒋歌,對OP
的擴展支持,Kernel
特化方面表現(xiàn)出眾委煤。
TensorFlow
最初由Google
大腦的研究員和工程師開發(fā)出來,用于機器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)方面的研究碧绞,于2015.10
宣布開源府框,在眾多深度學(xué)習(xí)框架中脫穎而出,在Github
上獲得了最多的Star
量讥邻。
本文將闡述TensorFlow
的系統(tǒng)架構(gòu)迫靖,幫助讀者加深理解TensorFlow
的工作機理。
本文假設(shè)讀者已經(jīng)了解
TensorFlow
的基本編程模型兴使,包括計算圖,OP
,Tensor
,Session
等基本概念系宜。
系統(tǒng)概述
TensorFlow
的系統(tǒng)結(jié)構(gòu)以C API
為界,將整個系統(tǒng)分為「前端」和「后端」兩個子系統(tǒng):
- 前端系統(tǒng):提供編程模型发魄,負責(zé)構(gòu)造計算圖盹牧;
- 后端系統(tǒng):提供運行時環(huán)境,負責(zé)執(zhí)行計算圖励幼。
如上圖所示汰寓,重點關(guān)注系統(tǒng)中如下4個基本組件,它們是系統(tǒng)分布式運行機制的核心苹粟。
Client
Client
是前端系統(tǒng)的主要組成部分有滑,它是一個支持多語言的編程環(huán)境。它提供基于計算圖的編程模型六水,方便用戶構(gòu)造各種復(fù)雜的計算圖俺孙,實現(xiàn)各種形式的模型設(shè)計辣卒。
Client
通過Session
為橋梁,連接TensorFlow
后端的「運行時」睛榄,并啟動計算圖的執(zhí)行過程荣茫。
Distributed Master
在分布式的運行時環(huán)境中,Distributed Master
根據(jù)Session.run
的Fetching
參數(shù)场靴,從計算圖中反向遍歷啡莉,找到所依賴的「最小子圖」。
然后旨剥,Distributed Master
負責(zé)將該「子圖」再次分裂為多個「子圖片段」咧欣,以便在不同的進程和設(shè)備上運行這些「子圖片段」。
最后轨帜,Distributed Master
將這些「子圖片段」派發(fā)給Work Service
魄咕;隨后Work Service
啟動「子圖片段」的執(zhí)行過程。
Worker Service
對于每以個任務(wù)蚌父,TensorFlow
都將啟動一個Worker Service
哮兰。Worker Service
將按照計算圖中節(jié)點之間的依賴關(guān)系,根據(jù)當(dāng)前的可用的硬件環(huán)境(GPU/CPU)苟弛,調(diào)用OP
的Kernel
實現(xiàn)完成OP
的運算(一種典型的多態(tài)實現(xiàn)技術(shù))喝滞。
另外,Worker Service
還要負責(zé)將OP
運算的結(jié)果發(fā)送到其他的Work Service
膏秫;或者接受來自其他Worker Service
發(fā)送給它的OP
運算的結(jié)果右遭。
Kernel Implements
Kernel
是OP
在某種硬件設(shè)備的特定實現(xiàn),它負責(zé)執(zhí)行OP
的運算缤削。
組件交互
如上圖所示窘哈,假設(shè)存在兩個任務(wù):
-
/job:ps/task:0
: 負責(zé)模型參數(shù)的存儲和更新 -
/job:worker/task:0
: 負責(zé)模型的訓(xùn)練或推理
接下來,我們將進一步抽絲剝繭僻他,逐漸挖掘出TensorFlow
計算圖的運行機制宵距。
客戶端
Client
基于TensorFlow
的編程接口,構(gòu)造計算圖吨拗。目前,TensorFlow
主流支持Python
和C++
的編程接口婿斥,并對其他編程語言接口的支持日益完善劝篷。
此時,TensorFlow
并未執(zhí)行任何計算民宿。直至建立Session
會話娇妓,并以Session
為橋梁,建立Client
與后端運行時的通道活鹰,將Protobuf
格式的GraphDef
發(fā)送至Distributed Master
哈恰。
也就是說只估,當(dāng)Client
對OP
結(jié)果進行求值時,將觸發(fā)Distributed Master
的計算圖的執(zhí)行過程着绷。
如下圖所示蛔钙,Client
構(gòu)建了一個簡單計算圖。它首先將w
與x
進行矩陣相乘荠医,再與截距b
按位相加吁脱,最后更新至s
。
Distributed Master
在分布式的運行時環(huán)境中彬向,Distributed Master
根據(jù)Session.run
的Fetching
參數(shù)兼贡,從計算圖中反向遍歷,找到所依賴的最小子圖娃胆。
然后Distributed Master
負責(zé)將該子圖再次分裂為多個「子圖片段」遍希,以便在不同的進程和設(shè)備上運行這些「子圖片段」。
最后里烦,Distributed Master
將這些圖片段派發(fā)給Work Service
凿蒜。隨后Work Service
啟動「本地子圖」的執(zhí)行過程。
Distributed Master
將會緩存「子圖片段」招驴,以便后續(xù)執(zhí)行過程重復(fù)使用這些「子圖片段」篙程,避免重復(fù)計算。
如上圖所示别厘,Distributed Master
開始執(zhí)行計算子圖虱饿。在執(zhí)行之前,Distributed Master
會實施一系列優(yōu)化技術(shù)触趴,例如「公共表達式消除」氮发,「常量折疊」等。隨后冗懦,Distributed Master
負責(zé)任務(wù)集的協(xié)同爽冕,執(zhí)行優(yōu)化后的計算子圖。
子圖片段
如上圖所示披蕉,存在一種合理的「子圖片段」劃分算法颈畸。Distributed Master
將模型參數(shù)相關(guān)的OP
進行分組,并放置在PS
任務(wù)上没讲。其他OP
則劃分為另外一組眯娱,放置在Worker
任務(wù)上執(zhí)行。
SEND/RECV節(jié)點
如上圖所示爬凑,如果計算圖的邊被任務(wù)節(jié)點分割徙缴,Distributed Master
將負責(zé)將該邊進行分裂,在兩個分布式任務(wù)之間插入SEND
和RECV
節(jié)點嘁信,實現(xiàn)數(shù)據(jù)的傳遞于样。
隨后疏叨,Distributed Master
將「子圖片段」派發(fā)給相應(yīng)的任務(wù)中執(zhí)行,在Worker Service
成為「本地子圖」穿剖,它負責(zé)執(zhí)行該子圖的上的OP
蚤蔓。
Worker Service
對于每個任務(wù),都將存在相應(yīng)的Worker Service
携御,它主要負責(zé)如下3個方面的職責(zé):
- 處理來自
Master
的請求昌粤; - 調(diào)度
OP
的Kernel
實現(xiàn),執(zhí)行本地子圖啄刹; - 協(xié)同任務(wù)之間的數(shù)據(jù)通信涮坐。
Worker Service
派發(fā)OP
到本地設(shè)備,執(zhí)行Kernel
的特定實現(xiàn)誓军。它將盡最大可能地利用多CPU/GPU
的處理能力袱讹,并發(fā)地執(zhí)行Kernel
實現(xiàn)。
另外昵时,TensorFlow
根據(jù)設(shè)備類型捷雕,對于設(shè)備間的SEND/RECV
節(jié)點進行特化實現(xiàn):
- 使用
cudaMemcpyAsync
的API實現(xiàn)本地CPU
與GPU
設(shè)備的數(shù)據(jù)傳輸; - 對于本地的
GPU
之間則使用端到端的DMA
壹甥,避免了跨host CPU
昂貴的拷貝過程救巷。
對于任務(wù)之間的數(shù)據(jù)傳遞,TensorFlow
支持多協(xié)議句柠,主要包括:
- gRPC over TCP
- RDMA over Converged Ethernet
Kernel Implements
TensorFlow
的運行時包含200
多個標準的OP
浦译,包括數(shù)值計算,多維數(shù)組操作溯职,控制流精盅,狀態(tài)管理等。每一個OP
根據(jù)設(shè)備類型都會存在一個優(yōu)化了的Kernel
實現(xiàn)谜酒。在運行時叹俏,運行時根據(jù)本地設(shè)備的類型,為OP
選擇特定的Kernel
實現(xiàn)僻族,完成該OP
的計算粘驰。
其中,大多數(shù)Kernel
基于Eigen::Tensor
實現(xiàn)述么。Eigen::Tensor
是一個使用C++
模板技術(shù)晴氨,為多核CPU/GPU
生成高效的并發(fā)代碼。但是碉输,TensorFlow
也可以靈活地直接使用cuDNN
實現(xiàn)更高效的Kernel
。
此外亭珍,TensorFlow
實現(xiàn)了矢量化技術(shù)敷钾,使得在移動設(shè)備枝哄,及其滿足高吞吐量,以數(shù)據(jù)為中心的應(yīng)用需求阻荒,實現(xiàn)更高效的推理挠锥。
如果對于復(fù)合OP
的子計算過程很難表示,或執(zhí)行效率低下侨赡,TensorFlow
甚至支持更高效的Kernle
實現(xiàn)的注冊蓖租,其擴展性表現(xiàn)相當(dāng)優(yōu)越。
技術(shù)棧
最后羊壹,按照TensorFlow
的軟件層次蓖宦,通過一張表格羅列TensorFlow
的技術(shù)棧,以便更清晰地對上述內(nèi)容做一個簡單回顧油猫。
開源技術(shù)書
https://github.com/horance-liu/tensorflow-internals