Dive into TensorFlow 系列（1）- 静态图运行原理

2022-11-11
北京
本文字数：4276 字
阅读完需：约 14 分钟

接触过 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__()方法的三个参数：

•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 类的继承关系如下图：

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 执行框架

为便于大家理解，我们先给出粗粒度的静态图执行原理如下：

1.客户端程序定义基于数据流的计算图，然后创建 session 并初始化计算引擎。

2.分布式 Master 依次完成四项工作：抽取实际执行的 subgraph、切分 subgraph 形成若干子图片段、调度子图片段至集群、每个子图片段执行前的初始化工作。

3.集群中的 Worker 节点调度子图片段中 Operation 的执行，与其他 Worker 节点通过 send/recv 节点对进行通信。

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 服务。