欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

TensorFlow引入了动态图机制Eager Execution

程序员文章站 2022-03-09 20:21:45
...

今天,我们为 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. tfe.enable_eager_execution()

  4. x = [[2.]]

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



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


  1. print(m)

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



动态模型的构建可使用 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)



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


梯度


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


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


  1. def square(x):

  2.  return tf.multiply(x, x)

  3. grad = tfe.gradients_function(square)

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

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



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. print(gradgrad(3.))  # [2.]



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


  1. def abs(x):

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

  3. grad = tfe.gradients_function(abs)

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

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



自定义梯度


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


下面的示例使用了自定义梯度。我们先来看函数 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. # The gradient computation works fine at x = 0.

  5. print(grad_log1pexp(0.))

  6. # [0.5]

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

  8. print(grad_log1pexp(100.))

  9. # [nan]



我们可以将自定义梯度应用于上述函数,简化梯度表达式。注意下面的梯度函数实现重用了前向传导中计算的 (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. # Gradient at x = 0 works as before.

  9. print(grad_log1pexp(0.))

  10. # [0.5]

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

  12. print(grad_log1pexp(100.))

  13. # [1.0]



建立模型


模型可以分成几类。此处我们要提的模型可以通过创建一个简单的两层网络对标准的 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



我们推荐使用 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)




注意我们在这里不需要任何的占位符或会话(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_)



然后是训练的循环过程:


  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)



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))



(注意:我们简化然后保存损失损失函数并直接调用 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。

相关标签: 动态图机制