AutoGraph 将 Python 转换为 TensorFlow 图

文 / 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:    
5            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):    
3        if tf.abs(a) <= delta:    
4            loss = a * a / 2    
5        else:    
6            loss = delta * (tf.abs(a) - delta / 2