TensorFlow技术内幕（六）：模型优化之Grappler

本章中分析TensorFlow的Grappler模块的实现。代码位于tensorflow/core/grappler。

上一章中分析session类的时候，已经介绍过了grappler模块的调用时机。

Grappler

Grappler是TensorFlow的优化模块。模块中的主要包括这些类：

这里写图片描述

图1：Grappler模块主要类

tensorflow.grappler.GrapplerItem表示待优化的TensforFlow模型，主要包括计算图、fetch节点、feed节点。

tensorflow.grappler.Cluster表示可以运行TensorFlow模型的硬件资源集合。一个进程同一时间只能创建一个Cluster.

tensorflow.grappler.GraphOptimizer是grappler中所有优化类的父类。

grappler模块的调用方式如下：

这里写图片描述

图2：grappler模块调用过程

tensorflow.grappler.MetaOptimizer.Optimize()作为所有优化实现类是入口，根据优化的配置决定是否调用之后的每个优化类。

ModelPruner

tensorflow.gappler.ModelPruner类的主要优化逻辑是裁剪计算图，剔除不需要的节点。

目前版本的实现中，主要剔除符合一定条件的”StopGradient”和”Identity”节点:

/* tensorflow/core/ops/array_ops.cc */
...
REGISTER_OP("StopGradient")
    .Input("input: T")
    .Output("output: T")
    .Attr("T: type")
    .SetShapeFn(shape_inference::UnchangedShape)
    .Doc(R"Doc(
Stops gradient computation.

When executed in a graph, this op outputs its input tensor as-is.
...

REGISTER_OP("Identity")
    .Input("input: T")
    .Output("output: T")
    .Attr("T: type")
    .SetShapeFn([](shape_inference::InferenceContext* c) {
      c->set_output(0, c->input(0));
      auto* handle_data = c->input_handle_shapes_and_types(0);
      if (handle_data != nullptr) {
        c->set_output_handle_shapes_and_types(0, *handle_data);
      }
      return Status::OK();
    })
    .Doc(R"Doc(
Return a tensor with the same shape and contents as the input tensor or value.
)Doc");
}

以上的定义中可以看出两个操作都是直接输出节点的输入。这两类节点也不是完全没有用处的，所以ModelPruner剔除前还要检查一些规则条件，比如明确绑定设备的这类节点不能剔除，有参与计算图流控制的节点不能提出，等等。

来看一个ModelPruner优化的例子：

