作者：劉光聰

鏈接：http://www.reibang.com/p/a5574ebcdeab

TensorFlow是什么？

TensorFlow基于數(shù)據(jù)流圖风秤，用于大規(guī)模分布式數(shù)值計(jì)算的開源框架骨稿。節(jié)點(diǎn)表示某種抽象的計(jì)算笨鸡，邊表示節(jié)點(diǎn)之間相互聯(lián)系的張量。

計(jì)算圖實(shí)例

TensorFlow支持各種異構(gòu)的平臺(tái)坦冠，支持多CPU/GPU形耗，服務(wù)器，移動(dòng)設(shè)備辙浑，具有良好的跨平臺(tái)的特性激涤；TensorFlow架構(gòu)靈活，能夠支持各種網(wǎng)絡(luò)模型判呕，具有良好的通用性倦踢；此外，TensorFlow架構(gòu)具有良好的可擴(kuò)展性侠草，對(duì)OP的擴(kuò)展支持辱挥，Kernel特化方面表現(xiàn)出眾。

TensorFlow最初由Google大腦的研究員和工程師開發(fā)出來边涕，用于機(jī)器學(xué)習(xí)和神經(jīng)網(wǎng)絡(luò)方面的研究晤碘，于2015.10宣布開源褂微，在眾多深度學(xué)習(xí)框架中脫穎而出，在Github上獲得了最多的Star量园爷。

本文將闡述TensorFlow的系統(tǒng)架構(gòu)宠蚂，幫助讀者加深理解TensorFlow的工作機(jī)理。

本文假設(shè)讀者已經(jīng)了解TensorFlow的基本編程模型童社，包括計(jì)算圖,OP,Tensor,Session等基本概念求厕。

系統(tǒng)概述

TensorFlow的系統(tǒng)結(jié)構(gòu)以C API為界，將整個(gè)系統(tǒng)分為「前端」和「后端」兩個(gè)子系統(tǒng)：

前端系統(tǒng)：提供編程模型扰楼，負(fù)責(zé)構(gòu)造計(jì)算圖呀癣；

后端系統(tǒng)：提供運(yùn)行時(shí)環(huán)境，負(fù)責(zé)執(zhí)行計(jì)算圖灭抑。

TensorFlow系統(tǒng)架構(gòu)

如上圖所示，重點(diǎn)關(guān)注系統(tǒng)中如下4個(gè)基本組件抵代，它們是系統(tǒng)分布式運(yùn)行機(jī)制的核心腾节。

Client

Client是前端系統(tǒng)的主要組成部分，它是一個(gè)支持多語言的編程環(huán)境荤牍。它提供基于計(jì)算圖的編程模型案腺，方便用戶構(gòu)造各種復(fù)雜的計(jì)算圖，實(shí)現(xiàn)各種形式的模型設(shè)計(jì)康吵。

Client通過Session為橋梁劈榨，連接TensorFlow后端的「運(yùn)行時(shí)」，并啟動(dòng)計(jì)算圖的執(zhí)行過程晦嵌。

Distributed Master

在分布式的運(yùn)行時(shí)環(huán)境中同辣，Distributed Master根據(jù)Session.run的Fetching參數(shù)，從計(jì)算圖中反向遍歷惭载，找到所依賴的「最小子圖」旱函。

然后，Distributed Master負(fù)責(zé)將該「子圖」再次分裂為多個(gè)「子圖片段」描滔，以便在不同的進(jìn)程和設(shè)備上運(yùn)行這些「子圖片段」棒妨。

最后，Distributed Master將這些「子圖片段」派發(fā)給Work Service含长；隨后Work Service啟動(dòng)「子圖片段」的執(zhí)行過程券腔。

Worker Service

對(duì)于每以個(gè)任務(wù)，TensorFlow都將啟動(dòng)一個(gè)Worker Service拘泞。Worker Service將按照計(jì)算圖中節(jié)點(diǎn)之間的依賴關(guān)系纷纫，根據(jù)當(dāng)前的可用的硬件環(huán)境(GPU/CPU)，調(diào)用OP的Kernel實(shí)現(xiàn)完成OP的運(yùn)算(一種典型的多態(tài)實(shí)現(xiàn)技術(shù))陪腌。

