TensorFlow 核心流程剖析 -- 2 神经网络模型的构建、分割和优化

与本章节相关的一些关键术语
模型的生成总体流程
graph的构建分割优化c部分
- Session的创建
- full_graph实例的创建和优化
- 根据feedfetch对full_graph进行修改和裁剪
- 根据device对full_graph_subset分割
  - 分割之前的完整模型
  - 分割之后的partitioned_graph

与本章节相关的一些关键术语

graph

我们知道，在tensorflow里，模型是以compuatation graph的形式存在，作为训练和inference的载体。下面简称graph。
graph的组成：
- node：即定义一个具体的计算操作，比如Add， MatMul，Conv等。每个node可以定义多种属性，包括它的计算操作（叫做Op），输入、输出数据类型，与计算相关的参数设置，比如Convolution时需要的padding\stride.
- tensor：是一个node输出的计算结果，比如MatMul计算完成后，输出一个多维的张量，就是一个tensor，tensor可以作为下一个node的数据输入。
- edge：连接各个node，目前有两种edege:
  - data_flow: 负责在node之间传递tensor数据
  - control_flow: 负责确定node之间的执行依赖关系
    多个node和他们之间的edge连接，就构成了一个graph，完整描述我们定义的神经网络模型。
GraphDef

用于定义graph的ProtoBuf协议格式，因其文本属性，可以将graph以这种格式保存至文本文件，实现训练模型的保存。也可以很方面的在不同设备、软件模块之间传输和解析GraphDef。另外，tensorboard可以读取GraphDef格式保存的文本文件，显示graph。
session

运行graph的主体。负责创建和管理graph及其所需的运行设备（device）资源。
device

运行graph的硬件资源，属于session，如本地的gpu device/ cpu device。
Executor

graph的具体执行者，属于session，当我们把graph分割到多个device执行的时候，也会生成对应的多个Executor实例。

模型的生成总体流程

graph的生成，源自通过python API中对graph中nodes的定义。

一般来讲，我们通过python API这样开始训练一个model：

定义graph和其中的node
创建session去Run这个graph

Graph的生成总体流程如下图：

graph的构建、分割、优化（c++部分）

graph的执行是在c++代码中完成的，在执行前，需要对graph进行构建、分割、优化。
而在开始构建一个graph之前，我们必须先创建一个session，作为运行这个graph的主体。

Session的创建

python API调用c++ API , 创建合适的session，用于运行graph。

运行graph的session有两类：

DirectSession：使用本地的devic作为运行资源，如本机中的gpu、cpu。可以将graph分割、分布到多个devices上运行，实现devices之间的并发执行，同样可以实现data parallelism/ model parallelism 这两种分布式训练。配置比较简单，我们主要结合这种session进行举例和讲解。
GrpcSession：使用远程主机的device作为计算资源，grpc作为远程调用的机制。使用cluster进行分布式训练的场景就需要使用这种session。实现上，graph按照worker分割的原理，类似于DirectSession按照device分割的原理, 每个worker上sub-graph的执行者也是Executor，这里不探讨。

为方便感兴趣的朋友进一步查看代码，列出主要函数调用路径：

Created with Raphaël 2.1.0c_api.ccc_api.ccsession.ccsession.ccTF_NewDeprecatedSession()NewSession()

full_graph实例的创建和优化

python部分定义好的graph，通过TF_ExtendGraph()接口，将graph以protobuf定义的 GraphDef格式，传递到c++ session中。
同步graph有两种场景：

graph增量式更新后，主动同步到c++ session；
每次session run时，都会先同步graph到c++ session，以保证run的是最新的graph；

主要函数调用路径：

Created with Raphaël 2.1.0c_api.ccc_api.ccdirect_session.ccdirect_session.ccsimple_graph_execution_state.ccsimple_graph_execution_state.ccTF_ExtendGraph()Extend()InitBaseGraph()SimpleGraphExecutionState::Extend()

full_graph生成和优化流程如下图所示：
注：之所以叫做full_graph,是因为此时还没有对graph根据feed/fetch裁剪、以及分布式的切割。

c++ session 接收到GraphDef后, 主要步骤描述如下：

根据是否是第一次创建graph, 决定是否进行extend。目前extend的版本，只支持增加node，不能修改、删除session中已绑定的graph。通过这个功能，我们可以很方便的对graph进行增量式的搭建。
直接对GraphDef进行优化，涉及到增加/删除部分node、GraphDef的重构。优化的策略目前主要有：
- pruning: 裁剪掉运行graph时无用的node，如StopGradient，Identity。
- layout：针对运行在cuda GPU上的DNN node，如MaxPool,Conv2D等，将input格式为NHWC的这些node转换为input为NCHW格式，以提高处理效率，比如cuDNN的大部分kernel在计算时，使用NCHW格式的tensor更为高效。
- constfolding：针对所有输入都为Constant OP的node，提前运行得到输出：创建临时kernel，输入各Contant的值，运行得到结果作为新的Constant，再重构图，以使session run时不会再运行上述node。
- memory：增加Identity nodes，将指定的tensor，swap to host memory，比如可以在CPU上debug 查看某个位于GPU memory的tensor的值。
- auto-parallel: 创建full_graph的replica，并将replicas分布到各个GPU上，实现data parallel。
具体算法可以参见函数 RunMetaOptimizer() in meta_optimizer.cc。
创建full_graph: 基于优化过的GraphDef，生成Graph实例，即full_graph，包含所有的node/edges。
node的部署(placement)：将所有node，根据用户创建node时指定的device、或者SimplePlacer策略，将合适的device分配给node。可用的devices列表由direct session在生成时创建，device可以是本地GPU/CPU，也可以是cluster中的worker job、ps job等，后面章节会单独讲解device。每个node，将会使用分配的device上的相应kernel实例来进行计算，所以在分配时，还需验证在该device上否具备node使用的kernel实例。目前开源的SimplePlacer策略比较简单，比如相邻node共device等。后续应发布根据node的计算资源、存储资源消耗，进行负载均衡等多种部署策略。

