今天,TensorFlow团队发布新功能“AutoGraph”,能自动将Python代码(包括控制流,print () 和其他Python原生特征)转换为纯TensorFlow图代码(pure TensorFlow graph code)。
不使用 Eager Execution编写TensorFlow代码需要进行一些元编程(metaprogramming) ——先编写一个创建图(Graph)的程序,稍后再执行这个Graph。这可能令人困惑,尤其是对开发者新手来说。一些特别棘手的情况涉及更复杂的模型,比如要使用 if 和 while 的模型,或者有 print () 等副作用或接受结构化输入的模型。
为什么我们需要Graph呢?Graph允许各种优化,例如删除常见的子表达式和融合内核(fusing kernel)。再者,Graph简化了分布式训练和部署到各种环境的过程,因为它们形成了独立于平台的模型计算过程。这对于模型在多个GPU或TPU上的分布式训练尤为重要,如果你通过TensorFlow Lite、移动端、物联网等其他平台分发模型,Graph也很重要。
下面是一个很简单的、你可能希望添加到Graph里的操作:
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解释器开销(interpreter overheads)或错过的程序优化机会,此类操作可能会很慢。
为了准备执行Graph,你需要重写这个以使用像 tf.cond () 这样的结构,但那样实现起来可能会耗时耗力而且很困难。AutoGraph可以为自动执行此类转换,将动态图编程的简易性保持很低的同时,获得基于Graph执行的性能优势。
在示例中,我们可以使用 autograph.convert () 来修饰函数,AutoGraph将自动生成 graph-ready 的代码。