作者：Xuechen Li，软件工程实习生

概述

Eager Execution 简化了 TensorFlow 中的模型构建体验，而图执行可以提供优化，使模型运行更快，内存效率更高。这篇博文展示了如何编写 TensorFlow 代码，以便使用 Eager Execution 和 tf.keras API 构建的模型可以转换为图，并最终使用 tf.estimator API 的支持部署在云 TPU 上。

我们使用可逆残差网络 (RevNet，Gomez 等人) 作为示例。以下部分假设您对卷积神经网络和 TensorFlow 有基本了解。本文的完整代码位于 此处（为了确保代码在所有设置中都能正常工作，强烈建议使用 tf-nightly 或 tf-nightly-gpu）。

RevNets

RevNets 类似于残差网络 (ResNet，He 等人)，区别在于它们是可逆的——给定输出，可以重建中间计算。其中一个好处是，我们可以通过重建激活来节省内存，而不是在训练期间将所有激活都存储在内存中（回想一下，我们需要中间结果来计算相对于输入的梯度，因为链式法则需要这样做）。这使我们能够与传统架构上的常规反向传播相比，拟合更大的批次大小并训练更深的模型。具体来说，这是通过使用一组巧妙构建的方程来定义网络来实现的
其中上面和下面的方程分别定义了正向计算及其逆运算。这里 x1 和 x2 是输入（从整体输入 x 中分割），y1 和 y2 是输出，F 和 G 是 ConvNets。这使我们能够在反向传播期间精确地重建激活，因此在训练期间我们不再需要存储它们。

使用 tf.keras.Model 定义正向和反向传递

假设我们有类“ResidualInner”来实例化函数 F 和 G，我们可以通过从 tf.keras.Model 子类化并通过覆盖 call 方法（如上面的方程所示）来定义可逆块和正向传递

class Residual(tf.keras.Model):
  def __init__(self, filters):
    super(Residual, self).__init__()
    self.f = ResidualInner(filters=filters, strides=(1, 1))
    self.g = ResidualInner(filters=filters, strides=(1, 1))

  def call(self, x, training=True):
    x1, x2 = tf.split(x, num_or_size_splits=2, axis=self.axis)
    f_x2 = self.f(x2, training=training)
    y1 = f_x2 + x1
    g_y1 = self.g(y1, training=training)
    y2 = g_y1 + x2
    return tf.concat([y1, y2], axis=self.axis)

这里的 training 参数用于确定批次归一化的状态。启用 Eager Execution 后，当 training=True 时，批次归一化的运行平均值会自动更新。当执行等效的图时，需要使用 get_updates_for 方法手动获取批次归一化更新。

为了构建节省内存的反向传递，我们使用 tf.GradientTape 作为上下文管理器，仅在需要时跟踪梯度

  def backward_grads(self, y, dy, training=True):
    dy1, dy2 = dy
    y1, y2 = y

    with tf.GradientTape() as gtape:
      gtape.watch(y1)
      gy1 = self.g(y1, training=training)
    grads_combined = gtape.gradient(
        gy1, [y1] + self.g.trainable_variables, output_gradients=dy2)
    dg = grads_combined[1:]
    dx1 = dy1 + grads_combined[0]
    x2 = y2 - gy1

    with tf.GradientTape() as ftape:
      ftape.watch(x2)
      fx2 = self.f(x2, training=training)
    grads_combined = ftape.gradient(
        fx2, [x2] + self.f.trainable_variables, output_gradients=dx1)
    df = grads_combined[1:]
    dx2 = dy2 + grads_combined[0]
    x1 = y1 - fx2

    x = x1, x2
    dx = dx1, dx2
    grads = df + dg

    return x, dx, grads

梯度计算的精确集可以在论文的算法 1 中找到（我们在代码中简化了使用变量 z1 的中间步骤）。该算法的设计使得在每个可逆块中，相对于输入和模型变量的梯度随着输入的重建而计算，给定输出和损失相对于输出的梯度。调用 tape.gradient(y, x) 计算 y 相对于 x 的梯度。我们也可以使用参数 output_gradients 来显式地应用链式法则。

使用 Eager Execution 加速原型设计

使用 Eager Execution 进行原型设计的一个显而易见的好处是，它是命令式的。我们可以立即获得结果，而不是先构建一个图，然后初始化一个会话来运行。

例如，我们通过比较可逆反向传播梯度和常规反向传播计算的梯度来验证我们的模型

block = Residual()
x = tf.random_normal(shape=(N, C, H, W))
dy = tf.random_normal(shape=(N, C, H, W))
with tf.GradientTape() as tape:
  tape.watch(x)
  y = block(x)
# Compute true grads
dx_true = tape.gradient(y, x, output_gradients=dy)

# Compute grads from reconstruction
dx, _ = block.backward_grads(x, y, dy)

# Check whether the difference is below a certain threshold
thres = 1e-6
diff_abs = tf.reshape(abs(dx - dx_true), [-1])
assert all(diff_abs < thres)

在上面的代码段中，dx_true 是常规反向传播返回的梯度，而 dx 是我们实现的可逆反向传播返回的梯度。Eager Execution 与原生 Python 集成在一起，因此可以将 all 和 abs 之类的函数直接应用于张量。

使用 tf.train.Checkpoint 存储和加载检查点

为了确保保存和加载检查点与 Eager Execution 和图执行都能正常工作，TensorFlow 团队建议使用 tf.train.Checkpoint API。

