一文讀懂TensorFlow的框架結構

2023-12-10旅遊

TensorFlow框架結構是指TensorFlow的系統設計和實作 ，它包括以下幾個方面：

TensorFlow基於數據流圖（Data Flow Graph） ，用於大規模分布式數值計算的開源框架1\。數據流圖由節點（Node）和邊（Edge）組成，節點表示計算單元（如算術運算、變量、常量等），邊表示數據單元（如張量、變量等）在節點間的傳遞。

TensorFlow以C API為界，將整個系統分為前端和後端兩個子系統 。前端系統提供編程模型，負責構造計算圖；後端系統提供執行時環境，負責執行計算圖。

TensorFlow的前端系統支持多種語言介面 ，如Python、C++、Java等，使用者可以透過這些介面定義計算圖和控制其執行。

TensorFlow的後端系統遵循良好的分層架構，可以進一步分解為四層：

執行時（Runtime）： 分別提供本地模式和分布式模式，並共享大部份設計和實作；

計算層（Computation Layer ）：由各個OP的Kernel實作組成，在執行時，Kernel實作執行OP的具體數學運算；

通訊層（Communication Layer） ：基於gRPC實作元件間的數據交換，並能夠在支持IB網絡的節點間實作RDMA通訊；

器材層（Device Layer） ：計算器材是OP執行的主要載體，TensorFlow支持多種異構的計算器材類別，如CPU、GPU等。

TensorFlow的模型構造和執行流程一般包括以下幾個步驟：

圖構建（Graph Construction）：使用者在客戶端（Client）中基於TensorFlow的多語言編程介面，添加算子（OP），完成計算圖的構造。
圖傳遞（Graph Transfer）：客戶端開啟會話（Session），透過它建立和主控節點（Master）之間的連線。執行Session.run()時，將構造好的Graph序列化為GraphDef後，以Protobuf的格式傳遞給Master。
圖剪枝（Graph Pruning）：Master根據Session.run()傳遞的Fetches和Feeds列表，反向遍歷全圖（Full Graph），實施剪枝，得到最小依賴子圖（Minimal Dependent Subgraph）。
圖分裂（Graph Partitioning）：Master將最小子圖分裂為多個圖分片（Graph Partition），並註冊到多個工作節點（Worker）上。一個Worker對應一個圖分片。
圖二次分裂（Graph Repartitioning）：Worker根據當前可用硬件資源，如CPU、GPU等，將圖分片按照OP算子器材約束規範（例如tf.device(‘/cpu:0’)），二次分裂到不同器材上。每個計算器材對應一個圖分片。
圖執行（Graph Execution）：對於每一個計算器材，Worker依照OP在Kernel中的實作，完成OP的運算。器材間數據通訊可以使用Send/Recv節點，而Worker間通訊，則使用gRPC或RDMA協定。