本章中分析TensorFlow的Grappler模塊的實(shí)現(xiàn)。

Grappler

Grappler是TensorFlow的優(yōu)化模塊共螺。模塊中的主要包括這些類：

圖1：Grappler模塊主要類

tensorflow.grappler.GrapplerItem表示待優(yōu)化的TensforFlow模型，主要包括計(jì)算圖、fetch節(jié)點(diǎn)、feed節(jié)點(diǎn)姆吭。

tensorflow.grappler.Cluster表示可以運(yùn)行TensorFlow模型的硬件資源集合。一個(gè)進(jìn)程同一時(shí)間只能創(chuàng)建一個(gè)Cluster.

tensorflow.grappler.GraphOptimizer是grappler中所有優(yōu)化類的父類唁盏。

grappler模塊的調(diào)用方式如下：

圖2：grappler模塊調(diào)用過(guò)程

tensorflow.grappler.MetaOptimizer.Optimize()作為所有優(yōu)化實(shí)現(xiàn)類是入口内狸，根據(jù)優(yōu)化的配置選線決定是否調(diào)用之后的每個(gè)優(yōu)化類。

ModelPruner

tensorflow.gappler.ModelPruner類的主要優(yōu)化邏輯是裁剪計(jì)算圖厘擂，剔除不需要的節(jié)點(diǎn)昆淡。

目前版本的實(shí)現(xiàn)中， 主要剔除符合一定條件的"StopGradient"和"Identity"節(jié)點(diǎn):

