AutoGraph:将Python转换为TensorFlow图
2018年07月19日 由 yuxiangyu 发表
619797
0
我们在这里向你介绍一个名为“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
由于这个修饰器(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 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
能够轻松地添加循环,控制流程以及更多图意味着可以轻松地将训练循环移动到图中。在下面链接中可以看到这种例子,我们获取RNN训练循环并通过sess.run()调用执行它。在需要将整个训练循环传递给加速器而不是通过CPU控制器管理训练的情况下,这可能非常有用。
AutoGraph开辟了构建和训练模型的新思路。我们期待根据开发者社区的建议为AutoGraph添加更多功能,所以请提出建议问题!
图性能与急切执行的对比
急切的执行非常有用,但图通常要快得多。虽然基准测试很复杂(并且取决于应用程序和硬件配置),但在这个简单的示例中,我们看到,从急切执行转换到大量使用if和whileAutoGraph代码时有显著的加速。
最终,AutoGraph允许您在GPU和云TPU等加速器硬件上使用动态和流控制较多的模型,这对于在大量数据上训练大型模型是必要的。
我们刚刚开始探索性能的过程。如果发现运行速度低于预期的图构造,请提出问题!
AutoGraph和急切执行的对比
在使用急切执行时,通过tf.contrib.eager.defun,你仍然可以为部分代码使用图执行。这需要你使用TensorFlow图操作,如tf.cond()。未来,AutoGraph将与defun无缝集成,使在简单的急切执行的Python中编写图代码。当这个实现可用时,你就可以通过有选择的将急切执行代码转换为图片段,以使用AutoGraph来加速研究。
结论
AutoGraph是一款可让你轻松构建直观,复杂的模型,在TensorFlow图中轻松运行的工具。它现在是个实验性工具,但我们希望尽快将其加入到TensorFlow的核心中。
更多TensorFlow系统模型点击“这里”下载