Press "Enter" to skip to content

【从零开始学深度学习编译器】番外一,Data Flow和Control Flow

 

本文作为从零开始学深度学习编译器的番外篇,介绍了一下深度学习框架的Data Flow和Control Flow,并基于TensorFlow解释了TensorFlow是如何在静态图中实现Control Flow。其实,目前除TensorFlow外其它DL框架基本上是不支持Control Flow的,以Pytorch为例,它支持的也仅仅是Python Control Flow,即在Python层写控制逻辑。这样就存在一个问题,如果我们要部署带Control Flow的模型就会比较困难,比如考虑一下如何部署RNN模型到不支持Python的设备上?对于Pytorch来说可能需要借助JIT+libtorch?但推理的效率可能无法得到保证,总之是一件比较麻烦的事情,再想想其它的一些框架似乎都不是很好处理这种情况。事实上TVM考虑到了这一点并推出了NNVM的二代Relay,使用Lambda演算来解决Control Flow的问题,等后续再继续学习吧。

 

这篇文章主要基于Google的SAM ABRAHAMS做的一个TensorFlow的Control Flow的分析,PPT可以在QQ群获取。

 

0x0. 前言

 

本来是想在讲TVM Relay的时候提一下DataFlow和ControlFlow的,但是担心读者看到解析代码的文章打开就关了,所以这里用一篇简短的文章来介绍一下深度学习框架中的DataFlow和ControlFlow。这也是TVM引入Relay的一个关键原因,希望读者能 把握住

 

0x1. DataFlow

 

我记得我接触的第一个深度学习框架是TensorFlow1.x,本科毕业设计也是基于TensorFlow完成的,因此这里我将以TensorFlow1.x为例介绍一下DataFlow。

 

假设现在我们要实现一个的逻辑,其中,,都是一个简单的实数,然后我们如果用Python来实现非常简单,大概长这样:

 

#coding=utf-8
import os
def cal(a, b, c):
    res = (a + b) * c
    print(res)
return res
print(cal(1.0, 2.0, 3.0))

 

输出结果是9.0。然后我们使用tf1.31.1同样来实现这个过程:

 

import tensorflow as tf
def cal(a, b, c):
    add_op = a + b
    print(add_op)
    mul_op = add_op * c
    init = tf.global_variables_initializer()
    sess = tf.Session()
    sess.run(init)
    mul_op_res = sess.run([mul_op])
return mul_op_res
a = tf.constant(1.0)
b = tf.constant(2.0)
c = tf.constant(3.0)
print(cal(a, b, c))

 

同样代码的输出是9.0。然后这两个示例是为了解释像TensorFlow这种框架,它的计算图是一个计算流图,是由数据来驱动的。在上面的程序中我们可以发现如果打印 add_op
我们获得的结果是一个 Tensor

 