另外涛酗，Worker Service還要負(fù)責(zé)將OP運(yùn)算的結(jié)果發(fā)送到其他的Work Service铡原；或者接受來自其他Worker Service發(fā)送給它的OP運(yùn)算的結(jié)果。

Kernel Implements

Kernel是OP在某種硬件設(shè)備的特定實(shí)現(xiàn)商叹，它負(fù)責(zé)執(zhí)行OP的運(yùn)算燕刻。

組件交互

組件交互

如上圖所示，假設(shè)存在兩個(gè)任務(wù)：

/job:ps/task:0: 負(fù)責(zé)模型參數(shù)的存儲(chǔ)和更新

/job:worker/task:0: 負(fù)責(zé)模型的訓(xùn)練或推理

接下來剖笙，我們將進(jìn)一步抽絲剝繭卵洗，逐漸挖掘出TensorFlow計(jì)算圖的運(yùn)行機(jī)制。

客戶端

Client基于TensorFlow的編程接口弥咪，構(gòu)造計(jì)算圖过蹂。目前，TensorFlow主流支持Python和C++的編程接口聚至，并對(duì)其他編程語言接口的支持日益完善酷勺。

此時(shí)，TensorFlow并未執(zhí)行任何計(jì)算扳躬。直至建立Session會(huì)話脆诉，并以Session為橋梁，建立Client與后端運(yùn)行時(shí)的通道贷币，將Protobuf格式的GraphDef發(fā)送至Distributed Master击胜。

也就是說，當(dāng)Client對(duì)OP結(jié)果進(jìn)行求值時(shí)，將觸發(fā)Distributed Master的計(jì)算圖的執(zhí)行過程。

如下圖所示届氢，Client構(gòu)建了一個(gè)簡(jiǎn)單計(jì)算圖。它首先將w與x進(jìn)行矩陣相乘辰斋，再與截距b按位相加，最后更新至s瘸味。

構(gòu)造計(jì)算圖

Distributed Master

在分布式的運(yùn)行時(shí)環(huán)境中亡呵，Distributed Master根據(jù)Session.run的Fetching參數(shù)，從計(jì)算圖中反向遍歷硫戈，找到所依賴的最小子圖锰什。

然后Distributed Master負(fù)責(zé)將該子圖再次分裂為多個(gè)「子圖片段」，以便在不同的進(jìn)程和設(shè)備上運(yùn)行這些「子圖片段」丁逝。

最后汁胆，Distributed Master將這些圖片段派發(fā)給Work Service。隨后Work Service啟動(dòng)「本地子圖」的執(zhí)行過程霜幼。

Distributed Master將會(huì)緩存「子圖片段」嫩码，以便后續(xù)執(zhí)行過程重復(fù)使用這些「子圖片段」，避免重復(fù)計(jì)算罪既。

執(zhí)行圖計(jì)算

如上圖所示铸题，Distributed Master開始執(zhí)行計(jì)算子圖铡恕。在執(zhí)行之前，Distributed Master會(huì)實(shí)施一系列優(yōu)化技術(shù)丢间，例如「公共表達(dá)式消除」探熔，「常量折疊」等。隨后烘挫，Distributed Master負(fù)責(zé)任務(wù)集的協(xié)同诀艰，執(zhí)行優(yōu)化后的計(jì)算子圖。

子圖片段

子圖片段

如上圖所示饮六，存在一種合理的「子圖片段」劃分算法其垄。Distributed Master將模型參數(shù)相關(guān)的OP進(jìn)行分組，并放置在PS任務(wù)上卤橄。其他OP則劃分為另外一組绿满，放置在Worker任務(wù)上執(zhí)行。

SEND/RECV節(jié)點(diǎn)

插入SEND/RECV節(jié)點(diǎn)

