深度学习编译器Data Flow和Control Flow

2021-11-25 07:01:15 阅读：222 来源： 互联网

标签：Control Python torch Flow 编译器 tf TensorFlow

深度学习编译器Data Flow和Control Flow

本文介绍了一下深度学习框架的Data Flow和Control Flow，基于TensorFlow解释了TensorFlow是如何在静态图中实现Control Flow的。支持在Python层直接写Control Flow的动态图，最后基于Pytorch介绍了如何将Python层的Control Flow导出到TorchScript模型以及ONNX模型。

1. 前言

1.1. DataFlow

以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 初级版本都是基于数据流图的。当然更多人称为静态图。

1.2. 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中的一个：

一个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模型时，根据输入跑一遍模型即tracing（这是以前的版本的做法，新版本的TensorFlow已经支持导出Python层的Control Flow），记录这个过程中发生了哪些操作。如果实现模型的过程中，有Python层的Control Flow（基于tracing机制），必然有一部分节点会丢弃。

Pytorch官方文档指出，当导出ONNX时，如果想导出Python层的控制流到计算图中，就需要包一层@jit.script。

大概就是如果想在Pytorch里面导出含有Python层控制流的模型时导出ONNX会丢失控制流，如果需要保留建议导出TorchScript模型或者使用基于script模型的导出方式

像Pytorch这种动态图框架，可以方便的使用Python层的Control Flow，但TensorFlow在1.x时代，为了解决这个问题，花费了不少努力，即TensorFlow1.x的原生控制流。

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引入到计算图中的深度学习框架，不是像动态图框架那样直接在Python层去做Control Flow，这方面必须给予一定的尊重。即使Pytorch目前在学术界已经比TensorFlow更加流行，但基于TensorFlow演化的各种工业级项目仍然发挥着作用。

3. 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可视化软件打开看一下：

可以看到直接导出Module，Python层的控制逻辑被丢掉（即for循环被完全展开），这是因为Pytorch在导出ONNX的时候默认使用了tracing机制

而当使用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'])

Pytorch模型中在Python层定义的Control Flow被保留下来了

4. 总结

这篇文章介绍了一下深度学习中的Data Flow和Control Flow，然后介绍了一下将Pytorch模型转为TorchScript的两种模式，并探索了要将Pytorch的Python层的Control Flow转换为ONNX应该怎么做。

5. 参考文献

https://mp.weixin.qq.com/s/Kt4xDLo-NRui8Whl0DqcSA

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

标签：Control,Python,torch,Flow,编译器,tf,TensorFlow
来源： https://www.cnblogs.com/wujianming-110117/p/15601118.html

本站声明： 1. iCode9 技术分享网（下文简称本站）提供的所有内容，仅供技术学习、探讨和分享；
2. 关于本站的所有留言、评论、转载及引用，纯属内容发起人的个人观点，与本站观点和立场无关；
3. 关于本站的所有言论和文字，纯属内容发起人的个人观点，与本站观点和立场无关；
4. 本站文章均是网友提供，不完全保证技术分享内容的完整性、准确性、时效性、风险性和版权归属；如您发现该文章侵犯了您的权益，可联系我们第一时间进行删除；
5. 本站为非盈利性的个人网站，所有内容不会用来进行牟利，也不会利用任何形式的广告来间接获益，纯粹是为了广大技术爱好者提供技术内容和技术思想的分享性交流网站。

ICode9

深度学习编译器Data Flow和Control Flow

深度学习编译器Data Flow和Control Flow

1. 前言

1.1. DataFlow

1.2. Control Flow

TensorFlow的原生控制流

控制依赖

条件分支

循环

3. Pytorch中的Control Flow

4. 总结

5. 参考文献