技术开发 频道

TensorFlow新工具 | AutoGraph

  【IT168 技术】前言

  我们(TensorFlow官方)想告诉你一个名为“AutoGraph”的新TensorFlow功能。 AutoGraph将Python代码(包括控制流,print()和其他Python原生特征)转换为纯TensorFlow图形代码。

  为什么我们需要图(graph)? 图允许各种优化,例如删除常见的子表达式和融合内核。此外,图形简化了分布式培训和部署到各种环境,因为它们形成了独立于平台的计算模型。这对于多个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

  使用 Eager Execution,这只是「正确运行」而已,但是此类操作可能会比较慢,因为 Python 解释器众所周知在实现地比较慢,且需要的计算比较复杂,这会令它错过许多程序优化的机会。

  为了给图执行做好准备,你需要重写代码,使用 tf.cond() 等语句,但是这很繁琐且难以实现。AutoGraph 可以自动完成该转换,保持 Eager 编程的简易性,同时还提升了计算图执行的性能。

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

  使用 AutoGraph,由于 decorator,下列代码:

  @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 连接起 Eager execution 和 Graph。AutoGraph 使用 Eager-style 的 Python 代码,然后将其转换成图生成代码。

  AutoGraph 不只是有用宏命令的集合,它还可以使用源代码转换来覆写 Python 语言的任意部分,包括控制流、函数应用和分配,生成样板代码,重构惯用 Python,以使转换成图的过程变得简单。

  使用任意编译器,都会对错误信息可读性产生担忧;为此,AutoGraph 可以创建错误信息,并堆叠揭示原始源代码中错误来源的多个轨迹,而不是仅仅显示生成代码的 reference。

  可运行示例

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

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

  我们还支持 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 控制器管理训练过程更好。

  Graph Performance 对比 Eager Execution

  Eager Execution 相当合用,但图更快。尽管对比基准较为复杂(由应用以及硬件配置决定),但在一些简单示例中我们可以看到,当从 Eager 转换到 AutoGraph 代码时有极大的加速,使用了大量 if 和 while 等语句。

  最终,AutoGraph 让你可以在 GPU 和 Cloud TPU 这样的加速器硬件上使用动态和流控制极严模型,这对在大量数据上训练大型模型非常有帮助。

  AutoGraph 和 Eager Execution

  虽然使用 Eager Execution,你也能通过 tf.contrib.eager.defun 对部分代码根据计算图执行。但这需要你使用 tf.cond() 这样计算图类的 TensorFlow ops。未来,AutoGraph 将无缝与 defun 融合,让你用简单的 eager-style Python 编写图代码。当成为现实时,通过选择性的把 eager 代码转换到图分段,你就可以期待使用 AutoGraph 加速热点了。

  结论

  AutoGraph是一款工具,可让您轻松构建直观,复杂的模型,在TensorFlow图中轻松运行。 这是一个现在在contrib中的实验工具,但我们希望尽快将其转移到核心TensorFlow中。

0
相关文章