如上圖所示窟扑，如果計(jì)算圖的邊被任務(wù)節(jié)點(diǎn)分割喇颁，Distributed Master將負(fù)責(zé)將該邊進(jìn)行分裂，在兩個(gè)分布式任務(wù)之間插入SEND和RECV節(jié)點(diǎn)辜膝，實(shí)現(xiàn)數(shù)據(jù)的傳遞无牵。

隨后漾肮，Distributed Master將「子圖片段」派發(fā)給相應(yīng)的任務(wù)中執(zhí)行厂抖，在Worker Service成為「本地子圖」，它負(fù)責(zé)執(zhí)行該子圖的上的OP克懊。

Worker Service

對(duì)于每個(gè)任務(wù)忱辅，都將存在相應(yīng)的Worker Service，它主要負(fù)責(zé)如下3個(gè)方面的職責(zé)：

處理來自Master的請(qǐng)求谭溉；

調(diào)度OP的Kernel實(shí)現(xiàn)墙懂，執(zhí)行本地子圖；

協(xié)同任務(wù)之間的數(shù)據(jù)通信扮念。

執(zhí)行本地子圖

Worker Service派發(fā)OP到本地設(shè)備损搬，執(zhí)行Kernel的特定實(shí)現(xiàn)。它將盡最大可能地利用多CPU/GPU的處理能力柜与，并發(fā)地執(zhí)行Kernel實(shí)現(xiàn)巧勤。

另外，TensorFlow根據(jù)設(shè)備類型弄匕，對(duì)于設(shè)備間的SEND/RECV節(jié)點(diǎn)進(jìn)行特化實(shí)現(xiàn)：

使用cudaMemcpyAsync的API實(shí)現(xiàn)本地CPU與GPU設(shè)備的數(shù)據(jù)傳輸颅悉；

對(duì)于本地的GPU之間則使用端到端的DMA，避免了跨host CPU昂貴的拷貝過程迁匠。

對(duì)于任務(wù)之間的數(shù)據(jù)傳遞剩瓶，TensorFlow支持多協(xié)議驹溃，主要包括：

gRPC over TCP

RDMA over Converged Ethernet

Kernel Implements

TensorFlow的運(yùn)行時(shí)包含200多個(gè)標(biāo)準(zhǔn)的OP，包括數(shù)值計(jì)算延曙，多維數(shù)組操作豌鹤，控制流，狀態(tài)管理等搂鲫。每一個(gè)OP根據(jù)設(shè)備類型都會(huì)存在一個(gè)優(yōu)化了的Kernel實(shí)現(xiàn)傍药。在運(yùn)行時(shí)，運(yùn)行時(shí)根據(jù)本地設(shè)備的類型魂仍，為OP選擇特定的Kernel實(shí)現(xiàn)拐辽，完成該OP的計(jì)算。

TensorFlow Core

其中擦酌，大多數(shù)Kernel基于Eigen::Tensor實(shí)現(xiàn)俱诸。Eigen::Tensor是一個(gè)使用C++模板技術(shù)，為多核CPU/GPU生成高效的并發(fā)代碼赊舶。但是睁搭，TensorFlow也可以靈活地直接使用cuDNN實(shí)現(xiàn)更高效的Kernel。

此外笼平，TensorFlow實(shí)現(xiàn)了矢量化技術(shù)园骆，使得在移動(dòng)設(shè)備，及其滿足高吞吐量寓调，以數(shù)據(jù)為中心的應(yīng)用需求锌唾，實(shí)現(xiàn)更高效的推理。

如果對(duì)于復(fù)合OP的子計(jì)算過程很難表示夺英，或執(zhí)行效率低下晌涕，TensorFlow甚至支持更高效的Kernle實(shí)現(xiàn)的注冊(cè)，其擴(kuò)展性表現(xiàn)相當(dāng)優(yōu)越痛悯。

技術(shù)棧

最后余黎，按照TensorFlow的軟件層次，通過一張表格羅列TensorFlow的技術(shù)棧载萌，以便更清晰地對(duì)上述內(nèi)容做一個(gè)簡(jiǎn)單回顧惧财。

相關(guān)學(xué)習(xí)資料移步：