tensorflow源碼分析


前言:

       一般來說,如果安裝tensorflow主要目的是為了調試些小程序的話,只要下載相應的包,然后,直接使用pip install tensorflow即可。

      但有時我們需要將Tensorflow的功能移植到其它平台,這時就無法直接安裝了。需要我們下載相應的Tensorflow源碼,自已動手編譯了。

正文:

       Tensorflow功能代碼龐大,結構復雜;如何快速了解源碼結構,就顯示尤為重要了。

       Tensorflow主體結構:

              

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

        前端:提供編譯模型,多語言接口支持,如:python,java,c++等。

        后端:提供運行環境,完成計算圖執行,大致可分為4層:

                運行層:分布式運行時和本地運行時,負責計算圖的接收,構造,編排等;

                計算層:提供各算子的內核實現,例如: 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包裝器

        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層的實現。

Tensorflow源碼結構

       Tensorflow源碼基本按照框架分層來組織文件,如下圖:

其中目錄core是tensorflow的核心,源碼結構如下:

Session:

      Session是連接前后端的橋梁,用戶可利用session使得client能夠與master的執行引擎建立連接,通過session.run()來觸發一次計算。

      Session創建時,系統會分配一些資源,如graph引用,連接的計算引擎名稱等,所以,計算完畢后,需要使用session.close()關閉session,避免引起內存泄漏,特別是graph無法釋放問題。可以顯式調用Session.close(),使用with上下文管理器,或使用InteractiveSession()

      session之間采用共享graph方式來提高 運行效率。一個session只能運行一個graph實例,但一個graph可以運行在多個session中。在session創建時,不會重新創建graph實例,而是默認graph引用計算加1.當session close時,引用計數減1.只有引用計數為0時,graph才會被回收。

      在后端master中,根據前端client調用tf.session(target='',graph=none,config=none)時指定的target,來創建不同的session.target為要連接的tf后端執行引擎,默認為空字符串。Session創建了抽象工廠模式,如果為空字符串,則創建本地DirectSession,如果以grpc://開頭,則創建分布式grpcSession。

DirectSession只能利用本地設備,將任務創建在本地的CPU和GPU上。而grpcSession可利用遠端分布式設備,將任務創建在其他機器的CPU,GPU上,然后通過grpc協議通信。

       Session生命周期,大致有4個階段:

              1、創建:通過tf.session()創建,進行系統資源分配,特別是graph引用計數加1;

              2、運行:通過session.run()觸發計算的執行,client會將整圖graph傳遞給master,由master進行執行;

              3、關閉:通過session.close()關閉,進行系統資源的回收,graph引用計數減1;

              4、銷毀:Python垃圾回收器進行GC時,調用 session.__del__()回收。

           基本都在python的basesession中,通過swig自動生成的函數符號映射關系,調用C層的實現。

 

未完待續。。。


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM