TensorFlow引入了动态图机制Eager Execution头条

ctolib 发布于1年前 阅读141次
0 条评论

PyTorch 的动态图一直是 TensorFlow 用户求之不得的功能,谷歌也一直试图在 TensorFlow 中实现类似的功能。最近,Google Brain 团队发布了 Eager Execution,一个由运行定义的新接口,让 TensorFlow 开发变得简单许多。在工具推出后,谷歌开发人员 Yaroslav Bulatov 对它的性能与 PyTorch 做了横向对比。

 

今天,我们为 TensorFlow 引入了「Eager Execution」,它是一个命令式、由运行定义的接口,一旦从 Python 被调用,其操作立即被执行。这使得入门 TensorFlow 变的更简单,也使研发更直观。

 

Eager Execution 的优点如下:

 

  • 快速调试即刻的运行错误并通过 Python 工具进行整合

  • 借助易于使用的 Python 控制流支持动态模型

  • 为自定义和高阶梯度提供强大支持

  • 适用于几乎所有可用的 TensorFlow 运算

 

Eager Execution 现在处于试用阶段,因此我们希望得到来自社区的反馈,指导我们的方向。

 

为了更好地理解 Eager Execution,下面让我们看一些代码。它很技术,熟悉 TensorFlow 会有所帮助。

 

使用 Eager Execution

 

当你启动 Eager Execution 时,运算会即刻执行,无需 Session.run() 就可以把它们的值返回到 Python。比如,要想使两个矩阵相乘,我们这样写代码:

 

 


 
  1. import tensorflow as tf

  2. import tensorflow.contrib.eager as tfe

  3.  

  4. tfe.enable_eager_execution()

  5.  

  6. x = [[2.]]

  7. m = tf.matmul(x, x)

  8.  

 

使用 print 或者 Python 调试器检查中间结果非常直接。

 

 


 
  1. print(m)

  2. # The 1x1 matrix [[4.]]

  3.  

 

动态模型的构建可使用 Python 控制流。下面是使用 TensorFlow 算术操作的考拉兹猜想(Collatz conjecture)的一个示例:

 

 


 
  1. a = tf.constant(12)

  2. counter = 0

  3. while not tf.equal(a, 1):

  4.  if tf.equal(a % 2, 0):

  5.    a = a / 2

  6.  else:

  7.    a = 3 * a + 1

  8.  print(a)

  9.  

 

这里,tf.constant(12) 张量对象的使用将把所有数学运算提升为张量运算,从而所有的返回值将是张量。

 

梯度

 

多数 TensorFlow 用户对自动微分(automatic differentiation)很感兴趣。因为每次调用都有可能出现不同的运算,可以理解为我们把所有的正向运算录到「磁带」上,然后在计算梯度时进行「倒放」。梯度计算完成后,「磁带」就没用了。

 

如果你熟悉 autograd 包,我们提供的 API 与之非常类似。例如:

 


 
  1. def square(x):

  2.  return tf.multiply(x, x)

  3.  

  4. grad = tfe.gradients_function(square)

  5.  

  6. print(square(3.))    # [9.]

  7. print(grad(3.))      # [6.]

  8.  

 

gradients_function 的调用使用一个 Python 函数 square() 作为参数,然后返回 Python callable,用于计算输入的 square() 偏导数。因此,为了得到输入为 3.0 时的 square() 导数,激活 grad(3.0),也就是 6。

 

同样的 gradient_function 调用可用于计算 square() 的二阶导数。

 

 


 
  1. gradgrad = tfe.gradients_function(lambda x: grad(x)[0])

  2.  

  3. print(gradgrad(3.))  # [2.]

  4.  

 

如前所述,控制流(control flow)会引起不同的运算,下面是一个示例:

 


 
  1. def abs(x):

  2.  return x if x > 0. else -x

  3.  

  4. grad = tfe.gradients_function(abs)

  5.  

  6. print(grad(2.0))  # [1.]

  7. print(grad(-2.0)) # [-1.]

  8.  

 

自定义梯度

 

用户或许想为运算或函数自定义梯度。这可能有用,原因之一是它为一系列运算提供了更高效、数值更稳定的梯度。

 

下面的示例使用了自定义梯度。我们先来看函数 log(1 + e^x),它通常用于计算交叉熵和 log 似然。

 


 
  1. def log1pexp(x):

  2.  return tf.log(1 + tf.exp(x))

  3. grad_log1pexp = tfe.gradients_function(log1pexp)

  4.  

  5. # The gradient computation works fine at x = 0.

  6. print(grad_log1pexp(0.))

  7. # [0.5]

  8. # However it returns a `nan` at x = 100 due to numerical instability.

  9. print(grad_log1pexp(100.))

  10. # [nan]

  11.  

 

我们可以将自定义梯度应用于上述函数,简化梯度表达式。注意下面的梯度函数实现重用了前向传导中计算的 (tf.exp(x)),避免冗余计算,从而提高梯度计算的效率。

 

 


 
  1. @tfe.custom_gradient

  2. def log1pexp(x):

  3.  e = tf.exp(x)

  4.  def grad(dy):

  5.    return dy * (1 - 1 / (1 + e))

  6.  return tf.log(1 + e), grad

  7. grad_log1pexp = tfe.gradients_function(log1pexp)

  8.  

  9. # Gradient at x = 0 works as before.

  10. print(grad_log1pexp(0.))

  11. # [0.5]

  12. # And now gradient computation at x=100 works as well.

  13. print(grad_log1pexp(100.))

  14. # [1.0]

  15.  

 

建立模型

 

模型可以分成几类。此处我们要提的模型可以通过创建一个简单的两层网络对标准的 MNIST 手写数字进行分类。

 

 


 
  1. class MNISTModel(tfe.Network):

  2.  def __init__(self):

  3.    super(MNISTModel, self).__init__()

  4.    self.layer1 = self.track_layer(tf.layers.Dense(units=10))

  5.    self.layer2 = self.track_layer(tf.layers.Dense(units=10))

  6.  def call(self, input):

  7.    """Actually runs the model."""

  8.    result = self.layer1(input)

  9.    result = self.layer2(result)

  10.    return result

  11.  

 

我们推荐使用 tf.layers 中的类别(而非函数),这是因为它们创建并包含了模型参数(变量,variables)。变量的有效期和层对象的有效期紧密相关,因此需要对它们进行追踪。

 

为什么要使用 tfe.Network?一个网络包含了多个层,是 tf.layer.Layer 本身,允许将 Network 的对象嵌入到其它 Network 的对象中。它还包含能够协助检查、保存和修复的工具。

 

即使没有训练模型,我们也可以命令式地调用它并检查输出:

 

 


 
  1. # Let's make up a blank input image

  2. model = MNISTModel()

  3. batch = tf.zeros([1, 1, 784])

  4. print(batch.shape)

  5. # (1, 1, 784)

  6. result = model(batch)

  7. print(result)

  8. # tf.Tensor([[[ 0.  0., ...., 0.]]], shape=(1, 1, 10), dtype=float32)

  9.  

 

 

注意我们在这里不需要任何的占位符或会话(session)。一旦数据被输入,层的参数就被设定好了。

 

训练任何模型都需要定义一个损失函数,计算梯度,并使用一个优化器更新参数。首先定义一个损失函数:

 


 
  1. def loss_function(model, x, y):

  2.  y_ = model(x)

  3.  return tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=y_)

  4.  

 

然后是训练的循环过程:

 

 


 
  1. optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)

  2. for (x, y) in tfe.Iterator(dataset):

  3.  grads = tfe.implicit_gradients(loss_function)(model, x, y)

  4.  optimizer.apply_gradients(grads)

  5.  

 

implicit_gradients() 计算损失函数关于计算使用的所有 TensorFlow 变量的导数。

 

我们可以按往常使用 TensorFlow 的方式将计算转移到 GPU 上:

 


 
  1. with tf.device("/gpu:0"):

  2.  for (x, y) in tfe.Iterator(dataset):

  3.    optimizer.minimize(lambda: loss_function(model, x, y))

  4.  

 

(注意:我们简化然后保存损失损失函数并直接调用 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 一起工作。请记住以下内容:

 

  • 一般对于 TensorFlow,我们建议如果你还没有从排队切换到使用 tf.data 进行输入处理,请抓紧做。它更容易使用,也更快。查看这篇博文(https://developers.googleblog.com/2017/09/introducing-tensorflow-datasets.html)和文档页(https://www.tensorflow.org/programmers_guide/datasets)会有所帮助。

  • 使用目标导向的层(比如 tf.layer.Conv2D() 或者 Keras 层),它们可以直接存储变量。

  • 你可以为大多数模型写代码,这对 eager execution 和图构建同样有效。也有一些例外,比如动态模型使用 Python 控制流改变基于输入的计算。

  • 一旦调用 tfe.enable_eager_execution(),它不可被关掉。为了获得图行为,需要建立一个新的 Python session。

 

开始使用

 

这只是预发布,还不完善。如果你想现在就开始使用,那么:

 

  • 安装 TensorFlow 的 nightly 版本(https://github.com/tensorflow/tensorflow#installation)

  • 查看 README(包括 known issues),地址:https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/README.md

  • 从 eager execution 用户指南(https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/g3doc/guide.md)中获取详细的指导

  • 在 GitHub 中查看 eager 示例(https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples)

  • 及时查看变更日志(https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/README.md#changelog)查看是否有更新

需要 登录 后回复方可回复, 如果你还没有账号你可以 注册 一个帐号。