为了存储模型，我们使用要存储的所有对象创建一个 tf.train.Checkpoint 实例。这可能包括我们的模型、我们使用的优化器、学习率计划以及全局步骤

checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer,
        learning_rate=learning_rate, global_step=global_step)

我们可以按如下方式保存和恢复特定训练实例

checkpoint.save(file_prefix)
checkpoint.restore(save_path)

使用 tf.contrib.eager.defun 提升 Eager Execution 的性能

Eager Execution 有时可能比执行等效的图慢，这是因为解释 Python 代码会产生开销。可以通过使用 tf.contrib.eager.defun 将由 TensorFlow 操作组成的 Python 函数编译为可调用的 TensorFlow 图来弥合这种性能差距。在训练深度学习模型时，通常有三个主要地方可以使用 tf.contrib.eager.defun：1）正向计算，2）梯度的反向计算，以及 3）将梯度应用于变量。例如，我们可以对正向传递和梯度计算进行 defun，如下所示

tfe = tf.contrib.eager
model.call = tfe.defun(model.call)
model.compute_gradients = tfe.defun(model.compute_gradients)

为了对优化器的应用梯度步骤进行 defun，我们需要将其包装在另一个函数中

def apply_gradients(optimizer, gradients, variables, global_step=None):
    optimizer.apply_gradients(
        zip(gradients, variables), global_step=global_step)
apply_gradients = tfe.defun(apply_gradients)

tf.contrib.eager.defun 正在积极开发中，使用它是一种不断发展的技术；有关更多信息，请查阅 其文档字符串。

使用 tf.contrib.eager.defun 包装 Python 函数会导致 Python 函数中的 TensorFlow API 调用构建一个图，而不是立即执行操作，从而实现整个程序优化。并非所有 Python 函数都能成功地转换为等效的图，特别是那些具有动态控制流的函数（例如，在 Tensor contents 上的 if 或 while）。tf.contrib.autograph 是一种相关工具，它可以增加可以转换为 TensorFlow 图的 Python 代码的覆盖范围。截至 2018 年 8 月，autograph 与 defun 的集成正在进行中。

使用 TFRecords 和 tf.data.Dataset 构建输入管道

Eager Execution 与 tf.data.Dataset API 兼容。我们可以读取 TFRecords 文件

dataset = tf.data.TFRecordDataset(filename)
dataset = dataset.repeat(epochs).map(parser).batch(batch_size)

为了提高性能，我们还可以使用 prefetch 函数并调整 num_parallel_calls。

在 Eager Execution 中遍历此数据集很简单，因为数据集包含图像、标签对。在这种情况下，我们甚至不需要显式地定义迭代器

for image, label in dataset:
  logits = model(image, training=True)
  ...

将 Keras 模型包装在 Estimator 中，并作为图执行

由于 tf.keras API 也支持图构建，因此使用 Eager Execution 构建的相同模型也可以用作提供给 Estimator 的图构造函数，代码只需做少量更改。要修改在 Eager Execution 中构建的 RevNet 示例，我们只需要将 keras 模型包装在一个 model_fn 中，并根据 tf.estimator API 使用它。

def model_fn(features, labels, mode, params):
  model = RevNet(params["hyperparameters"])
  if mode == tf.estimator.ModeKeys.TRAIN:
    optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)
    logits, saved_hidden = model(features, training=True)
    grads, loss = model.compute_gradients(saved_hidden, labels, training=True)
    with tf.control_dependencies(model.get_updates_for(features)):
      train_op = optimizer.apply_gradients(zip(grads, model.trainable_variables))
    return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)

tf.estimator API 所需的 input_fn 可以像往常一样使用 tf.data API 定义，从 TFRecords 读取。

将 Keras 模型包装在 TPU Estimator 中，用于云 TPU 训练

将模型和输入管道包装在 Estimator 中，可以使模型在 云 TPU 上运行。

所需步骤如下

1. 设置云 TPU 特定的 配置
   
2. 从 tf.estimator.Estimator 切换到 tf.contrib.tpu.TPUEstimator
   
3. 将通常的优化器包装在 tf.contrib.tpu.CrossShardOptimizer 中
   

有关具体演示，请查看 RevNet 示例文件夹中的 TPU estimator 脚本。我们预计在将来使用 tf.contrib.tpu.keras_to_tpu_model 将进一步简化使 Keras 模型在 TPU 上运行的过程。

可选：模型性能

tf.GradientTape 与梯度计算的简化相结合，无需额外的正向传递，这使我们能够实现 RevNet 的可逆反向传播，其计算开销仅为常规反向传播的 25%。

蓝色和橙色曲线分别表示常规反向传播和可逆反向传播的示例/秒，随着全局步骤的增加。该图来自在模拟 ImageNet 数据上训练的 RevNet-104，批次大小为 32，使用单个 Tesla P100。 为了验证内存节省，我们在训练过程中绘制内存使用情况。蓝色和黑色曲线分别表示常规反向传播和可逆反向传播。该图记录了 RevNet-104 图模式训练 100 次迭代，模拟 ImageNet 数据，批次大小为 128。该图由 mprof 生成，训练在 CPU 上执行，因此我们可以在常规反向传播中使用相同的批次大小。

结论

以 RevNet 为例，我们演示了如何使用 Eager Execution 和 tf.keras API 快速地对机器学习模型进行原型设计。这简化了模型构建体验，此外，只需做少量额外工作，我们就可以将模型转换为 Estimator 并将其部署在云 TPU 上，以实现高性能。您可以在 此处 找到本文的完整代码。此外，请务必查看使用 Eager Execution 的 其他示例。