为TensorFlow引入即时执行

Posted 2021-04-24 谷歌开发者

我们宣布为 TensorFlow 引入即时执行。即时执行是一个由运行定义的命令式接口，在这个接口中，运算将在从 Python 调用时立即执行。这样可以让 TensorFlow 的入门变得更加简单，并让研发变得更加直观。

即时执行的好处包括：

• 快速调试即时运行时错误以及与 Python 工具集成

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

• 为自定义和高阶渐变提供强大支持

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

• 即时执行目前以实验性功能形式提供，因此，我们希望获得来自社区的反馈，以便指引我们的工作方向。

为了更好地理解这一功能，我们来看一些代码。它的技术性非常强；熟悉 TensorFlow 会有所帮助。

使用即时执行

启用即时执行后，运算将立即执行并将它们的值返回到 Python，无需 Session.run()。例如，要将两个矩阵相乘，我们可以这样编写代码：

import tensorflow as tf
import tensorflow.contrib.eager as tfe

tfe.enable_eager_execution()

x = [[2.]]
m = tf.matmul(x, x)

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

print(m)
# The 1x1 matrix [[4.]]

动态模型可以通过 Python 流控制构建。下面是一个使用 TensorFlow 的算数运算的考拉兹猜想示例：

a = tf.constant(12)
counter = 0
while not tf.equal(a, 1):
  if tf.equal(a % 2, 0):
    a = a / 2
  else:
    a = 3 * a + 1
  print(a)

在这里，tf.constant(12) 张量对象的使用会将所有数学运算提升为张量运算，因此，所有返回值都是张量。

渐变

大多数 TensorFlow 用户都对自动微分感兴趣。因为每次调用都可能发生不同的运算，我们将所有正向运算记录到磁带中，然后在计算渐变时进行倒放。在计算完渐变后，我们将舍弃磁带。

如果您熟悉 autograd 软件包，就会发现 API 将非常相似。例如：

def square(x):
  return tf.multiply(x, x)

grad = tfe.gradients_function(square)

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

gradients_function 调用将 Python 函数 square() 用作一个参数并返回一个 Python 可调用对象，这个对象可以根据输入计算 square() 的偏导数。因此，要获得 square() 在输入为 3.0 时的导数，请调用 grad(3.0)，结果为 6。

相同的 gradients_function 调用可用于获得平方的二阶导数：

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

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

如前文所述，控制流会引起不同的运算，如下面的示例中所示。

def abs(x):
  return x if x > 0. else -x

grad = tfe.gradients_function(abs)

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

自定义渐变

用户可能希望为某个运算或函数定义自定义渐变。这可能非常有用，原因之一是它为一系列运算提供了一种更高效、数值更稳定的渐变。

下面的示例演示了自定义渐变的使用。我们先来看一下 log(1 + ex) 函数，它通常用于计算交叉熵和对数似然性。

def log1pexp(x):
  return tf.log(1 + tf.exp(x))
grad_log1pexp = tfe.gradients_function(log1pexp)

# The gradient computation works fine at x = 0.
print(grad_log1pexp(0.))
# [0.5]
# However it returns a `nan` at x = 100 due to numerical instability.
print(grad_log1pexp(100.))
# [nan]

我们可以为上面的函数使用自定义渐变，从分析角度简化渐变表达式。注意下面的渐变函数实现重用了在正向传递期间计算的表达式 (tf.exp(x))，通过避免冗余计算提高了渐变计算的效率。

@tfe.custom_gradient
def log1pexp(x):
  e = tf.exp(x)
  def grad(dy):
    return dy * (1 - 1 / (1 + e))
  return tf.log(1 + e), grad
grad_log1pexp = tfe.gradients_function(log1pexp)

# Gradient at x = 0 works as before.
print(grad_log1pexp(0.))
# [0.5]
# And now gradient computation at x=100 works as well.
print(grad_log1pexp(100.))
# [1.0]

构建模型

模型可以分成几类。下面的模型类创建了一个（简单的）两层网络来对标准的 MNIST 手写数字进行分类。

class MNISTModel(tfe.Network):
  def __init__(self):
    super(MNISTModel, self).__init__()
    self.layer1 = self.track_layer(tf.layers.Dense(units=10))
    self.layer2 = self.track_layer(tf.layers.Dense(units=10))
  def call(self, input):
    """Actually runs the model."""
    result = self.layer1(input)
    result = self.layer2(result)
    return result