Tensor("add:0", shape=(), dtype=float32

 

这是因为,TensorFlow1.x实现的这个计算函数首先在内存中构造了一个数据流图,长这样:

上面tensorflow程序对应的数据流图

我们回看一下Python的实现,实际上在执行 res = (a + b) * c
这行代码时,已经计算出了 res
的值,因为Python这种过程语言的数学计算是由代码驱动的。而TensorFlow不一样,它首先构造了数据流图,然后对这个计算流图进行绑定数据,让这个数据在这个图里面流起来,这是显示调用 sess.run
来获得输出的。

 

像TensorFlow这种基于数据流图(DataFlow)进行计算的深度学习框架不少,比如早期的Theano,2020年开源的国内深度学习框架 OneFlow
,PaddlePaddle1.x 初级版本都是基于数据流图的。当然更多人将其称为静态图。

 

0x2. Control Flow

 

实际上大多数深度学习框架都不能支持Control Flow,例如Pytorch要写控制流,可能就需要开发者在Python端去写Python Control Flow。这一节我将结合TensorFlow的Control来为大家解析一下Control Flow的难点以及TensorFlow的一些解决方案。这里的内容理解主要基于这篇博客(https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/),感兴趣的同学可以去查看原文。

 

在计算机科学中,控制流(Control Flow)定义了独立语句,指令,函数调用等执行或者求值的顺序。举个例子,我们要实现一个本机控制流,即我们需要根据函数A的输出值选择运行函数B或者C中的一个:

一个Control Flow的例子

然后要实现这个控制流,最Naive的方式在是Python端写If/Else语句,即Python端的Control Flow,然后在不同条件下使用session.run()来求取不同分支的值。对于TensorFlow来说,大概是这样:

这里获取A的值只是将其反馈回来

然后这个Python层的Control Flow并不会在计算图中被表示出来,即:

黄色部分在计算图中实际上是被删掉了,因为早期的TensorFlow无法表示这种控制逻辑

我们可以看到上面的实现是比较烂的,这是因为我们使用 sess.run
对A进行求值之后,没做任何修改又放回了原始的计算图,而TensorFlow 计算图与 Python 交换数据频繁时会严重拖慢运算速度。除了性能的问题,在Python层做Control Flow,你会发现在计算图中并没有表示 Python 逻辑,如果你将 graph 导出,实际上是看不到这些 if/else 语句的,因此网络结构信息会丢失。

 

这个问题趟过Pytorch导出ONNX的读者应该知道,我们如果想导出一个完整的检测模型,带了NMS后处理那幺必须找一张可以正常输出目标的图片作为输入。如果随机输出很可能后处理那部分在导出时就会被丢掉,这就是因为在Pytorch实现检测模型的时候在Python层用了If这种Control Flow。而Pytorch在导出ONNX模型时是根据输出结果进行反推,保留那些对结果有贡献的节点作为导出的计算图,删除那些没有贡献的节点。我们想一下,如果实现模型的过程中有Python层的Control Flow,那幺必然有一部分节点会被丢弃。

 

其实在Pytorch推出ScriptModule之后我就一直在想为什幺Pytorch不主推和ONNX的交互,直到我在官方文档看到下面这段话(当然这也可能并不是Pytorch的想法,我YY的):

大概就是如果想在Pytorch里面导出含有Python层控制流的模型时导出ONNX会丢失控制流,如果需要保留建议导出TorchScript模型

因此,我们作为用户自然希望深度学习框架支持控制流了,这样我们就不需要像上面那个例子那样很naive的在Python层去实现控制逻辑了。

 

TensorFlow的原生控制流

 

TensorFlow提供了几个运算符用于原生控制流,如下:

TensorFlow提供了几个运算符用于原生控制流

那幺使用这些原生控制流好处是什幺呢?

 

高效。TensorFlow 计算图与 Python 交换数据比较慢,计算图如果是端到端的,才能将数据传输开销降到最低,运行速度更快。

 

灵活。静态计算图可以使用动态模块加强,计算图逻辑是自包含的。Pytorch目前比TensorFlow更受欢迎的主要原因就是前者为动态计算图,可以在运行时修改计算图。TensorFlow 利用控制流可以在一个静态定义的计算图中实现类似动态计算图的功能。

 

兼容。通过 TensorBoard 调试和检查计算图,无缝通过 TensorFlow Serving 部署,也可以利用自动微分,队列和流水线机制。

 

控制依赖

 

TensorFlow会记录每一个运算符的依赖,然后基于依赖进行调度计算。也就是说一个运算符当且仅当它的依赖都完成之后才会执行一次。任何两个完成依赖的运算符可以以任意顺序进行。但这种设定可能会引发 竞争
,比如:

控制依赖引发竞争

其中 var 为一个变量,在对 bot 求值时,var 本身自增 2,然后将自增后的值返回。这时 top 语句执行顺序就会对 out 结果产生不同影响,结果不可预知。

 

为了解决这个问题,开发者可以人为的加入bot和top的依赖关系,让指定运算符先完成,如下图所示:

人为的加入bot和top的依赖关系,让指定运算符先完成

这里表示的就是如果我们需要保证读取的值是最新的,就需要新增下图中虚线箭头表示的依赖关系,即下图中上方蓝色圆圈依赖下方蓝色圆圈的运算完成,才能进行计算。

加入依赖关系之后,计算图长这样

条件分支

 

接下来看一下条件分支,即TensorFlow如何处理我们在这一节开头提出来的那个例子?

TensorFlow提供了两个条件控制OP,即tf.cond和tf.case

下面的代码中,利用了tf.cond实现条件分支,在 a < b 为真时,对 out 求值会执行 tf.add(3, 3);否则,执行 tf.square(3)。

使用tf.cond实现条件分支

上面这段代码等价于: tf.cond(a < b, lambda: tf.add(3, 3), lambda: tf.sqaure(3))

 

然后生成的计算图如下所示:

带有条件控制流的计算图

当并列的分支比较多时,我们可以使用tf.case来处理,例如:

并列的条件分支>2个时,使用tf.case来控制

循环

 

TensorFlow提供了 tf.while_loop
来构造循环块,感觉和RNN类似的结构有这个需求,例如:

tf.while_loop可以实现循环控制流解决RNN这种计算图结构的控制逻辑

下面的代码实现了一个基础的循环例子,即循环100次。

使用tf.while_loop在静态图中实现循环控制流

总的来说,TensorFlow应该是首个将Control Flow引入到计算图中的深度学习框架,这方面必须给予一定的尊重。即使Pytorch目前在学术界已经比TensorFlow更加流行,但基于TensorFlow演化的各种工业级项目仍然发挥着作用。

 

0x3. 总结

 

这篇文章介绍了一下深度学习中的Data Flow和Control Flow,其实也想学习一下Pytorch的JIT技术并放一起介绍的。来看看Pytorch是如何通过JIT技术将Python层的控制流导出的,但目前还没学习暂时就算了。这篇文章作为从零开始学深度学习编译器的番外篇,可能会让我们在深入到TVM的Relay时能理解为什幺要发明Relay这个IR,其中一个主要原因就是为了解决这里提到的控制流问题。

 

0x4. 参考

https://blog.csdn.net/lvxingzhe123456/article/details/82597095

Logical Graphs: Native Control Flow Operations in TensorFlow

https://mp.weixin.qq.com/s/6uVeEHcQeaPN_qEhHvcEoA

 

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注