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 為界，將整個系統分為「前端」和「后端」兩個子系統：

• 前端系統：提供編程模型，負責構造計算圖；
• 后端系統：提供運行時環境，負責執行計算圖。

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技術棧

 

來自：http://www.jianshu.com/p/a5574ebcdeab