【GiantPandaCV 导语】本文作为从零开始学深度学习编译器的番外篇,介绍了一下深度学习框架的 Data Flow 和 Control Flow,并基于 TensorFlow 解释了 TensorFlow 是如何在静态图中实现 Control Flow 的。而对于动态图来说,是支持在 Python 层直接写 Control Flow 的,最后基于 Pytorch 介绍了如何将 Python 层的 Control Flow 导出到 TorchScript 模型以及 ONNX 模型。
0x0. 前言 ¶
本来是想在讲 TVM Relay 的时候提一下 DataFlow 和 ControlFlow 的,但是担心读者看到解析代码的文章打开就关了,所以这里用一篇简短的文章来介绍一下深度学习框架中的 DataFlow 和 ControlFlow。
0x1. DataFlow¶
我记得我接触的第一个深度学习框架是 TensorFlow1.x,本科毕业设计也是基于 TensorFlow 完成的,因此这里我将以 TensorFlow1.x 为例介绍一下 DataFlow。
假设现在我们要实现一个(a+b)∗c 的逻辑,其中a,b,c 都是一个简单的实数,然后我们如果用 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 实现的这个计算函数首先在内存中构造了一个数据流图,长这样:
我们回看一下 Python 的实现,实际上在执行res = (a + b) * c
这行代码时,已经计算出了res
的值,因为 Python 这种过程语言的数学计算是由代码驱动的。而 TensorFlow 不一样,它首先构造了数据流图,然后对这个计算流图进行绑定数据,让这个数据在这个图里面流起来,这是显示调用sess.run
来获得输出的。
像 TensorFlow 这种基于数据流图(DataFlow)进行计算的深度学习框架不少,比如早期的 Theano,2020 年开源的国内深度学习框架 OneFlow,PaddlePaddle1.x 初级版本都是基于数据流图的。当然更多人将其称为静态图。
0x2. Control Flow¶
这一节我将结合 TensorFlow1.x 的 Control Flow 来为大家解析一下 Control Flow 的难点以及 TensorFlow 的一些解决方案。这里的内容理解主要基于这篇博客 (https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/),感兴趣的同学可以去查看原文。
在计算机科学中,控制流(Control Flow)定义了独立语句,指令,函数调用等执行或者求值的顺序。举个例子,我们要实现一个本机控制流,即我们需要根据函数 A 的输出值选择运行函数 B 或者 C 中的一个:
然后要实现这个控制流,最 Naive 的方式在是 Python 端写 If/Else 语句,即 Python 端的 Control Flow,然后在不同条件下使用 session.run() 来求取不同分支的值。对于 TensorFlow 来说,大概是这样:
然后这个 Python 层的 Control Flow 并不会在计算图中被表示出来,即:
我们可以看到上面的实现是比较烂的,这是因为我们使用sess.run
对 A 进行求值之后,没做任何修改又放回了原始的计算图,而 TensorFlow 计算图与 Python 交换数据频繁时会严重拖慢运算速度。除了性能的问题,在 Python 层做 Control Flow,你会发现在计算图中并没有表示 Python 逻辑,如果你将 graph 导出,实际上是看不到这些 if/else 语句的,因此网络结构信息会丢失。
这个问题趟过 Pytorch 导出 ONNX 的读者应该知道,我们如果想导出一个完整的检测模型,带了 NMS 后处理那么必须找一张可以正常输出目标的图片作为输入。如果随机输出很可能后处理那部分在导出时就会被丢掉,这就是因为在 Pytorch 实现检测模型的时候在 Python 层用了 If 这种 Control Flow。而 Pytorch 在导出 ONNX 模型时是根据输入跑一遍模型即 tracing(这是以前的版本的做法,新版本的 TensorFlow 已经支持导出 Python 层的 Control Flow),然后记录这个过程中发生了哪些操作。我们想一下,如果实现模型的过程中有 Python 层的 Control Flow(基于 tracing 机制),那么必然有一部分节点会被丢弃。
Pytorch 官方文档指出当我们导出 ONNX 的时候如果想导出 Python 层的控制流到计算图中就需要包一层@jit.script
像 Pytorch 这种动态图框架可以方便的使用 Python 层的 Control Flow,但 TensorFlow 在 1.x 时代为了解决这个问题是花费了不少努力的,即 TensorFlow1.x 的原生控制流。
TensorFlow 的原生控制流 ¶
TensorFlow 提供了几个运算符用于原生控制流,如下:
那么使用这些原生控制流好处是什么呢?
- 高效。TensorFlow 计算图与 Python 交换数据比较慢,计算图如果是端到端的,才能将数据传输开销降到最低,运行速度更快。
- 灵活。静态计算图可以使用动态模块加强,计算图逻辑是自包含的。Pytorch 目前比 TensorFlow 更受欢迎的主要原因就是前者为动态计算图,可以在运行时修改计算图。TensorFlow 利用控制流可以在一个静态定义的计算图中实现类似动态计算图的功能。
- 兼容。通过 TensorBoard 调试和检查计算图,无缝通过 TensorFlow Serving 部署,也可以利用自动微分,队列和流水线机制。
控制依赖 ¶
TensorFlow 会记录每一个运算符的依赖,然后基于依赖进行调度计算。也就是说一个运算符当且仅当它的依赖都完成之后才会执行一次。任何两个完成依赖的运算符可以以任意顺序进行。但这种设定可能会引发竞争,比如:
其中 var 为一个变量,在对 bot 求值时,var 本身自增 2,然后将自增后的值返回。这时 top 语句执行顺序就会对 out 结果产生不同影响,结果不可预知。
为了解决这个问题,开发者可以人为的加入 bot 和 top 的依赖关系,让指定运算符先完成,如下图所示:
这里表示的就是如果我们需要保证读取的值是最新的,就需要新增下图中虚线箭头表示的依赖关系,即下图中上方蓝色圆圈依赖下方蓝色圆圈的运算完成,才能进行计算。
条件分支 ¶
接下来看一下条件分支,即 TensorFlow 如何处理我们在这一节开头提出来的那个例子?
下面的代码中,利用了 tf.cond 实现条件分支,在 a <b 为真时,对 out 求值会执行 tf.add(3, 3);否则,执行 tf.square(3)。
上面这段代码等价于:tf.cond(a < b, lambda: tf.add(3, 3), lambda: tf.sqaure(3))
然后生成的计算图如下所示:
当并列的分支比较多时,我们可以使用 tf.case 来处理,例如:
循环 ¶
TensorFlow 提供了tf.while_loop
来构造循环块,感觉和 RNN 类似的结构有这个需求,例如:
下面的代码实现了一个基础的循环例子,即循环 100 次。
总的来说,TensorFlow 应该是首个将 Control Flow 引入到计算图中的深度学习框架,而不是像动态图框架那样直接在 Python 层去做 Control Flow,这方面必须给予一定的尊重。即使 Pytorch 目前在学术界已经比 TensorFlow 更加流行,但基于 TensorFlow 演化的各种工业级项目仍然发挥着作用。
0x3. Pytorch 中的 Control Flow¶
在 Pytorch 这种动态图框架中,支持直接在 Python 端写 Control Flow,并且可以将这些控制逻辑放到计算图中。这里以 TorchScript 为例,当我们尝试将 Pytorch 模型转为 TorchScript 时,有两种方式,一种是 trace,另外一种是 script。对于 trace 模式,适合 Python 层没有 Control Flow 的计算图,举例如下:
#coding=utf-8
import torch
import torch.nn as nn
class MyModule(nn.Module):
def __init__(self):
super(MyModule,self).__init__()
self.conv1 = nn.Conv2d(1,3,3)
def forward(self,x):
x = self.conv1(x)
return x
model = MyModule() # 实例化模型
trace_module = torch.jit.trace(model,torch.rand(1,1,224,224))
print(trace_module.code) # 查看模型结构
output = trace_module (torch.ones(1, 1, 224, 224)) # 测试
print(output)
# trace_modult('model.pt')
打印trace_module
的代码可以看到:
def forward(self,
input: Tensor) -> Tensor:
return (self.conv1).forward(input, )
而 script 模式则适用于计算图在 Python 层有 Control Flow 的情况,比如:
#coding=utf-8
import torch
import torch.nn as nn
class MyModule(nn.Module):
def __init__(self):
super(MyModule,self).__init__()
self.conv1 = nn.Conv2d(1,3,3)
self.conv2 = nn.Conv2d(2,3,3)
def forward(self,x):
b,c,h,w = x.shape
if c ==1:
x = self.conv1(x)
else:
x = self.conv2(x)
return x
model = MyModule()
# 这样写会报错,因为有控制流
# trace_module = torch.jit.trace(model,torch.rand(1,1,224,224))
# 此时应该用script方法
script_module = torch.jit.script(model)
print(script_module.code)
output = script_module(torch.rand(1,1,224,224))
打印script_module
的代码可以看到 TorchScript 模型包含了在上面 Python 层定义的 Control Flow:
def forward(self,
x: Tensor) -> Tensor:
b, c, h, w, = torch.size(x)
if torch.eq(c, 1):
x0 = (self.conv1).forward(x, )
else:
x0 = (self.conv2).forward(x, )
return x0
然后我们来实验一下将上面带有 Control Flow 的 Module 导出 ONNX,这里以 Pytorch 官方文档提供的一个带循环的 Control Flow 的示例为例:
import torch
# Trace-based only
class LoopModel(torch.nn.Module):
def forward(self, x, y):
for i in range(y):
x = x + i
return x
model = LoopModel()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)
torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True)
这样就可以成功导出名字为loop
的 ONNX 模型,使用 Netron 可视化软件打开看一下:
而当我们使用 script 模式时,导出的 ONNX 就会保留 Python 层的 Control Flow 并将其转换成 ONNX 中的 Loop OP。示例代码以及 Netron 可视化结果如下:
import torch
# Mixing tracing and scripting
@torch.jit.script
def loop(x, y):
for i in range(int(y)):
x = x + i
return x
class LoopModel2(torch.nn.Module):
def forward(self, x, y):
return loop(x, y)
model = LoopModel2()
dummy_input = torch.ones(2, 3, dtype=torch.long)
loop_count = torch.tensor(5, dtype=torch.long)
torch.onnx.export(model, (dummy_input, loop_count), 'loop.onnx', verbose=True,
input_names=['input_data', 'loop_range'])
0x4. 总结 ¶
这篇文章介绍了一下深度学习中的 Data Flow 和 Control Flow,然后介绍了一下将 Pytorch 模型转为 TorchScript 的两种模式,并探索了要将 Pytorch 的 Python 层的 Control Flow 转换为 ONNX 应该怎么做。
0x5. 参考 ¶
- https://blog.csdn.net/lvxingzhe123456/article/details/82597095
- https://www.altoros.com/blog/logical-graphs-native-control-flow-operations-in-tensorflow/
- https://mp.weixin.qq.com/s/6uVeEHcQeaPN_qEhHvcEoA
欢迎关注 GiantPandaCV, 在这里你将看到独家的深度学习分享,坚持原创,每天分享我们学习到的新鲜知识。(• ̀ω•́)✧
有对文章相关的问题,或者想要加入交流群,欢迎添加 BBuf 微信:
为了方便读者获取资料以及我们公众号的作者发布一些 Github 工程的更新,我们成立了一个 QQ 群,二维码如下,感兴趣可以加入。
本文总阅读量476次