由于 tf.layer 创建并包含模型参数（变量），我们建议在 tf.layer 中使用类，而不是函数。变量生命周期与层对象的生命周期关联，因此，务必跟踪它们。

为什么要使用 tfe.Network？一个 Network 是一个层容器并且本身是一个 tf.layer.Layer，这样，可以将 Network 对象嵌入到其他 Network 对象中。它还包含可以协助检查、保存和恢复的实用程序。

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

# Let's make up a blank input image
model = MNISTModel()
batch = tf.zeros([1, 1, 784])
print(batch.shape)
# (1, 1, 784)
result = model(batch)
print(result)
# tf.Tensor([[[ 0.  0., ...., 0.]]], shape=(1, 1, 10), dtype=float32)

请注意，我们不需要任何占位符或会话。在我们第一次传入输入时，层参数的大小就会设置好。

为了训练任何模型，我们需要定义一个损失函数来优化、计算渐变，并使用一个优化器来更新变量。首先，下面是一个损失函数：

def loss_function(model, x, y):
  y_ = model(x)
  return tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=y_)

然后，我们的训练循环如下所示：

optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
for (x, y) in tfe.Iterator(dataset):
  grads = tfe.implicit_gradients(loss_function)(model, x, y)
  optimizer.apply_gradients(grads)

implicit_gradients() 将相对于计算期间使用的所有 TensorFlow 变量计算 loss_function 的导数。

我们可以按照往常处理 TensorFlow 的相同方式将计算转移到 GPU 中：

with tf.device("/gpu:0"):
  for (x, y) in tfe.Iterator(dataset):
    optimizer.minimize(lambda: loss_function(model, x, y))

（注：我们简化了损失函数的存储并直接调用 optimizer.minimize，但是，您也可以使用上面的 apply_gradients() 函数；它们是等效的。）

结合使用即时执行与图表

即时执行让开发和调试更具互动性，但是，TensorFlow 图表在分布式训练、性能优化和生产部署方面也有很多优势。

在启用即时执行时执行运算的相同代码将构建一个图表来说明未启用即时执行时的计算。要将您的模型转换成图表，只需在未启用即时执行的新 Python 会话中运行相同的代码，如我们在 MNIST 示例中看到的一样。模型变量的值可保存并从检查点恢复，这让我们可以在即时（命令式）与图表（声明式）编程之间轻松切换。因此，可以轻松导出在启用即时执行时开发的模型，用于生产部署。

在不久的将来，我们会提供实用程序，让您可以选择性地将部分模型转换成图表。这样，您可以融合部分计算（例如自定义 RNN 细胞的内部结构）来实现高性能，同时保持即时执行的灵活性和可读性。

怎样改写我的代码？

即时执行的使用方法对当前 TensorFlow 用户来说应当是直观的。目前只有少量特定于即时执行的 API；大多数的现有 API 和运算都适用于即时执行。下面是一些需要注意的方面：

• 一如我们通常对 TensorFlow 的建议，如果您还没有从队列转到使用 tf.data 来处理输入，我们建议您立即行动。TensorFlow 的使用更加简单，并且速度通常更快。

• 使用以对象为导向的层（例如 tf.layer.Conv2D()）或 Keras 层；这些层可以显式存储变量。

• 对于大多数模型，您可以编写代码，让它同时适用于即时执行和图表构建。也有一些例外，例如，使用 Python 控制流根据输入改变计算的动态模型就不行。

• 一旦您调用 tfe.enable_eager_execution()，它将无法关闭。要获取图表行为，请启动新的 Python 会话。

入门和未来

这仍是一个预览版，因此，您可能会遇到一些不尽人意的地方。如果您想立即开始体验，请执行以下操作：

• 安装 Nightly 版本的 TensorFlow。

• 查看 README（包括已知问题）

• 在即时执行用户指南中获取详细的说明

• 浏览 GitHub 中的即时执行示例

• 关注 changelog 获取更新。

关于即时执行，我们还有很多要与大家分享的，我们非常高兴，或者说非常渴望您立即试用这一功能。热烈欢迎大家提供反馈。