開源書籍
關(guān)于TensorFlow的內(nèi)核詳細(xì)原理與運行機制，請查閱我的開源技術(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)：提供編程模型，負(fù)責(zé)構(gòu)造計算圖；
- 后端系統(tǒng)：提供運行時環(huán)境，負(fù)責(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負(fù)責(zé)將該「子圖」再次分裂為多個「子圖片段」，以便在不同的進(jìn)程和設(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還要負(fù)責(zé)將OP運算的結(jié)果發(fā)送到其他的Work Service；或者接受來自其他Worker Service發(fā)送給它的OP運算的結(jié)果。
Kernel Implements
Kernel是OP在某種硬件設(shè)備的特定實現(xiàn)，它負(fù)責(zé)執(zhí)行OP的運算。
組件交互
如上圖所示，假設(shè)存在兩個任務(wù)：
-
/job:ps/task:0: 負(fù)責(zé)模型參數(shù)的存儲和更新 -
/job:worker/task:0: 負(fù)責(zé)模型的訓(xùn)練或推理
接下來，我們將進(jìn)一步抽絲剝繭，逐漸挖掘出TensorFlow計算圖的運行機制。
客戶端
Client基于TensorFlow的編程接口，構(gòu)造計算圖。目前，TensorFlow主流支持Python和C++的編程接口，并對其他編程語言接口的支持日益完善。
此時，TensorFlow并未執(zhí)行任何計算。直至建立Session會話，并以Session為橋梁，建立Client與后端運行時的通道，將Protobuf格式的GraphDef發(fā)送至Distributed Master。
也就是說，當(dāng)Client對OP結(jié)果進(jìn)行求值時，將觸發(fā)Distributed Master的計算圖的執(zhí)行過程。
如下圖所示，Client構(gòu)建了一個簡單計算圖。它首先將w與x進(jìn)行矩陣相乘，再與截距b按位相加，最后更新至s。
Distributed Master
在分布式的運行時環(huán)境中，Distributed Master根據(jù)Session.run的Fetching參數(shù)，從計算圖中反向遍歷，找到所依賴的最小子圖。
然后Distributed Master負(fù)責(zé)將該子圖再次分裂為多個「子圖片段」，以便在不同的進(jìn)程和設(shè)備上運行這些「子圖片段」。
最后，Distributed Master將這些圖片段派發(fā)給Work Service。隨后Work Service啟動「本地子圖」的執(zhí)行過程。
Distributed Master將會緩存「子圖片段」，以便后續(xù)執(zhí)行過程重復(fù)使用這些「子圖片段」，避免重復(fù)計算。
如上圖所示，Distributed Master開始執(zhí)行計算子圖。在執(zhí)行之前，Distributed Master會實施一系列優(yōu)化技術(shù)，例如「公共表達(dá)式消除」，「常量折疊」等。隨后，Distributed Master負(fù)責(zé)任務(wù)集的協(xié)同，執(zhí)行優(yōu)化后的計算子圖。
子圖片段
如上圖所示，存在一種合理的「子圖片段」劃分算法。Distributed Master將模型參數(shù)相關(guān)的OP進(jìn)行分組，并放置在PS任務(wù)上。其他OP則劃分為另外一組，放置在Worker任務(wù)上執(zhí)行。
SEND/RECV節(jié)點
如上圖所示，如果計算圖的邊被任務(wù)節(jié)點分割，Distributed Master將負(fù)責(zé)將該邊進(jìn)行分裂，在兩個分布式任務(wù)之間插入SEND和RECV節(jié)點，實現(xiàn)數(shù)據(jù)的傳遞。
隨后，Distributed Master將「子圖片段」派發(fā)給相應(yīng)的任務(wù)中執(zhí)行，在Worker Service成為「本地子圖」，它負(fù)責(zé)執(zhí)行該子圖的上的OP。
Worker Service
對于每個任務(wù)，都將存在相應(yīng)的Worker Service，它主要負(fù)責(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é)點進(jìn)行特化實現(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多個標(biāo)準(zhǔn)的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
