天天看點

Dive into TensorFlow系列(1)-靜态圖運作原理

作者:京東雲

接觸過TensorFlow v1的朋友都知道,訓練一個TF模型有三個步驟:定義輸入和模型結構,建立tf.Session執行個體sess,執行sess.run()啟動訓練。不管是因為曆史遺留代碼或是團隊保守的模組化規範,其實很多算法團隊仍在大量使用TF v1進行日常模組化。我相信很多算法工程師執行sess.run()不下100遍,但背後的運作原理大家是否清楚呢?不管你的回答是yes or no,今天讓我們一起來探個究竟。

學習靜态圖運作原理能幹什麼?掌握它對我們TF實踐中的錯誤排查、程式定制、性能優化至關重要,是必備的前置知識。

一、何為靜态圖?

衆所周知,TensorFlow程式有兩種運作選擇,即靜态圖模式與動态圖模式。

1.1 靜态圖

靜态圖采用聲明式程式設計範式(先編譯後執行),根據前端語言(如python)描述的神經網絡結構和參數資訊建構固定的靜成計算圖,靜态圖在執行期間不依賴前端語言,而是由TF架構負責排程執行,是以非常适合做神經網絡模型的部署。使用者定義的靜态圖經序列化後用GraphDef表達,其包含的資訊有:網絡連接配接、參數設定、損失函數、優化器等。

有了完整的靜态圖定義後,TF編譯器将計算圖轉化成IR(中間表示)。初始IR會經TF編譯器一系列的轉換和優化政策生成等價的計算圖。編譯器前端轉換和優化包括:自動微分、常量折疊、公共子表達式消除;編譯器後端與硬體相關,其轉換和優化包括:代碼指令生成和編譯、算子選擇、記憶體配置設定、記憶體複用等。

綜上所述,靜态圖的生成過程可用下圖簡要概适:

Dive into TensorFlow系列(1)-靜态圖運作原理



1.2 動态圖

動态圖采用指令式程式設計範式,即編譯與執行同時發生。動态圖采用前端語言的解釋器對使用者代碼進行解析,然後利用TF架構的算子分發功能,使得算子立即執行并向前端傳回計算結果。當模型接收輸入資料後,TF開始動态生成圖拓撲結構,添加輸入節點并将資料傳輸給後續節點。如果動态圖中含有條件控制邏輯,會立即計算邏輯判斷結果并确定後續資料流向,是以動态圖完整的拓撲結構在執行前是未知的。另外,當模型根據新的batch訓練時,原有的圖結構則失效,必須根據輸入和控制條件重新生成圖結構。

綜上所述,動态圖生成過程可用下圖簡要概括:

Dive into TensorFlow系列(1)-靜态圖運作原理



1.3 比較

為了友善大家深入了解動/靜态圖原理及異同點,梳理相關資訊如下表:

特性 靜态圖 動态圖
即時擷取中間結果
代碼調試難度
控制流實作方式 TF特定文法 可采用前端語言文法
性能 多種優化政策,性能好 優化受限,性能差
記憶體占用 記憶體占用少 記憶體占用多
部署情況 可直接部署 不可直接部署

二、Session是幹啥的?

2.1 Session定義

tf.Session代表使用者程式和C++運作時之間的連接配接。一個Session類對象session可以用來通路本機計算裝置,也可通路TF分布式運作時環境中的遠端裝置。session也能緩存tf.Graph資訊,使得相同計算邏輯的多次執行得以高效實作。

tf.Session的構造方法定義如下:

def __init__(self, target='', graph=None, config=None):
    """Creates a new TensorFlow session.
    If no `graph` argument is specified when constructing the session,
    the default graph will be launched in the session. If you are
    using more than one graph (created with `tf.Graph()` in the same
    process, you will have to use different sessions for each graph,
    but each graph can be used in multiple sessions. In this case, it
    is often clearer to pass the graph to be launched explicitly to
    the session constructor.
    Args:
      target: (Optional.) The execution engine to connect to.
        Defaults to using an in-process engine. See
        @{$distributed$Distributed TensorFlow}
        for more examples.
      graph: (Optional.) The `Graph` to be launched (described above).
      config: (Optional.) A [`ConfigProto`](https://www.tensorflow.org/code/tensorflow/core/protobuf/config.proto)
        protocol buffer with configuration options for the session.
    """
    super(Session, self).__init__(target, graph, config=config)
    # NOTE(mrry): Create these on first `__enter__` to avoid a reference cycle.
    self._default_graph_context_manager = None
    self._default_session_context_manager = None           

我們來看一下__init__()方法的三個參數:

•target:預設為空,代表session僅可通路本機上的計算裝置。如果設定grpc://樣式的URL,則可以通路TF server對應機器的計算裝置。

•graph:預設執行目前default graph中的op算子。如果使用者程式中包含多個計算圖,則在建立session時必須指定是哪個計算圖。

•config:通過指定tf.ConfigProto來控制session的行為。常見的配置選項有:裝置退化allow_soft_placement、分布式叢集配置cluster_def、圖優化政策graph_options.optimizer_options、GPU記憶體逐漸增長gpu_options.allow_growth。

2.2 Session.run()

tf.Session.run()實際是調用tf.BaseSession.run()方法,其函數簽名如下:

def run(self, fetches, feed_dict=None, options=None, run_metadata=None):           

run()方法的參數說明如下:

•fetches:指定要執行的tf.Operation或評估的tf.Tensor,可以是單個元素或是清單、字典。

•feed_dict:一個占位符到填充值的映射。

•options:RunOptions的protocol buffer。

•run_metadata:RunMetadata的protocol buffer,用來收集執行過程的中繼資料資訊。

當Session指定fetches後,根據要擷取的結果決定tf.Graph實際執行的subgraph(并非整個tf.Graph都要執行)。執行靜态圖還有三個要點:

