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編譯器一系列的轉換和優化策略生成等價的計算圖，編譯器前端轉換和優化包括：自動微分、常量折疊、公共子運算式消除；編譯器后端與硬體相關，其轉換和優化包括：代碼指令生成和編譯、算子選擇、記憶體分配、記憶體復用等，

綜上所述，靜態圖的生成程序可用下圖簡要概適：

1.2 動態圖

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

綜上所述，動態圖生成程序可用下圖簡要概括：

1.3 比較

為了方便大家深入理解動/靜態圖原理及異同點，梳理相關資訊如下表：

二、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__()方法的三個引數：

2.2 Session.run()

tf.Session.run()實際是呼叫tf.BaseSession.run()方法，其函式簽名如下：

def run(self, fetches, feed_dict=None, options=None, run_metadata=https://www.cnblogs.com/Jcloud/p/None):

run()方法的引數說明如下：

當Session指定fetches后，根據要獲取的結果決定tf.Graph實際執行的subgraph（并非整個tf.Graph都要執行），執行靜態圖還有三個要點：

2.3 Session類前后端設計

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

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的類圖關系如下所示：

三、靜態圖執行程序

3.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的處理程序還涉及MasterSession的創建，以及GrpcSession與MasterSession的互動與標識問題，篇幅所限，不展開了，

四、總結

作為Dive into TensorFlow系列第一講，本文由淺入深、系統講解了靜態圖及其運行原理，以及支撐這些功能的架構設計與部分原始碼決議，回到文章開頭提到的用戶讀懂全文能有什么收益？（嘗試提幾點）

參考資料

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

標籤：其他

上一篇：HTTP報錯的狀態碼原因

下一篇：帶你從0到1開發AI影像分類應用