![这里写图片描述](https://img-blog.csdn.net/20180608132450632?watermark/2/text/aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dhb2ZlaXBhb3Bhb3Rhbmc=/font/5a6L5L2T/fontsize/400/fill/I0JBQkFCMA==/dissolve/70) 图3：ModelPruner优化前模型 ![这里写图片描述](https://img-blog.csdn.net/2018060813250788?watermark/2/text/aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dhb2ZlaXBhb3Bhb3Rhbmc=/font/5a6L5L2T/fontsize/400/fill/I0JBQkFCMA==/dissolve/70) 图4：ModelPruner优化前模型

ConstantFolding

tensorflow.grappler.ConstantFolding类的主要优化逻辑是做常量的折叠，所谓的常量折叠是将计算图中可以预先可以确定输出值的节点替换成常量，并对计算图进行一些结构简化的操作。

tensorflow.grappler.ConstantFolding.Optimize函数主要调用了三个方法：MaterializeShapes，FoldGraph 和 SimplifyGraph。

目前版本中，MaterializeShapes函数处理”Shape”, “Size”, “Rank”三类运算节点：

/* tensorflow/core/ops/array_ops.cc */
...

REGISTER_OP("Shape")
    .Input("input: T")
    .Output("output: out_type")
    .Attr("T: type")
    .Attr("out_type: {int32, int64} = DT_INT32")
    .SetShapeFn(ShapeShapeFn)
    .Doc(R"doc(
Returns the shape of a tensor.

This operation returns a 1-D integer tensor representing the shape of `input`.

...

REGISTER_OP("Size")
    .Input("input: T")
    .Output("output: out_type")
    .Attr("T: type")
    .Attr("out_type: {int32, int64} = DT_INT32")
    .SetShapeFn(shape_inference::ScalarShape)
    .Doc(R"doc(
Returns the size of a tensor.

This operation returns an integer representing the number of elements in
`input`.
...

REGISTER_OP("Rank")
    .Input("input: T")
    .Output("output: int32")
    .Attr("T: type")
    .SetShapeFn(shape_inference::ScalarShape)
    .Doc(R"doc(
Returns the rank of a tensor.

This operation returns an integer representing the rank of `input`.
...

三类节点的输出都取决与输入Tensor的形状，而与具体的输入取值没关系，所以输出可能是可以提前计算出来的。

MaterializeShapes函数将可以提前计算出输出值的这三类节点全部替换成 Const 节点。

FoldGraph函数中做折叠计算图的操作。如果一个节点的输入都是常量，那么它的输出也是可以提前计算的，基于这个原理不断地用常量节点替换计算图中的此类节点，直到没有任何可以替换的节点为止。

目前版本中，SimplifyGraph函数主要处理Sum,Prod,Min,Max,Mean,Any,All这几类运算节点。这几类节点的共同点是都沿着输入Tensor的一定维度做一定的运算，或是求和或是求平均，等等。SimplifyGraph将符合一定条件的这几类节点替换为Identity节点。

来看一个ConstantFolding的例子：

![这里写图片描述](https://img-blog.csdn.net/20180608132524230?watermark/2/text/aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dhb2ZlaXBhb3Bhb3Rhbmc=/font/5a6L5L2T/fontsize/400/fill/I0JBQkFCMA==/dissolve/70) 图5：ConstantFolding优化前模型 ![这里写图片描述](https://img-blog.csdn.net/20180608132536866?watermark/2/text/aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2dhb2ZlaXBhb3Bhb3Rhbmc=/font/5a6L5L2T/fontsize/400/fill/I0JBQkFCMA==/dissolve/70) 图6：ConstantFolding优化前模型

LayoutOptimizer

tensorflow.grappler.LayoutOptimizer类的主要优化逻辑是改变一些运算节点的输入数据格式来提高运算效率，这些运算节点包括：

“AvgPool”,”AvgPoolGrad”,”Conv2D”,”Conv2DBackpropFilter”,”Conv2DBackpropInput”,
“BiasAdd”,”BiasAddGrad”,”FusedBatchNorm”,”FusedBatchNormGrad”,
“MaxPool”,”MaxPoolGrad”。

这几类节点的输入数据支持NHWC和NCHW两种格式，但是在GPU的上的运算核实现上，采用NNCHW格式运算效率要高，LayoutOptimizer的优化就是将GPU设备上的这几类节点的输入格式从NHWC转换为NCHW。

说明：
1, 这类操作的输入一般是一批图片数据，NHWC表示输入格式为
[batch,height,width,channel],NCHW表示输入格式为
[batch,height,width,channel]。
2，之所以有两类格式的存在，是因为在CPU和GPU上两类OpKernel，
要求的最优格式不一样；也因为这个，tf中默认采用的格式是NHWC.

LayoutOptimizer 采用的优化方法是在适当的位置插入Transpose运算节点：

/* tensorflow/core/ops/array_ops.cc */
...

REGISTER_OP("Transpose")
    .Input("x: T")
    .Input("perm: Tperm")
    .Output("y: T")
    .Attr("T: type")
    .Attr("Tperm: {int32, int64} = DT_INT32")
    .SetShapeFn([](InferenceContext* c) {
      ShapeHandle input = c->input(0);
      ... ...
      )
    .Doc(R"doc(
Shuffle dimensions of x according to a permutation.

The output `y` has the same rank as `x`. The shapes of `x` and `y` satisfy:
  `y.shape[i] == x.shape[perm[i]] for i in [0, 1, ..., rank(x) - 1]`
)doc");
...

最后看一个LayoutOptimizer优化的例子：

这里写图片描述

图7：LayoutOptimizer优化前模型

这里写图片描述

图8：LayoutOptimizer优化前模型

MemoryOptimizer

tensorflow.grappler.MemoryOptimizer的优化逻辑是降低设备的内存的占用。

在模型的调用计算过程中，计算产生的中间结果是不需要保存的，我们的目标是得出fetch的结果。

但是在模型训练过程中，节点运算产生的中间结果可能是需要保留的，因为在计算梯度的时候需要用，这就造成了设备内存的占用问题，而且模型越大，占用的内存就越多，而如GPU之类是设备内存是很稀缺和宝贵的资源。MemoryOptimizer就是为了解决这个问题的。

MemoryOptimizer采用的方法就是把一些中间结果交换到host主机的CPU内存中，等到需要用的时候，再交换进设备内存。

具体的实现是在计算图中适当的位置插入一对相连的Identity节点，一个分配到HOST CPU设备，另一个分配到如GPU的设备中，并设置好合适的上下游依赖。这样上游的中间计算结果就会被传出到HOST GPU中，然后在下游节点需要的时候通过这对节点交换回设备内存中。这个逻辑比较简单，就不做过多解释了，来看一个优化的例子：

这里写图片描述

图9：MemoryOptimizer优化前模型

这里写图片描述

图10：MemoryOptimizer优化后模型

AutoParallel

tensorflow.grappler.AutoParallel的优化逻辑是通过重构原来的计算图，使得模型的训练过程实现数据并行，准确的说是多个batch的数据能并行训练，而不用等前一个batch训练完成。

实际上，tensorflow的分布式模式，已经支持多个batch同时训练，目的一样，与AutoParallel的实现方式不一样。

分布式的数据并发模式中，存在多份一样的模型，共享一份待训练的参数，而AutoParallel的优化中，只存在一个模型，AutoParallel通过过修改模型的结构来实现的并发。

下面通过一个例子来学习AutoParaller的优化逻辑:

这里写图片描述

图11：AutoParaller优化前的模型

图3是AutoParaller优化前的模型，分号前的是节点的名称，分号后面的部分是节点的运算名称。

dequeue节点会每次从fifo节点中取出一定数量的数据，与constant_a相加后最后作为apply_gradient节点的输入之一。apply_gradient节点的运算是:

v a r - = a d d * l e a r n i n g_r a t e .

$var -= add * learning\_rate.$

这个模型简单模拟的一下模型训练的过程，add节点代表梯度的计算，不过真实训练中，计算梯度的子网络比这里的要复杂。

先面我们把这个模型输入到AutoParaller中，并设置并发度为2：

这里写图片描述

图12：AutoParaller优化后的模型

图4为经过AutoParaller优化后的模型。图中的虚线代表控制输入依赖，实现代表数据输入依赖。

可以看出，原始模型中的一些节点保留了下来，例如fifo,constant_b等等，一些节点被复制了一份，例如add, learning_rate,最后新添加了一些节点，例如AutoParaller-Div-apply_gradient, AutoParallerl-Contol-Fetch.

可以看出，两份ApplyGradientDescent节点可以并发运行。分别执行运算：

v a r - = \frac{a d d_{(a u t o p a r a l l e l - r e p l i c a - 0)} * l e a r n i n g_r a t e_{(a u t o p a r a l l e l - r e p l i c a - 0)}}{2}

$var -= \frac{add_{(autoparallel-replica-0)} * learning\_rate_{(autoparallel-replica-0)}}{2}$

v a r - = \frac{a d d_{(a u t o p a r a l l e l - r e p l i c a - 1)} * l e a r n i n g_r a t e_{(a u t o p a r a l l e l - r e p l i c a - 1)}}{2}

$var -= \frac{add_{(autoparallel-replica-1)} * learning\_rate_{(autoparallel-replica-1)}}{2}$