/* 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");
}

以上的定義中可以看出兩個(gè)操作都是直接輸出節(jié)點(diǎn)的輸入刽严。這兩類節(jié)點(diǎn)也不是完全沒(méi)有用處的昂灵，所以ModelPruner剔除前還要檢查一些規(guī)則條件，比如明確綁定設(shè)備的這類節(jié)點(diǎn)不能剔除，有參與計(jì)算圖流控制的節(jié)點(diǎn)不能提出眨补，等等管削。

來(lái)看一個(gè)ModelPruner優(yōu)化的例子：

圖3：ModelPruner優(yōu)化前模型

圖4：ModelPruner優(yōu)化前模型

ConstantFolding

tensorflow.grappler.ConstantFolding類的主要優(yōu)化邏輯是做常量的折疊，所謂的常量折疊是將計(jì)算圖中可以預(yù)先可以確定輸出值的節(jié)點(diǎn)替換成常量撑螺，并對(duì)計(jì)算圖進(jìn)行一些結(jié)構(gòu)簡(jiǎn)化的操作含思。

tensorflow.grappler.ConstantFolding.Optimize函數(shù)主要調(diào)用了三個(gè)方法：MaterializeShapes，F(xiàn)oldGraph 和 SimplifyGraph甘晤。

目前版本中含潘，MaterializeShapes函數(shù)處理"Shape", "Size", "Rank"三類運(yùn)算節(jié)點(diǎn)：

/* 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`.
...

三類節(jié)點(diǎn)的輸出都取決與輸入Tensor的形狀，而與具體的輸入取值沒(méi)關(guān)系安皱，所以輸出可能是可以提前計(jì)算出來(lái)的调鬓。

MaterializeShapes函數(shù)將可以提前計(jì)算出輸出值的這三類節(jié)點(diǎn)全部替換成 Const 節(jié)點(diǎn)。

FoldGraph函數(shù)中做折疊計(jì)算圖的操作酌伊。如果一個(gè)節(jié)點(diǎn)的輸入都是常量，那么它的輸出也是可以提前計(jì)算的缀踪，基于這個(gè)原理不斷地用常量節(jié)點(diǎn)替換計(jì)算圖中的此類節(jié)點(diǎn)居砖，直到?jīng)]有任何可以替換的節(jié)點(diǎn)為止。

目前版本中?驴娃，SimplifyGraph函數(shù)主要處理Sum,Prod,Min,Max,Mean,Any,All這幾類運(yùn)算節(jié)點(diǎn)奏候。這幾類節(jié)點(diǎn)的共同點(diǎn)是都沿著輸入Tensor的一定維度做一定的運(yùn)算，或是求和或是求平均唇敞，等等蔗草。SimplifyGraph將符合一定條件的這幾類節(jié)點(diǎn)替換為Identity節(jié)點(diǎn)。

來(lái)看一個(gè)ConstantFolding的例子：

圖5：ConstantFolding優(yōu)化前模型

圖6：ConstantFolding優(yōu)化前模型

LayoutOptimizer

tensorflow.grappler.LayoutOptimizer類的主要優(yōu)化邏輯是改變一些運(yùn)算節(jié)點(diǎn)的輸入數(shù)據(jù)格式來(lái)提高運(yùn)算效率疆柔，這些運(yùn)算節(jié)點(diǎn)包括：

"AvgPool","AvgPoolGrad","Conv2D","Conv2DBackpropFilter","Conv2DBackpropInput",
"BiasAdd","BiasAddGrad","FusedBatchNorm","FusedBatchNormGrad",
"MaxPool","MaxPoolGrad"咒精。

這幾類節(jié)點(diǎn)的輸入數(shù)據(jù)支持NHWC和NCHW兩種格式，但是在GPU的上的運(yùn)算核實(shí)現(xiàn)上旷档，采用NNCHW格式運(yùn)算效率要高模叙，LayoutOptimizer的優(yōu)化就是將GPU設(shè)備上的這幾類節(jié)點(diǎn)的輸入格式從NHWC轉(zhuǎn)換為NCHW。

說(shuō)明：
1, 這類操作的輸入一般是一批圖片數(shù)據(jù)鞋屈，NHWC表示輸入格式為
[batch,height,width,channel],NCHW表示輸入格式為
[batch,height,width,channel]范咨。
2，之所以有兩類格式的存在厂庇，是因?yàn)樵贑PU和GPU上兩類OpKernel渠啊，
要求的最優(yōu)格式不一樣；也因?yàn)檫@個(gè)权旷，tf中默認(rèn)采用的格式是NHWC.

LayoutOptimizer 采用的優(yōu)化方法是在適當(dāng)?shù)奈恢貌迦隩ranspose運(yùn)算節(jié)點(diǎn)：

/* 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");
...

最后看一個(gè)LayoutOptimizer優(yōu)化的例子：

圖7：LayoutOptimizer優(yōu)化前模型

圖8：LayoutOptimizer優(yōu)化前模型

MemoryOptimizer

tensorflow.grappler.MemoryOptimizer的優(yōu)化邏輯是降低設(shè)備的內(nèi)存的占用替蛉。

在模型的調(diào)用計(jì)算過(guò)程中，計(jì)算產(chǎn)生的中間結(jié)果是不需要保存的，我們的目標(biāo)是得出fetch的結(jié)果灭返。

但是在模型訓(xùn)練過(guò)程中盗迟，節(jié)點(diǎn)運(yùn)算產(chǎn)生的中間結(jié)果可能是需要保留的，因?yàn)樵谟?jì)算梯度的時(shí)候需要用熙含，這就造成了設(shè)備內(nèi)存的占用問(wèn)題罚缕，而且模型越大，占用的內(nèi)存就越多怎静，而如GPU之類是設(shè)備內(nèi)存是很稀缺和寶貴的資源邮弹。MemoryOptimizer就是為了解決這個(gè)問(wèn)題的。

MemoryOptimizer采用的方法就是把一些中間結(jié)果交換到host主機(jī)的CPU內(nèi)存中蚓聘，等到需要用的時(shí)候腌乡，再交換進(jìn)設(shè)備內(nèi)存。

具體的實(shí)現(xiàn)是在計(jì)算圖中適當(dāng)?shù)奈恢貌迦胍粚?duì)相連的Identity節(jié)點(diǎn)夜牡，一個(gè)分配到HOST CPU設(shè)備与纽，另一個(gè)分配到如GPU的設(shè)備中，并設(shè)置好合適的上下游依賴塘装。這樣上游的中間計(jì)算結(jié)果就會(huì)被傳出到HOST GPU中急迂，然后在下游節(jié)點(diǎn)需要的時(shí)候通過(guò)這對(duì)節(jié)點(diǎn)交換回設(shè)備內(nèi)存中。這個(gè)邏輯比較簡(jiǎn)單蹦肴，就不做過(guò)多解釋了僚碎，來(lái)看一個(gè)優(yōu)化的例子：

圖9：MemoryOptimizer優(yōu)化前模型

<center>圖9：MemoryOptimizer優(yōu)化前模型</center>

圖10：MemoryOptimizer優(yōu)化后模型

AutoParallel

tensorflow.grappler.AutoParallel的優(yōu)化邏輯是通過(guò)重構(gòu)原來(lái)的計(jì)算圖，使得模型的訓(xùn)練過(guò)程實(shí)現(xiàn)數(shù)據(jù)并行阴幌，準(zhǔn)確的說(shuō)是多個(gè)batch的數(shù)據(jù)能并行訓(xùn)練勺阐，而不用等前一個(gè)batch訓(xùn)練完成。

實(shí)際上矛双，tensorflow的分布式模式渊抽，已經(jīng)支持多個(gè)batch同時(shí)訓(xùn)練，目的一樣背零，與AutoParallel的實(shí)現(xiàn)方式不一樣腰吟。

分布式的數(shù)據(jù)并發(fā)模式中，存在多份一樣的模型徙瓶，共享一份待訓(xùn)練的參數(shù)毛雇，而AutoParallel的優(yōu)化中，只存在一個(gè)模型侦镇，AutoParallel通過(guò)過(guò)修改模型的結(jié)構(gòu)來(lái)實(shí)現(xiàn)的并發(fā)灵疮。

下面通過(guò)一個(gè)例子來(lái)學(xué)習(xí)AutoParaller的優(yōu)化邏輯:

圖11：AutoParaller優(yōu)化前的模型

圖3是AutoParaller優(yōu)化前的模型，分號(hào)前的是節(jié)點(diǎn)的名稱壳繁，分號(hào)后面的部分是節(jié)點(diǎn)的運(yùn)算名稱震捣。

dequeue節(jié)點(diǎn)會(huì)每次從fifo節(jié)點(diǎn)中取出一定數(shù)量的數(shù)據(jù)荔棉，與constant_a相加后最后作為apply_gradient節(jié)點(diǎn)的輸入之一。apply_gradient節(jié)點(diǎn)的運(yùn)算是:

$$var -= add * learning_rate.$$

這個(gè)模型簡(jiǎn)單模擬的一下模型訓(xùn)練的過(guò)程蒿赢，add節(jié)點(diǎn)代表梯度的計(jì)算润樱，不過(guò)真實(shí)訓(xùn)練中，計(jì)算梯度的子網(wǎng)絡(luò)比這里的要復(fù)雜羡棵。

先面我們把這個(gè)模型輸入到AutoParaller中壹若，并設(shè)置并發(fā)度為2：

圖12：AutoParaller優(yōu)化后的模型

圖4為經(jīng)過(guò)AutoParaller優(yōu)化后的模型。圖中的虛線代表控制輸入依賴皂冰，實(shí)現(xiàn)代表數(shù)據(jù)輸入依賴店展。

可以看出，原始模型中的一些節(jié)點(diǎn)保留了下來(lái)秃流，例如fifo,constant_b等等赂蕴，一些節(jié)點(diǎn)被復(fù)制了一份，例如add, learning_rate,最后新添加了一些節(jié)點(diǎn)舶胀，例如AutoParaller-Div-apply_gradient, AutoParallerl-Contol-Fetch.

可以看出概说，兩份ApplyGradientDescent節(jié)點(diǎn)可以并發(fā)運(yùn)行。分別執(zhí)行運(yùn)算：

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

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