1. 程式人生 > >Tensorflow原始碼解析1 -- 核心架構和原始碼結構

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

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主要的工具。