TensorFlow技術(shù)內(nèi)幕（六）：模型優(yōu)化之Grappler

本章中分析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)用過程

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)也不是完全沒有用處的，所以ModelPruner剔除前還要檢查一些規(guī)則條件，比如明確綁定設(shè)備的這類節(jié)點(diǎn)不能剔除，有參與計(jì)算圖流控制的節(jié)點(diǎn)不能提出，等等。

來看一個(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)簡化的操作。

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的形狀，而與具體的輸入取值沒關(guān)系，所以輸出可能是可以提前計(jì)算出來的。

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)。

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

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

LayoutOptimizer

tensorflow.grappler.LayoutOptimizer類的主要優(yōu)化邏輯是改變一些運(yùn)算節(jié)點(diǎn)的輸入數(shù)據(jù)格式來提高運(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。

說明：
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ì)算過程中，計(jì)算產(chǎn)生的中間結(jié)果是不需要保存的，我們的目標(biāo)是得出fetch的結(jié)果。

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

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í)候通過這對(duì)節(jié)點(diǎn)交換回設(shè)備內(nèi)存中。這個(gè)邏輯比較簡單，就不做過多解釋了，來看一個(gè)優(yōu)化的例子：

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

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

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

AutoParallel

tensorflow.grappler.AutoParallel的優(yōu)化邏輯是通過重構(gòu)原來的計(jì)算圖，使得模型的訓(xùn)練過程實(shí)現(xiàn)數(shù)據(jù)并行，準(zhǔn)確的說是多個(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通過過修改模型的結(jié)構(gòu)來實(shí)現(xiàn)的并發(fā)。

下面通過一個(gè)例子來學(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è)模型簡單模擬的一下模型訓(xùn)練的過程，add節(jié)點(diǎn)代表梯度的計(jì)算，不過真實(shí)訓(xùn)練中，計(jì)算梯度的子網(wǎng)絡(luò)比這里的要復(fù)雜。

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

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

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

可以看出，原始模型中的一些節(jié)點(diǎn)保留了下來，例如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} $$

?著作權(quán)歸作者所有,轉(zhuǎn)載或內(nèi)容合作請(qǐng)聯(lián)系作者
【社區(qū)內(nèi)容提示】社區(qū)部分內(nèi)容疑似由AI輔助生成，瀏覽時(shí)請(qǐng)結(jié)合常識(shí)與多方信息審慎甄別。
平臺(tái)聲明：文章內(nèi)容（如有圖片或視頻亦包括在內(nèi)）由作者上傳并發(fā)布，文章內(nèi)容僅代表作者本人觀點(diǎn)，簡書系信息發(fā)布平臺(tái)，僅提供信息存儲(chǔ)服務(wù)。

相關(guān)閱讀更多精彩內(nèi)容

友情鏈接更多精彩內(nèi)容