发布者 Alex Wiltschko，Dan Moldovan，Wolff Dobson


我们想告诉您一个名为“AutoGraph”的新 TensorFlow 功能。AutoGraph 将 Python 代码（包括控制流、print() 和其他 Python 原生功能）转换为纯 TensorFlow 图代码。

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

我们为什么需要图呢？图允许各种优化，例如删除公共子表达式和融合内核。此外，图简化了分布式训练和部署到各种环境，因为它们形成了与平台无关的计算模型。这对于在多个 GPU 或 TPU 上进行分布式训练，或通过 TensorFlow Lite 将您的模型分布到其他平台（如移动或物联网）尤其重要。

以下是一个您可能想添加到图中的操作的简单示例

def huber_loss(a):
  if tf.abs(a) <= delta:
    loss = a * a / 2
  else:
    loss = delta * (tf.abs(a) - delta / 2)
  return loss

使用急切执行，这将“正常工作”，但是此类操作可能由于 Python 解释器开销或错过的程序优化机会而很慢。

为了使它准备好进行图执行，您需要将其重写以使用诸如 tf.cond() 之类的结构，但这可能很繁琐且难以实现。AutoGraph 可以自动为您执行此转换，保持急切编程的简便性，同时获得基于图的执行的性能优势。

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

使用 AutoGraph，这段代码

@autograph.convert()
def huber_loss(a):
  if tf.abs(a) <= delta:
    loss = a * a / 2
  else:
    loss = delta * (tf.abs(a) - delta / 2)
  return loss

由于装饰器，这段代码将在执行时变为以下代码。

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 中打开它，或在 GitHub 中查看它。
这里，我们使用循环和分支检查 哥德巴赫猜想。请注意，为了多样性，我们没有使用装饰器，而是使用 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 counter

graph_mode_collatz = autograph.to_graph(collatz)
# The code is human-readable, too
print(autograph.to_code(collatz))

collatz_tensor = graph_mode_collatz(tf.constant(n))

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

def f(n):
  if n >= 0:
    while n < 5:
      n += 1
      print(n)
  return 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

能够轻松地将循环、控制流和更多内容添加到图中意味着可以轻松地将训练循环移入图中。在这个 笔记本 中可以找到一个示例，我们在其中使用单个 sess.run() 调用执行 RNN 训练循环。这在您需要将整个训练循环传递给加速器而不是通过 CPU 控制器管理训练的情况下很有用。

AutoGraph 开启了构建和训练模型的新思维方式。我们期待根据开发者社区的建议为 AutoGraph 添加更多功能，因此请提交 问题 并提出建议！

图性能与急切执行

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

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

我们才刚刚开始探索性能。如果您发现图结构

AutoGraph 与急切执行

在使用急切执行时，您仍然可以通过 tf.contrib.eager.defun 为代码的某些部分使用图执行。这要求您使用图 TensorFlow 操作（如 tf.cond()）。将来，AutoGraph 将与 defun 无缝集成，允许以纯急切式 Python 编写图代码。当该实现可用时，您可以期望使用 AutoGraph 通过选择性地将急切代码转换为图片段来加快热点。

结论

AutoGraph 是一种工具，可以让您轻松地构建直观的复杂模型，这些模型可以在 TensorFlow 图中轻松运行。这现在是 contrib 中的一种实验性工具，但我们预计很快就会将其移入核心 TensorFlow。

告诉我们您使用 AutoGraph 的体验！如果您有反馈、建议或想法，请提交 问题 并向 TensorFlow 开发人员小组 发送消息。

致谢

我们要感谢 Andrew Johnson、Bart van Merriënboer、Zachary Nado 和 Alex Passos 的核心贡献。我们还要感谢以下同事：Akshay Agrawal、Mark Daoust、Josh Levenberg、Dougal Maclaurin、Rajat Monga、Mahima Pushkarna、Alexey Radul、D. Sculley 和 Asim Shankar。