TF新工具AutoGraph：将Python转换为TensorFlow图

编译：yxy

出品：ATYUN订阅号

我们在这里向你介绍一个名为“AutoGraph”的TensorFlow新功能。AutoGraph将Python代码（包括控制流print()和其他Python原生特性）转换为纯的TensorFlow图代码。

在不使用急切执行的情况下编写TensorFlow代码需要你进行一些元编程，即编写一个创建图的程序，然后稍后执行这个图。这个过程可能使人感到混乱，特别是对于新开发人员并不友好。一些特别棘手的情况涉及更复杂的模型，例如使用if和while的，或具有副作用（如，print()），或接受结构化输入的。

为什么我们需要图？图可以进行各种优化，例如删除常见的子表达式和融合内核。此外，简化了对各种环境的分布式训练和部署，因为它们形成了独立于平台的计算模型。这对于多个GPU或TPU上的分布式训练，或者通过TensorFlow Lite在移动或物联网等其他平台上分发模型而言尤为重要。

以下是你想要添加到图中的操作的一个非常简单的示例：

def huber_loss(a): if tf.abs(a)

通过急切执行，它可以“正常工作”，但是由于Python解释器的负担，可能错过的程序优化机会，此类操作可能会很慢。

要使图迅速执行，需要重写它以使用像tf.cond()这样的结构，但这个实现繁琐而困难。AutoGraph可以为自动执行这个转换，保持了急切编程的简易性，同时获得了基于图的执行的性能优势。

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

使用AutoGraph，这段代码：

@autograph.convert()def huber_loss(a): if tf.abs(a)

由于这个修饰器（decorator），在执行时变为此代码。

def tf__huber_loss(a): with tf.name_scope('huber_loss'): def if_true(): with tf.name_scope('if_true'): loss = a * a / 2 return loss, def if_false(): with tf.name_scope('if_false'): loss = delta * (tf.abs(a) - delta / 2) return loss, loss = ag__.utils.run_cond(tf.less_equal(tf.abs(a), delta), if_true, if_false) return loss

然后你可以像调用TensorFlow操作一样调用你的代码：

with tf.Graph().as_default(): x_tensor = tf.constant(9.0) # The converted function works like a regular op: tensors in, tensors out. huber_loss_tensor = huber_loss(x_tensor) with tf.Session() as sess: print('TensorFlow result: %2.2f\n' % sess.run(huber_loss_tensor))

如你所见，AutoGraph填补了急切执行和图之间的空白。AutoGraph接收急切执行风格的Python代码并将其转换为生成图的代码。

AutoGraph不仅仅是有用的宏命令的集合；它使用源代码转换来覆盖Python语言的任何部分，包括控制流，函数应用和赋值，生成样板代码，以及重构通顺的Python以使其易于转换为图。

对于任何编译器来说，我们都要考虑错误消息的可读性；为此，AutoGraph设置为创建错误消息和堆栈跟踪，以显示原始源代码中的错误源，而不是仅显示对生成的代码的引用。

运行示例

那么，AutoGraph可以为您做什么？以下是一些代码示例，现在可以直接转换为图代码，无需任何更改。如果你想查看所有这些，可以访问下方链接。

Colab：https://colab.research.google.com/github/tensorflow/models/blob/master/samples/core/guide/autograph.ipynb

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

在这里，我们使用循环和分支证实Collatz猜想。注意，它是多样性的，所以我们不用修饰器，我们使用AutoGraph的 .to_graph()函数将其转换为图。

def collatz(a): counter = 0 while a != 1: if a % 2 == 0: a = a // 2 else: a = 3 * a + 1 counter = counter + 1 return countergraph_mode_collatz = autograph.to_graph(collatz)# The code is human-readable, tooprint(autograph.to_code(collatz))collatz_tensor = graph_mode_collatz(tf.constant(n))

AutoGraph可以支持任意的嵌套控制流，例如：

def f(n): if n >= 0: while n

AutoGraph允许你将元素追加到循环内的数组中。想要这样做，我们使用一些AutoGraph的辅助，set_element_type和stack。

def f(n): z = [] # We ask you to tell us the element dtype of the list autograph.set_element_type(z, tf.int32) for i in range(n): z.append(i) # when you're done with the list, stack it # (this is just like np.stack) return autograph.stack(z)

我们还支持诸如break、continue、甚至print和assert之类的结构。转换完成后，此片段的Python assert将转换为使用适当的tf.Assert的图。

def f(x): assert x != 0, 'Do not pass zero!' return x * x

能够轻松地添加循环，控制流程以及更多图意味着可以轻松地将训练循环移动到图中。在下面链接中可以看到这种例子，我们获取RNN训练循环并通过sess.run()调用执行它。在需要将整个训练循环传递给加速器而不是通过CPU控制器管理训练的情况下，这可能非常有用。

AutoGraph开辟了构建和训练模型的新思路。我们期待根据开发者社区的建议为AutoGraph添加更多功能，所以请提出建议问题！

图性能与急切执行的对比

急切的执行非常有用，但图通常要快得多。虽然基准测试很复杂（并且取决于应用程序和硬件配置），但在这个简单的示例中，我们看到，从急切执行转换到大量使用if和whileAutoGraph代码时有显著的加速。

最终，AutoGraph允许您在GPU和云TPU等加速器硬件上使用动态和流控制较多的模型，这对于在大量数据上训练大型模型是必要的。

我们刚刚开始探索性能的过程。如果发现运行速度低于预期的图构造，请提出问题！

AutoGraph和急切执行的对比

结论

AutoGraph是一款可让你轻松构建直观，复杂的模型，在TensorFlow图中轻松运行的工具。它现在是个实验性工具，但我们希望尽快将其加入到TensorFlow的核心中。