Py学习  »  Python

AutoGraph 将 Python 转换为 TensorFlow 图

TensorFlow • 5 年前 • 594 次点击  

文 / Alex Wiltschko、Dan Moldovan 和 Wolff Dobson



今天,我们将介绍一项名为 “AutoGraph” 的 TensorFlow 新功能。AutoGraph 可以将 Python 代码(包括控制流、print() 和其他 Python 原生功能)转换为纯 TensorFlow 图代码。


要在不使用即时执行的情况下编写 TensorFlow 代码,您需要进行一些元编程,即编写一个用于创建图的程序,然后再执行该图。这种方式可能会令人困惑,特别是对于新开发者而言。一些特别棘手的情况涉及更复杂的模型,例如使用 if 和 while 的模型,或者具有 print() 等副作用或接受结构化输入的模型。


那么,我们为什么需要图呢?图允许各种优化,如移除常见的子表达式和融合内核。此外,由于图形成了一种独立于平台的计算模型,可简化分布式训练和针对各种环境的部署。这对于多个 GPU 或 TPU 上的分布式训练或者通过 TensorFlow Lite 在移动或物联网等其他平台上分发模型尤为重要。


下面这个简单示例显示了您可能要添加到图的运算:


1    def huber_loss(a):    

2        if tf.abs(a) <= delta:    

3            loss = a * a / 2    

4        else:    

           loss = delta * (tf.abs(a) - delta / 2)    

6        return loss


如果使用即时执行,也可以满足要求,但是由于 Python 解释器开销或错过程序优化机会,这样可能会很慢。


要准备好执行图,您需要进行重写以使用类似 tf.cond() 的语句,但这可能比较繁琐并且难以实现。AutoGraph 可以为您自动执行此转换,既保持了即时编程的简易性,同时又获得了图执行的性能优势。


在本例中,我们可以使用 autograph.convert() 装饰函数,AutoGraph 将自动生成可生成图的代码。


使用 AutoGraph 时,由于装饰器的原因,以下代码:


1    @autograph.convert()    

2    def huber_loss(a):    

       if tf.abs(a) <= delta:    

4            loss = a * a / 2    

5        else:    

6            loss = delta * (tf.abs(a) - delta / 2)    

       return loss 


在执行时变为以下代码。


1    def tf__huber_loss(a):    

2        with tf.name_scope('huber_loss'):    

3    

           def if_true():    

5                with tf.name_scope('if_true'):    

6                    loss = a * a / 2    

7                    return loss,    

8

9                def if_false():    

10                with tf.name_scope('if_false'):    

11                    loss = delta * (tf.abs(a) - delta / 2)    

12                    return loss,    

13                loss = ag__.utils.run_cond(tf.less_equal(tf.abs(a),  delta), if_true,  

14                    if_false)    

15                return loss 


然后,您可以将这些代码视为 TensorFlow 运算进行调用:


1    with tf.Graph().as_default():    

2        x_tensor = tf.constant(9.0)    

3

4        # The converted function works like a regular op: tensors in, tensors out.    

5        huber_loss_tensor = huber_loss(x_tensor)    

6

7        with tf.Session() as sess:    

8            print('TensorFlow result: %2.2f\n' % sess.run(huber_loss_tensor))    


如您所见,AutoGraph 填补了即时执行与图之间的差距。AutoGraph 接收即时风格的 Python 代码并将其转换为生成图的代码。


AutoGraph 不仅仅是一组有用的宏,它还使用源代码转换来允许替换 Python 语言的 任意 部分,包括控制流、函数应用和赋值、生成模板代码,以及重构惯用 Python 以便轻松转换成图。


使用任何编译器,都会对错误消息的可读性产生担忧;为此,AutoGraph 设置为创建错误消息和堆叠追踪,以显示原始源代码中的错误源,而不是仅显示对生成代码的引用。



可运行示例

那么,AutoGraph 可以为我们做什么呢?以下是一些代码示例,现在可以直接转换为图代码而无需任何更改。如果您想实际运行所有这些操作,我们提供了一个笔记本,您可以在 Colab 中打开或在 GitHub 中查看。

注:GitHub 链接https://github.com/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb


接下来,我们使用循环和分支来检查 Collatz 猜想。请注意,出于多样性考虑,我们将不使用装饰器,而使用 AutoGraph 的 .to_graph() 函数将其转换为图。


1    def collatz(a):    

2            counter = 0    

3            while a != 1:    

                   if a % 2 == 0:    

                       a = a // 2    

6                    else:    

7                        a = 3 * a + 1    

8                    counter = counter + 1    

9            return counter

10 

11        graph_mode_collatz = autograph.to_graph(collatz)    

12        # The code is human-readable, too    

13        print(autograph.to_code(collatz))    

14

15        collatz_tensor = graph_mode_collatz(tf.constant(n))    


AutoGraph 可以支持任意的嵌套控制流,例如:


1    def f(n):    

2        if n >= 0:    

3            while n < 5:    

4                n += 1    

               print(n)    

       return n    


AutoGraph 允许在循环内向数组追加元素。为此,我们将使用一些 AutoGraph 辅助工具,set_element_type 和 stack。


1    def f(n):    

       z = []    

3        # We ask you to tell us the element dtype of the list    

4        autograph.set_element_type(z, tf.int32)    

5        for i in range(n):    

           z.append(i)    

7        # when you're done with the list, stack it    

8        # (this is just like np.stack)    

9        return autograph.stack(z) 


我们还支持 break、continue,甚至 print 和 assert 等语句。转换后,此代码段的 Python assert 转换为使用适当的 tf.Assert 的图。


   def f(x):    

2        assert x != 0, 'Do not pass zero!'    

3        return x * x    


能够轻松向图添加循环和控制流等,这意味着可以轻松将训练循环转移到图中。您可以在此笔记本中找到相关示例,其中,我们接受一个 RNN 训练循环并用一个 sess.run() 调用执行。在需要将整个训练循环传递给加速器而不是通过 CPU 控制器管理训练的情况下,这样做十分有用。

注:笔记本链接

https://colab.research.google.com/github/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb#scrollTo=4LfnJjm0Bm0B


AutoGraph 开启了构建和训练模型的新思路。我们期待根据开发者社区的建议为 AutoGraph 添加更多功能,因此欢迎您提交问题并给出建议!



图性能与即时执行

即时执行非常简单易用,但图的速度通常要更快。虽然二者的比较基准较为复杂(并且取决于应用和硬件配置),但在这个简单的示例中,我们可以看到,从即时切换到大量使用 if 和 while 的 AutoGraph 代码时,速度有了显著提升。


最终,AutoGraph 允许您在加速器硬件(如 GPU 和 Cloud TPU)上使用动态和流控制较多的模型,这在基于大量数据训练大型模型时十分必要。


我们刚刚开始着手分析性能。如果您发现某个图语句的运行速度低于预期,请提交问题!



AutoGraph 和即时执行

虽然使用即时执行时,您仍然可以通过 tf.contrib.eager.defun 对部分代码使用图执行,但这要求您使用 tf.cond() 等 TensorFlow 图运算。将来,AutoGraph 将与 defun 无缝集成,允许用简单的即时式 Python 语言编写图代码。实现此功能后,您可以通过选择性地将即时代码转换为图片段来使用 AutoGraph 加速热点。



结论

AutoGraph 是一款新工具,可帮您轻松构建能够在 TensorFlow 图中轻松运行的直观、复杂的模型。此工具目前为 contrib 中的实验性工具,但我们希望尽快将它加入到核心 TensorFlow 中。


跟我们分享您的 AutoGraph 使用体验!如果您有任何反馈、建议或想法,请提交问题并向 TensorFlow 开发者群组发送消息。

注:TensorFlow 开发者群组链接https://groups.google.com/a/tensorflow.org/forum/#!forum/developers


更多 AI 相关阅读:

· 使用 TensorFlow 和 DLTK 进行生物医学影像分析

· Seedbank - 发现机器学习示例

· TensorFlow 版本 1.10.0 发布啦!



今天看啥 - 高品质阅读平台
本文地址:http://www.jintiankansha.me/t/OQB1PrQMqt
Python社区是高质量的Python/Django开发社区
本文地址:http://www.python88.com/topic/22480
 
594 次点击