•訓練階段用一個靜态圖,而預測/評估階段用另一個靜态圖。

•一個session執行個體隻能運作一個graph執行個體,但一個graph可以運作在多個session中。

•session之間可通過共享graph的方式來提高運作效率。建立session時若不指定graph,則運作的是default graph。如果新建立的session也不指定graph,則隻需要對default graph的引用計數加1即可;當此session close時,default graph引用計數減1。

2.3 Session類前後端設計

首先我們看一下和使用者直接打交道的前端Session,具體分為普通Session和互動式InteractiveSession。前者全稱為tf.Session,需要在啟動之前先建構完整的計算圖;後者全稱為tf.InteractiveSession,它是先建構一個session,然後再定義各種操作,适用于shell和IPython等互動式環境。這兩個類均繼承自BaseSession,這個基類實作了整個生命周期的所有會話邏輯(相關代碼在tensorflow/python/client/session.py中)。前端Session類的繼承關系如下圖:

Dive into TensorFlow系列(1)-靜态圖運作原理



TensorFlow後端會根據前端tf.Session(target='', graph=None, config=None)建立時指定的target來建立不同的後端Session。target是要連接配接的TF後端執行引擎,預設為空字元串。後端Session的建立采用抽象工廠模式,如果為空字元串,則建立本地DirectionSession;如果是grpc://開頭的URL串,則建立分布式GrpcSession。DirectSession隻能利用本地裝置,将任務排程到本地的CPU/GPU裝置上;GrpcSession可利用遠端裝置,将任務分發到不同機器的CPU/GPU上,然後機器之間通過gRPC進行通信。顯而易見,DirectionSession的定義應在core/common_runtime/direction_session.h中;GrpcSession的定義在core/distributed_runtime/rpc/grpc_session.h中。後端Session的類圖關系如下所示:

Dive into TensorFlow系列(1)-靜态圖運作原理



三、靜态圖執行過程

3.1 執行架構

為便于大家了解,我們先給出粗粒度的靜态圖執行原理如下:

1.用戶端程式定義基于資料流的計算圖,然後建立session并初始化計算引擎。

2.分布式Master依次完成四項工作:抽取實際執行的subgraph、切分subgraph形成若幹子圖檔段、排程子圖檔段至叢集、每個子圖檔段執行前的初始化工作。

3.叢集中的Worker節點排程子圖檔段中Operation的執行,與其他Worker節點通過send/recv節點對進行通信。

Dive into TensorFlow系列(1)-靜态圖運作原理



3.2 若幹執行細節

靜态圖的實際執行過程要比3.1節描述的複雜得多。由于本篇的初衷不是做源碼的完整剖析,是以我們僅就Client向Master的處理過程做詳細說明,旨在讓讀者親身體會一下互動過程的複雜性。

Client建立GrpcSession,控制Client會話的生命周期;Master運作時被MasterSession控制。GrpcSession通過抽象工廠模式得到,首先得到工廠類GrpcSessionFactory的對象,并用SessionFactory句柄factory存儲。然後通過factory的多态方法生成GrpcSession,如果target為grpc://的話。Master本質上是一個Server,每個Server均有一個MasterService和一個WorkerService。

Client通過GrpcSession調用Master節點的MasterService,這個過程需借助MasterInterface才可完成。MasterInterface用來和MasterService進行通信,它有兩種不同的場景實作:

•如果Client和Master在同一個程序中,則用LocalMaster實作程序内的直接通信。

•GrpcRemoteMaster則使用gRPC來和MasterService進行通信,此時Master和Client在兩個不同的程序中。GrpcRemoteMaster的角色是gRPC用戶端,它通過stub通路遠端Master節點上的MasterService服務。

Dive into TensorFlow系列(1)-靜态圖運作原理



如果讀者想對上述過程做更為深入的了解,可以參考幾個關鍵類的源碼:

•GrpcSession:core/distributed_runtime/rpc/grpc_session.h

•LocalMaster:core/distributed_runtime/local_master.h

•GrpcRemoteMaster:core/distributed_runtime/rpc/grpc_remote_master.cc

•GrpcMasterService:core/distributed_runtime/rpc/grpc_master_service.cc

其實Client到Master的處理過程還涉及MasterSession的建立,以及GrpcSession與MasterSession的互動與辨別問題。篇幅所限,不展開了。

四、總結

作為Dive into TensorFlow系列第一講,本文由淺入深、系統講解了靜态圖及其運作原理,以及支撐這些功能的架構設計與部分源碼解析。回到文章開頭提到的使用者讀懂全文能有什麼收益?(嘗試提幾點)

•明白預設session能運作預設靜态圖的原理,及常見的錯誤排查與調試方法。能根據場景需要靈活選擇動/靜态圖計算模式。

•如果一個靜态圖由幾個獨立子圖建構,我們建議對每個子圖分别建構tf.Graph對象。

•了解3.1小節對後續深入掌握op的placement、graph partition、基于gRPC的send/recv算子對做程序間通信有方向性指引作用。

作者:李傑

參考資料

1.Graphs and Sessions: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/graphs.md

2.《機器學習系統:設計與實作》: https://openmlsys.github.io/chapter_computational_graph/index.html

3.前後端連接配接的橋梁Session: https://www.likecs.com/show-306440850.html

4.TensorFlow v1.15.5源碼: https://github.com/tensorflow/tensorflow/tree/v1.15.5/tensorflow/core/graph

5.TensorFlow Architecture: https://github.com/tensorflow/docs/blob/master/site/en/r1/guide/extend/architecture.md

6.TensorFlow分布式環境Session: https://www.cnblogs.com/rossiXYZ/p/16065124.html

繼續閱讀