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

在分散式的執行時環境中，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

對於每個任務，都將存在相應的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的計算。

TensorFlow Core

其中，大多數Kernel基於Eigen::Tensor實現。Eigen::Tensor是一個使用C++模板技術，為多核CPU/GPU生成高效的併發程式碼。但是，TensorFlow也可以靈活地直接使用cuDNN實現更高效的Kernel。

此外，TensorFlow實現了向量化技術，使得在移動裝置，及其滿足高吞吐量，以資料為中心的應用需求，實現更高效的推理。

如果對於複合OP的子計算過程很難表示，或執行效率低下，TensorFlow甚至支援更高效的Kernle實現的註冊，其擴充套件性表現相當優越。

技術棧

最後，按照TensorFlow的軟體層次，通過一張表格羅列TensorFlow的技術棧，以便更清晰地對上述內容做一個簡單回顧。

TensorFlow技術棧