TensorFlow 日常更新：加入 AutoGraph,简化构图流程

2018 年 7 月 19 日 专知

【导读】谷歌发布了一个新的 TensorFlow 工具「AutoGraph」，旨在快速将Python 代码转化到TensorFlow体系里。目前测试来看，这一工具简化了构图流程，加强了 TensorFlow 调用Python 时的性能。

作者 | Alex Wiltschko, Dan Moldovan, Wolff Dobson

整理 | Huaiwen

在正常版的 TensorFlow中，代码构建流程是：先构建算法的计算图，然后feed数据, 然后执行计算图。当然因为当前 PyTorch等动态计算图的流行，TensorFlow 也发布了 Eager Execution，它可以帮助用户自动构建计算图。但一般的 TensorFlow 还是常使用静态计算图的方式，因为它的构建逻辑与部署都非常有优势。

然而，静态计算图比较难理解，它与日常那种从上至下一句一句执行的逻辑相违背。尤其是在一些涉及复杂模型场景中，例如使用 if 和 while 等 Python 语句，或使用 print() 与接受结构化输入等，Tensorflow 总会让人感到深深的无力。

总于，最新加入的 Feature：AutoGraph 可以自动搞定上述事情。

使用 autograph.convert() 装饰器来装饰函数，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 op 一样：

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 看上去就是一种翻译机制，把 Python 代码，翻译成 TensorFlow 的规范。实际上，并不是，AutoGraph 要高端一些。它不只是宏命令的集合，它还可以使用源代码转换来覆写 Python 语言的任意部分，包括控制流、函数应用和分配，生成样板代码，重构惯用 Python，以使转换成图的过程变得简单。

那么，AutoGraph 可以为我们做什么呢？以下有一些示例代码，它们可以直接转换为图代码而不需要任何的改写。如果你想实际运行这些操作，谷歌在这个 GitHub 的 Colab 中提供了一个 notebook 可供使用。

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

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

以下我们使用循环和分支来测试「科拉兹猜想」。注意，考虑到多样性，我们将不使用 decorator，而使用 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)

AutoGraph还支持 break、continue，甚至 print 和 assert 等语句。当转换完成后，这个片段的 Python assert 使用合适的 tf.Assert 将其转换为 TensorFlow 计算图。

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

具备轻易地添加循环、控制流等到图上的能力意味着可以很容易将训练循环转移到图中。可以在这个 Colab 的 notebook 中找到一个示例，其中使用了一个 RNN 训练循环，并用一个 sess.run() 调用来执行它。当你需要传递一个完整的训练循环到加速器时，这很有用，比通过 CPU 控制器管理训练过程更好。

AutoGraph 打开了构建和训练模型的新思路。谷歌在未来将基于开发者社区建议尝试添加更多的功能到 AutoGraph 上。

原文链接：

https://medium.com/tensorflow/autograph-converts-python-into-tensorflow-graphs-b2a871f87ec7

-END-

专 · 知

人工智能领域主题知识资料查看与加入专知人工智能服务群：

【专知AI服务计划】专知AI知识技术服务会员群加入与人工智能领域26个主题知识资料全集获取。欢迎微信扫一扫加入专知人工智能知识星球群，获取专业知识教程视频资料和与专家交流咨询！