• 微信公众号:美女很有趣。 工作之余,放松一下,关注即送10G+美女照片!

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

开发技术 开发技术 2小时前 2次浏览

深度学习编译器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.02.03.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实现的这个计算函数,先在内存中构造了一个数据流图:

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

 

 上面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中的一个:

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

 

 一个Control Flow的例子

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

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

 

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

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

 深度学习编译器Data Flow和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提供了几个运算符用于原生控制流,如下:

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

 

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

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

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

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

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

 

控制依赖

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

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

 

 控制依赖引发竞争

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

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

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

 

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

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

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

 

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

条件分支

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

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

 

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

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

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

 

 使用tf.cond实现条件分支

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

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

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

 

 带有条件控制流的计算图

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

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

 

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

循环

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

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

 

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

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

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

 

 使用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(11224224)) 测试
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(23, 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可视化软件打开看一下:

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

 

 

可以看到直接导出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(23, 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

 

 


程序员灯塔
转载请注明原文链接:深度学习编译器Data Flow和Control Flow
喜欢 (0)