Tensorflow原始碼解析1 -- 核心架構和原始碼結構

摘要： 1 主流深度學習框架對比 當今的軟體開發基本都是分層化和模組化的，應用層開發會基於框架層。比如開發Linux Driver會基於Linux kernel，開發Android app會基於Android Framework。深度學習也不例外，框架層為上層模型開發提供了強大的多語言介面、穩定的...

1 主流深度學習框架對比

當今的軟體開發基本都是分層化和模組化的，應用層開發會基於框架層。比如開發Linux Driver會基於Linux kernel，開發Android app會基於Android Framework。深度學習也不例外，框架層為上層模型開發提供了強大的多語言介面、穩定的執行時、高效的運算元，以及完備的通訊層和裝置層管理層。因此，各大公司早早的就開始了深度學習框架的研發，以便能佔領市場。當前的框架有數十種之多，主流的如下（截止到2018年11月）

顯然TensorFlow是獨一無二的王者。第二名Keras，它是對TensorFlow或Theano介面的二次封裝，嚴格意義上並不是一個獨立的深度學習框架。TensorFlow目前也已經集成了Keras，使得安裝了TensorFlow的使用者就可以直接使用Keras了。

TensorFlow之所以能夠從數十種框架中脫穎而出，主要優點有

1. 出身高貴，是谷歌出品的。但其他很多框架出身也不差，例如PyTorch之於Facebook，MXNET之於Amazon
2. 2015年就開源了，比較早的俘獲了一大批開發者。這個確實是tf的一大先發優勢，但PyTorch的前身Caffe，以及MXNET開源時間都不晚，而且Caffe流行時間比tf早，後來才被趕超的。更有Theano這樣的絕對老前輩。由此可見，軟體開源是多麼重要。目前流行的深度學習框架也基本都開源了。
3. 支援的開發語言多，支援Python Java Go C++等多種流行語言。相比某些框架，確實是優勢很大。相比MXNET則小巫見大巫了。MXNET早期發展的一個主要方向就是前端多語言的支援，連MATLAB R Julia等語言都支援了。
4. 執行效率高。早期的時候，其實tf的執行效率比很多框架都要低一些的。
5. 安裝容易，使用者上手快，文件齊全，社群活躍。這個是tf的一個較大優勢，特別是社群方面，也就是我們常說的生態優勢。網際網路頭部集中效應十分明顯，體現在開源軟體上也是一樣。這也是我認為最大的一個優勢。

總結起來，TensorFlow雖然每個方面都不是絕對領先的優勢，但貴在每個方面都做的不錯，因此最終能夠一騎絕塵，獨領風騷。

學習Tensorflow框架核心，可以理解前端介面語言的支援，session生命週期，graph的構建、分裂和執行，operation的註冊和執行，模組間資料通訊，本地執行和分散式執行模式，以及CPU GPU TPU等異構裝置的封裝支援等。學習這些，對於模型的壓縮 加速 優化等都是大有裨益的。

2 TensorFlow系統架構

TensorFlow設計十分精巧，基於分層和模組化的設計思想進行開發的。框架如下圖

整個框架以C API為界，分為前端和後端兩大部分。

1. 前端：提供程式設計模型，多語言的介面支援，比如Python Java C++等。通過C API建立前後端的連線，後面詳細講解。

2. 後端：提供執行環境，完成計算圖的執行。進一步分為4層

   1. 執行時：分為分散式執行時和本地執行時，負責計算圖的接收，構造，編排等。
   2. 計算層：提供各op運算元的核心實現，例如conv2d, relu等
   3. 通訊層：實現元件間資料通訊，基於GRPC和RDMA兩種通訊方式
   4. 裝置層：提供多種異構裝置的支援，如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協議。

3 前端多語言實現 - swig包裝器

TensorFlow提供了很多種語言的前端介面，使得使用者可以通過多種語言來完成模型的訓練和推斷。其中Python支援得最好。這也是TensorFlow之所以受歡迎的一大原因。前端多語言是怎麼實現的呢？這要歸功於swig包裝器。

swig是個幫助使用C或者C++編寫的軟體能與其它各種高階程式語言進行嵌入聯接的開發工具。在TensorFlow使用bazel編譯時，swig會生成兩個wrapper檔案

1. pywrap_tensorflow_internal.py：對接上層Python呼叫
2. pywrap_tensorflow_internal.cc：對接底層C API呼叫。

pywrap_tensorflow_internal.py 模組被匯入時，會載入_pywrap_tensorflow_internal.so動態連結庫，它裡面包含了所有執行時介面的符號。而pywrap_tensorflow_internal.cc中，則註冊了一個函式符號表，實現Python介面和C介面的對映。執行時，就可以通過對映表，找到Python介面在C層的實現了。

4 tensorflow 原始碼結構

TensorFlow原始碼基本也是按照框架分層來組織檔案的。如下

其中core為tf的核心，它的原始碼結構如下

5 總結

TensorFlow框架設計精巧，程式碼量也很大，我們可以從以下部分逐步學習

1. TensorFlow核心架構和原始碼結構。先從全域性上對框架進行理解。
2. 前後端連線的橋樑--Session，重點理解session的生命週期，並通過相關原始碼可以加深理解Python前端如何呼叫底層C實現。
3. TensorFlow核心物件—Graph。圖graph是TensorFlow最核心的物件，基本都是圍繞著它來進行的。graph的節點為運算元operation，邊為資料tensor。
4. TensorFlow圖的節點 -- Operation。operation是圖graph的節點，承載了計算運算元。
5. TensorFlow圖的邊 -- Tensor。Tensor是圖graph的邊，承載了計算的資料。
6. TensorFlow本地執行時。
7. TensorFlow分散式執行時。和本地執行時有一些共用的介面，但區別也很大。
8. TensorFlow裝置層。主要了解裝置層的定義規範，以及實現。
9. TensorFlow佇列和並行運算。
10. TensorFlow斷點檢查checkpoint，模型儲存Saver，以及視覺化tensorboard。這三個為TensorFlow主要的工具。

原文連結

本文為雲棲社群原創內容，未經允許不得轉載。