TensorFlow系統架構
TensorFlow設計十分精巧,基於分層和模塊化的設計思想進行開發的。框架如下圖
整個框架以C API為界,分為前端和后端兩大部分。
- 前端:提供編程模型,多語言的接口支持,比如Python Java C++等。通過C API建立前后端的連接,后面詳細講解。
- 后端:提供運行環境,完成計算圖的執行。進一步分為4層
- 運行時:分為分布式運行時和本地運行時,負責計算圖的接收,構造,編排等。
- 計算層:提供各op算子的內核實現,例如conv2d, relu等
- 通信層:實現組件間數據通信,基於GRPC和RDMA兩種通信方式
- 設備層:提供多種異構設備的支持,如CPU GPU TPU FPGA等
模型構造和執行流程
TensorFlow的一大特點是,圖的構造和執行相分離。用戶添加完算子,構建好整圖后,才開始進行訓練和執行,也就是圖的執行。大體流程如下
1. 圖構建:用戶在client中基於TensorFlow的多語言編程接口,添加算子,完成計算圖的構造。
2. 圖傳遞:client開啟session,通過它建立和master之間的連接。執行session.run()時,將構造好的graph序列化為graphDef后,以protobuf的格式傳遞給master。
3.圖剪枝:master根據session.run()傳遞的fetches和feeds列表,反向遍歷全圖full graph,實施剪枝,得到最小依賴子圖
4. 圖分裂:master將最小子圖分裂為多個Graph Partition,並注冊到多個worker上。一個worker對應一個Graph Partition。
5. 圖二次分裂:worker根據當前可用硬件資源,如CPU GPU,將Graph Partition按照op算子設備約束規范(例如tf.device(’/cpu:0’),二次分裂到不同設備上。每個計算設備對應一個Graph Partition。
6. 圖運行:對於每一個計算設備,worker依照op在kernel中的實現,完成op的運算。設備間數據通信可以使用send/recv節點,而worker間通信,則使用GRPC或RDMA協議。
前端多語言實現 - swig包裝器
swig是個幫助使用C或者C++編寫的軟件能與其它各種高級編程語言進行嵌入聯接的開發工具。在TensorFlow使用bazel編譯時,swig會生成兩個wrapper文件
- pywrap_tensorflow_internal.py:對接上層Python調用
- pywrap_tensorflow_internal.cc:對接底層C API調用。
pywrap_tensorflow_internal.py 模塊被導入時,會加載_pywrap_tensorflow_internal.so動態鏈接庫,它里面包含了所有運行時接口的符號。而pywrap_tensorflow_internal.cc中,則注冊了一個函數符號表,實現Python接口和C接口的映射。運行時,就可以通過映射表,找到Python接口在C層的實現了。
tensorflow 源碼結構
TensorFlow源碼基本也是按照框架分層來組織文件的。如下
其中core為tf的核心,它的源碼結構如下