PyTorch 的动态图一直是 TensorFlow 用户梦寐以求的功能,谷歌也一直试图在 TensorFlow 中实现相似的功能。最近,Google Brain 团队发布了 Eager Execution,一个由运行定义的新接口,让 TensorFlow 开发变得简单许多。在工具推出后,谷歌开发人员 Yaroslav Bulatov 对它的性能与 PyTorch 作了横向对比。
今天,咱们为 TensorFlow 引入了「Eager Execution」,它是一个命令式、由运行定义的接口,一旦从 Python 被调用,其操做当即被执行。这使得入门 TensorFlow 变的更简单,也使研发更直观。html
Eager Execution 的优势以下:python
Eager Execution 如今处于试用阶段,所以咱们但愿获得来自社区的反馈,指导咱们的方向。git
为了更好地理解 Eager Execution,下面让咱们看一些代码。它很技术,熟悉 TensorFlow 会有所帮助。github
使用 Eager Execution算法
当你启动 Eager Execution 时,运算会即刻执行,无需 Session.run() 就能够把它们的值返回到 Python。好比,要想使两个矩阵相乘,咱们这样写代码:编程
使用 print 或者 Python 调试器检查中间结果很是直接。性能优化
动态模型的构建可以使用 Python 控制流。下面是使用 TensorFlow 算术操做的考拉兹猜测(Collatz conjecture)的一个示例:网络
这里,tf.constant(12) 张量对象的使用将把全部数学运算提高为张量运算,从而全部的返回值将是张量。session
梯度app
多数 TensorFlow 用户对自动微分(automatic differentiation)很感兴趣。由于每次调用都有可能出现不一样的运算,能够理解为咱们把全部的正向运算录到「磁带」上,而后在计算梯度时进行「倒放」。梯度计算完成后,「磁带」就没用了。
若是你熟悉 autograd 包,咱们提供的 API 与之很是相似。例如:
gradients_function 的调用使用一个 Python 函数 square() 做为参数,而后返回 Python callable,用于计算输入的 square() 偏导数。所以,为了获得输入为 3.0 时的 square() 导数,激活 grad(3.0),也就是 6。
一样的 gradient_function 调用可用于计算 square() 的二阶导数。
如前所述,控制流(control flow)会引发不一样的运算,下面是一个示例:
自定义梯度
用户或许想为运算或函数自定义梯度。这可能有用,缘由之一是它为一系列运算提供了更高效、数值更稳定的梯度。
下面的示例使用了自定义梯度。咱们先来看函数 log(1 + e^x),它一般用于计算交叉熵和 log 似然。
咱们能够将自定义梯度应用于上述函数,简化梯度表达式。注意下面的梯度函数实现重用了前向传导中计算的 (tf.exp(x)),避免冗余计算,从而提升梯度计算的效率。
创建模型
模型能够分红几类。此处咱们要提的模型能够经过建立一个简单的两层网络对标准的 MNIST 手写数字进行分类。
咱们推荐使用 tf.layers 中的类别(而非函数),这是由于它们建立并包含了模型参数(变量,variables)。变量的有效期和层对象的有效期紧密相关,所以须要对它们进行追踪。
为何要使用 tfe.Network?一个网络包含了多个层,是 tf.layer.Layer 自己,容许将 Network 的对象嵌入到其它 Network 的对象中。它还包含可以协助检查、保存和修复的工具。
即便没有训练模型,咱们也能够命令式地调用它并检查输出:
注意咱们在这里不须要任何的占位符或会话(session)。一旦数据被输入,层的参数就被设定好了。
训练任何模型都须要定义一个损失函数,计算梯度,并使用一个优化器更新参数。首先定义一个损失函数:
而后是训练的循环过程:
implicit_gradients() 计算损失函数关于计算使用的全部 TensorFlow 变量的导数。
咱们能够按往常使用 TensorFlow 的方式将计算转移到 GPU 上:
(注意:咱们简化而后保存损失损失函数并直接调用 optimizer.minimize,但你也可使用上面的 apply_gradients() 方法,它们是等价的。)
使用 Eager 和 Graphs
Eager execution 使开发和调试互动性更强,可是 TensorFlow graph 在分布式训练、性能优化和生产部署中也有不少优点。
启用 eager execution 时,执行运算的代码还能够构建一个描述 eager execution 未启用时的计算图。为了将模型转换成图,只须要在 eager execution 未启用的 Python session 中运行一样的代码。示例:https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/mnist。咱们能够从检查点保存和修复模型变量值,这容许咱们在 eager(命令式)和 graph(声明式)编程之间轻松转换。这样,启用 eager execution 开发出的模型能够轻松导出到生产部署中。
在不久的未来,咱们将提供工具,能够选择性地将模型的某些部分转换成 graph。用这种方式,你就能够融合部分计算(如自定义 RNN 细胞的内部)实现高性能,同时还能保持 eager execution 的灵活性和可读性。
如何改写个人代码?
Eager execution 的使用方法对现有 TensorFlow 用户来讲应是直观的。目前只有少许针对 eager 的 API;大多数现有的 API 和运算须要和启用的 eager 一块儿工做。请记住如下内容:
开始使用
这只是预发布,还不完善。若是你想如今就开始使用,那么:
性能测试
Eager Execution 目前仅处于开发的前期,它的性能究竟如何?Google Brain 的工程师 Yaroslav Bulatov 对这一新工具作出了评测。
TensorFlow 此前最使人诟病的问题就是它必须将计算定义为静态图。
咱们在谷歌大脑的工做之一就是解决这类需求,并最终以命令式版本开源。可是这依赖于私有/不稳定的 API,并且这些 API 的维护成本会愈来愈高昂。
幸运的是,PyTorch 知足了研究员的需求,而且现在的 TensorFlow 也官方支持执行模式而不须要定义图。
目前,Eager Execution 仍在积极开发中,但在最近发布的可用版本很是有用,咱们能够试用一下:
请注意,此操做并不须要处理图,Session 就能够当即执行。若想应用 GPU 加速,请先将 tensor 拷贝至指定设备。
端口命令代码
你能够将一个已有的 numpy/pytorch/matlab 的命令式代码重写成正确的 API 调用。例如,
我已使用 PyTorch 实现的 l-BFGS 做为练习,第一次在 GPU 上并行跑两个实验时(PyTorch & Eager),我获得前 8 位小数相同的结果。这使我大吃一惊,前所未闻。
使用已有的基于图的代码
若是你的代码不依赖于特定的 API,例如 graph_editor,你可使用现有的代码并在 eager execution 模式下运行。
还有一个实验性的函数「graph_callable」,能够将任意 tensorflow 子图做为一个能够调用的函数。它仍然处于开发阶段,但我能获得一个有效的例子来讲明,该例子将 tensorflow /models 中的 resnet_model 包装成一个 graph_callable。下面是一个随机批大小训练这个模型的例子。
一旦该功能上线,它应该有助于提升程序性能,具体可参考下文的性能部分。
拓展了梯度
原始 tf.gradients_function 的新衍生版本反映了autograd 的梯度。你能够调用在一个已有函数内调用「gradients_function」N 次得到 N 阶导数,即
还有一个原始「custom_gradient」函数,这使得建立自定义梯度更容易。例如,假设咱们想要平方函数,但在后向传播时增长了噪声。
效果以下:
你会看到版本二收敛更慢,可是一旦收敛,它的泛化能力更好。
这种梯度修正对于实现如 KFAC 的高级优化算法时十分有用。想一想我早期所讲,KFAC 在简单网络中至关于激活函数和反向传播值白化的梯度降低。
这就能够理解为梯度在其两边乘上了白化的矩阵
假设你已经将这些矩阵保存为 m1,m2,那么你自定义的乘操做能够是这样的:
注意,true_grad1, true_grad2 函数是乘法操做的反向传播实现,请参考 Mike Giles 的第 4 页「An extended collection of matrix derivative results for forward and reverse mode algorithmic differentiation」(https://people.maths.ox.ac.uk/gilesm/files/NA-08-01.pdf)
你能够经过使用 kfac_matmul 替代采用梯度降低算法恢复原来的 kfac,或者你能够尝试新的变种方法,利用动量和 Adam。
这里(https://gist.github.com/yaroslavvb/eb02440272ddcbea549f1e47e4023376)有一个端到端的运行在 Eager execution 模式下的 KFAC 样例。
性能
Eager Execution 模式使你的程序执行慢一点或慢不少的程度取决于你的计算高运算强度的卷积仍是矩阵相乘。
作纯矩阵乘法(超过 1 毫秒的时间)是没有太大的差异,不管你用 tensorflow 快速模式,pytorch 或 tensorflow 经典模式。
另外一方面,端到端的例子更易受影响。
在测试中,当运行环境设置为 O(n^(1.5)) 操做,如 matmul/conv 时,Eager Execution 的速度要比 PyTorch 慢 20%,或者在大量 O(n) 操做如矢量添加的例子中,比 PyTorch 慢 2-5 倍。
做为一个简单的例子,咱们使用吴恩达提出的 UFLDL 来训练 MNIST 自编码器。在批尺寸=60k,I-BFGS 的 history=5 时,大量的计算效能都被花在了自编码器正向传播上,Eager 的版本要比 PyTorch 慢 1.4 倍。
在批尺寸为 60k,I-BFGS 的 history=100 的设置下,两个回环在每一步 I-BFGS(点积和向量增长)中执行「两步递归」,Eager 版本的模型速度下降了 2.5 倍,而 PyTorch 仅受轻微影响。
最后,若是咱们将批尺寸减小到 10k,咱们能够看到每次迭代的速度都要慢 5 倍,偶尔甚至会慢 10 倍,这多是由于垃圾回收策略形成的。
结论
虽然目前 Eager Execution 的表现还不够强大,但这种执行模式可让原型设计变得容易不少。对于在 TensorFlow 中构建新计算任务的开发者而言,这种方式必将很快成为主流。
原文地址:
选自Google Brain
做者:Asim Shankar & Wolff Dobson
机器之心编译
本文为机器之心编译,转载请联系本公众号得到受权。