2.4　TensorFlow架構

2.4.1　TensorFlow架構概述

TensorFlow運行時是一個跨平臺庫，圖2.10展示了它的一般架構，其中的C API可將不同語言的用戶級代碼與核心運行時分開。TensorFlow在設計時借鑒了Android的成功架構，即其生態構建使得層次和背景各異的小微開發者能夠簡單設計應用，并且易于遷移到類似生態。

TensorFlow架構包括以下核心組件。

• 客戶端：將計算定義為數據流圖和使用會話啟動圖執行。
• 分布式主服務：根據Session.run()的參數定義，從圖中裁剪出特定子圖；將子圖分成在不同進程和設備中運行的多個塊；將各個圖塊分發給工作器服務；通過工作器服務開始執行圖塊。
• 工作器服務（每項任務一個）：使用適合可用硬件（CPU、GPU等）的內核實現來安排圖操作的執行，向其他工作器服務發送操作結果以及從它們那里接收操作結果。
• 內核實現：對單個圖操作執行計算。

圖2.11說明了這些組件的交互。“/job:worker/task:0”和“/job:ps/task:0”都是工作器服務的任務。PS代表參數服務器（負責存儲和更新模型參數的任務）。其他任務會在優化這些參數時向這些參數發送更新。任務之間的這種特殊分工不是必需的，但對于分布式訓練來說很常見。

注意，分布式主服務和工作器服務僅存在于分布式TensorFlow中。TensorFlow的單進程版本包含一個特殊的會話實現，它可以執行分布式主服務執行的所有操作，但只與本地進程中的設備進行通信。

以下各部分將更詳細地介紹TensorFlow核心層，并逐步說明如何處理示例圖。

2.4.2　TensorFlow客戶端架構

用戶負責編寫可構建計算圖的客戶端TensorFlow程序。該程序可以直接組建各項操作，也可以使用Keras API之類的便利庫組建神經網絡層和其他更高級別的抽象物。TensorFlow支持多種客戶端語言，Python和C++是首先支持的，C++是TensorFlow框架的實現語言，而Python在數據處理工程師群體中有著廣泛的使用基礎，目前針對Java和JavaScript的API庫也日益成熟。隨著功能的日益成熟，通常會將它們移植到C++中，以便用戶可以通過所有客戶端語言訪問經過優化的實現。大多數訓練庫層面仍然只支持Python，主要目的是讓科研人員更加專注于機器學習算法的構建。

客戶端會創建一個會話，以便將圖定義作為tf.GraphDef協議緩沖區發送到分布式主服務。當客戶端評估圖中的一個或多個節點時，會調用分布式主服務來啟動計算。

在圖2.12中，客戶端構建了一個圖，將權重（w）應用于特征向量（x），添加偏差項（b）并將結果保存在變量（s）中。

2.4.3　TensorFlow分布式主服務架構

分布式主服務包含圖的低階操作，包括：對圖進行裁剪以獲得評估客戶端請求的節點所需的子圖，劃分圖以獲取每臺參與設備的圖塊，以及緩存這些塊以便可以在后續步驟中重復使用。

由于主服務可以看到某一步的總體計算，因此它會應用常見的子表達式消除或常量折疊等標準優化。然后，它會在一組任務中協調經過優化的子圖的執行，如圖2.13所示。

圖2.14顯示了示例圖的可能劃分。分布式主服務已對模型參數進行分組，以便將它們一起放在參數服務器上。

如果圖邊緣在劃分時被切割掉，分布式主服務就會插入發送和接收節點，以便在分布式任務之間傳遞信息（見圖2.15）。

然后，分布式主服務會將各個圖塊傳送到分布式任務，如圖2.16所示。

2.4.4　TensorFlow工作器服務架構

在工作器服務架構中，每項任務中的工作器服務會處理來自主服務的請求，為包含本地子圖的操作安排內核的執行時間，以及調解任務之間的直接通信。

我們會優化工作器服務，以便以較低的開銷運行大型圖。當前的實現每秒可以執行數萬個子圖，這可讓大量副本生成快速、精細的訓練步。工作器服務會將各個內核分派給本地設備并盡可能并行運行這些內核，例如通過使用多個CPU核或GPU流。

本地CPU和GPU設備之間的傳輸使用cudaMemcpyAsync() API來重疊計算和數據傳輸，如圖2.17所示。兩個本地GPU之間的傳輸使用對等DMA，以避免通過主機CPU進行代價高昂的復制。對于任務之間的傳輸，TensorFlow會使用多種協議，包括：基于TCP的gRPC和基于聚合以太網的RDMA。此外，還提供對NVIDIA的NCCL庫的初步支持，以進行多GPU通信。

2.4.5　TensorFlow內核架構

內核是TensorFlow真正的能力所在，運行時包含200多個標準操作，包括數學、數組、控制流和狀態管理操作。其中，每項操作都具有針對各種設備優化的內核實現。許多操作內核都是使用Eigen::Tensor實現的，后者使用C++模板為多核CPU和GPU生成高效的并行代碼。不過，我們可以隨意使用cuDNN實現更高效內核的庫。另外，還實現了量化，能夠在移動設備和高吞吐量數據中心應用等環境中實現更快的推斷，并使用GEMMLOWP低精度矩陣庫加快量化計算。

如果將子計算表示為操作組合很困難或效率低下，那么用戶可以注冊提供高效實現（用C++編寫）的其他內核。例如，建議為一些對性能要求苛刻的操作注冊自己的混合內核，例如ReLU和S型激活函數及其對應的梯度。另外，XLA編譯器具有自動內核混合的實驗性實現。