根据feed/fetch对full_graph进行修改和裁剪

每次session run，需要指定feed和fetch 的tensors，我们知道，tensor是node的输出，所以可以根据tensor定位到full_graph中的feed/fetch nodes：以feed nodes作为起点，fetch nodes作为终点，圈定的full_graph要执行的范围，并不一定等同于整个full_graph，比如我们可以指定full_graph的中间节点作为本次run的起点、终点。
通过裁剪，可以提高执行的效率，从而确定本次run要执行的full_graph的哪一部分，将不需要执行的部分剪掉，提取需要执行的部分，整个动作就是full_graph的裁剪。
主要有两个步骤：
1。 feed／fetch node的修改：为feed 增加Recv node、Source node，为fetch 增加Send node、Sink node，同时修改增加的node、feed/fetch node 与上下游node的关系。Source node 即作为full_graph 本次run的起点，Sink node为终点。
2。裁减：从fetch node（+ target node，如果有指定）出发，广度优先的方式沿着data_flow edge/control_flow edge 逆向（记住，data_flow/control_flow是有向连接）遍历full_graph，没有覆盖到的node即为对本次fetch tensor无关的node，删除这些node，完成对full_graph的裁减。

为了便于讲解，我们把这时候得到的graph叫做full_graph_subset.

主要函数调用路径：

Created with Raphaël 2.1.0c_api.ccc_api.ccdirect_session.ccdirect_session.ccsimple_graph_execution_state.ccsimple_graph_execution_state.ccsubgraph.ccsubgraph.ccTF_Run()Run()GetOrCreateExecutors()CreateGraphs()BuildGraph()RewriteGraphForExecution()

根据device对full_graph_subset分割

这时候我们得到的本次需要run的full_graph_subset，其中每个node都已经包含了分配的device信息，接下来我们要把full_graph_subset分割为多个Graph实例，每一个Graph实例与每一个具体device一一对应。分割后得到的各graph，我们叫做partitioned_graph。

下面举例，用一个简单的模型：

layer1 + layer2 + softmax_loss，其中layer1 和 layer2 架构一样：W*x + b -> Relu
由于没有训练需求，不添加BP的node。
我们 layer1 定义为在cpu:0上运行，layer2 和 softmax_loss layer在 gpu:0 上运行。
feed tensor有[x, y_], fetch tensor是 softmax_loss/Mean node的输出tensor。

我们可以把上述主要步骤中得到的各个graph实例，都逐一转换为GraphDef格式保存至文本文件，再用tensorboard来看一下分割的结果：

有必要先讲一下分割graph时，为跨devices的edge，增加Send/Recv node 和Rendezvous(汇合点)的机制，以实现cpu<->gpu, gpu<->gpu等场景的数据传输，如下图：

分割之前的完整模型：

以layer1中的MatMul这个node为例，查看其Device属性，可以看到位于指定的cpu:0。

查看 softmax_loss layer中的Mean node，可以看到其位于指定的gpu:0 。

分割之后的partitioned_graph:

先看一下分割到cpu:0 上的partitioned_graph，可以看到为 cpu:0 <-> cpu:0 数据传输创建的send／recv nodes，以及创建的集合点 Rendezvous。send／recv nodes 通过Rendezvous, 与gpu:0 传输数据。
同时还可以看到为feed／fetch创建的 send／recv nodes: 这里相当于定义了一个清晰的边界，用户程序指定的feed / fetch tensors，必须通过Rendezvous 传递到graph内部。注意: feed/fetch tensor 只能在cpu:0 上。

再看一下分割到gpu:0 上的 partitioned_graph, 可以看到同样的机制。这里的Rendezvous 和cpu:0 上使用的是同一个实例，实现了gpu:0 <-> cpu: 0 数据传输。

接下来会为每个partitioned_graph创建local executor,作为其本地的执行者。

上述例子对应的处理流程如下：

主要函数调用路径：

Created with Raphaël 2.1.0c_api.ccc_api.ccdirect_session.ccdirect_session.ccgraph_partition.ccgraph_partition.ccTF_Run()Run()GetOrCreateExecutors()CreateGraphs()Partition()AddControlFlow()AddSend()AddRecv()

对于原来full_graph_subset改动比较大的，有两个地方：

为while_loop的跨devices运行增加control_loop. 这个逻辑比较繁琐，有兴趣的朋友可以直接看代码AddControlFlow()。
为跨devices的edge，增加Send/Recv node 和Rendezvous(汇合点)机制。

至此， graph已经做好了运行的准备。后续文章再讲解executor并发执行的机制。