使用 Eager Execution 编码并运行图表：以通过 RevNet 优化代码为例

文 / 软件工程实习生 Xuechen Li

Eager Execution 可简化 TensorFlow 中的模型构建体验，而 Graph Execution 可提供优化，以加快模型运行速度及提高存储效率。本篇博文展示了如何编写 TensorFlow 代码，以便将借助 tf.keras API 并使用 Eager Execution 构建的模型转换为图表，最终借助 tf.estimator API 的支持，在 Cloud TPU 上部署此模型。

注：tf.keras 链接

https://www.tensorflow.org/guide/keras

tf.estimator 链接

https://www.tensorflow.org/guide/estimators

我们使用可逆残差网络（RevNet、Gomez 等）作为示例。接下来的部分假设读者对卷积神经网络和 TensorFlow 有基本了解。您可以在此处找到本文的完整代码（为确保代码在所有设置中正常运行，强烈建议您使用 tf-nightly 或 tf-nightly-gpu）。

RevNet

RevNet 与残差网络（ResNet、He 等）类似，只不过他们是可逆的，在给定输出的情况下可重建中间计算。此技术的好处之一是我们可以通过重建激活来节省内存，而不是在训练期间将其全部存储在内存中（回想一下，由于链式法则有此要求，因此我们需要中间结果来计算有关输入的梯度）。相比传统架构上的一般反向传播，这使我们可以适应较大的批次大小，并可训练更具深度的模型。具体来说，此技术的实现方式是通过使用一组巧妙构建的方程来定义网络：

其中顶部和底部方程组分别定义正演计算和其反演计算。这里的 x1 和 x2 是输入（从整体输入 x 中拆分出来），y1 和 y2 是输出，F 和 G 是 ConvNet。这使我们能够在反向传播期间精准重建激活，如此一来，在训练期间便无需再存储这些数据。

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

假设我们使用 “ResidualInner” 类来实例化函数 F 和 G，我们可以通过子类化 tf.keras.Model 来定义可逆代码块，并通过替换上面的方程中所示的 call 方法来定义正向传递：

1    class Residual(tf.keras.Model):    

2        def __init__(self, filters):    

3            super(Residual, self).__init__()    

4            self.f = ResidualInner(filters=filters, strides=(1, 1))

5            self.g = ResidualInner(filters=filters, strides=(1, 1))

6

7        def call(self, x, training=True):    

8            x1, x2 = tf.split(x, num_or_size_splits=2, axis=self.axis)

9            f_x2 = self.f(x2, training=training)    

10            y1 = f_x2 + x1    

11            g_y1 = self.g(y1, training=training)    

12            y2 = g_y1 + x2    

13            return tf.concat([y1, y2], axis=self.axis) 

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

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

注：tf.GradientTape 链接

https://www.tensorflow.org/api_docs/python/tf/GradientTape

1        def backward_grads(self, y, dy, training=True):    

2            dy1, dy2 = dy    

3            y1, y2 = y

4

5            with tf.GradientTape() as gtape:    

6                gtape.watch(y1)    

7                gy1 = self.g(y1, training=training)    

8            grads_combined = gtape.gradient(    

9                    gy1, [y1] + self.g.trainable_variables, output_gradients=dy2)    

10            dg = grads_combined[1:]    

11            dx1 = dy1 + grads_combined[0]

12            x2 = y2 - gy1

13

14            with tf.GradientTape() as ftape:    

15                ftape.watch(x2)    

16                fx2 = self.f(x2, training=training)    

17            grads_combined = ftape.gradient(    

18                    fx2, [x2] + self.f.trainable_variables,output_gradients=dx1)    

19            df = grads_combined[1:]    

20            dx2 = dy2 + grads_combined[0] 

21            x1 = y1 - fx2

22

23            x = x1, x2    

24            dx = dx1, dx2    

25            grads = df + dg 

26

27            return x, dx, grads    

您可以在论文的 “算法 1” 中找到确切的一组梯度计算（我们在代码中简化了使用变量 z1 的中间步骤）。此算法经过精心设计，在给定输出和有关输出的损失梯度的情况下，我们可以在每个可逆代码块内，计算有关输入和模型变量的梯度及重建输入。调用 tape.gradient(y, x)，即可计算有关 x 的 y 梯度。我们也可使用参数 output_gradients 来明确应用链式法则。

使用 Eager Execution 来加快原型设计速度

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

例如，我们通过由一般反向传播计算的梯度来比较可逆反向传播梯度，从而验证我们的模型：

1    block = Residual()    

2    x = tf.random_normal(shape=(N, C, H, W))    

3    dy = tf.random_normal(shape=(N, C, H, W))    

4    with tf.GradientTape() as tape:    

5        tape.watch(x)    

6        y = block(x)    

7    # Compute true grads    

8    dx_true = tape.gradient(y, x, output_gradients=dy)

9

10    # Compute grads from reconstruction    

11    dx, _ = block.backward_grads(x, y, dy)

12

13    # Check whether the difference is below a certain 14    threshold    

thres = 1e-6    

15    diff_abs = tf.reshape(abs(dx - dx_true), [-1])    

16    assert all(diff_abs < thres)  

在上面的片段中，dx_true 是一般反向传播返回的梯度，而 dx 是执行可逆反向传播后返回的梯度。Eager Execution 整合了原生 Python，如此一来，all 和 abs 等函数便可直接应用于 Tensor。

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

为确保能够使用 Eager Execution 和 Graph Execution 存储和加载检查点，TensorFlow 团队建议您使用 tf.train.Checkpoint API。

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

1    checkpoint = tf.train.Checkpoint(model=model, optimizer=optimizer,    

2                        learning_rate=learning_rate, global_step=global_step)

我们可以按照下面的方法存储和还原特定的已训练实例：

1    checkpoint.save(file_prefix)    

2    checkpoint.restore(save_path)    

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

由于解读 Python 代码会产生开销，Eager Execution 有时会比执行等效图要慢。通过使用 tf.contrib.eager.defun 将由 TensorFlow 运算组成的 Python 函数编译成可调用的 TensorFlow 图表，可以弥补这种性能差距。在训练深度学习模型时，我们通常可以在三个主要位置应用 tf.contrib.eager.defun：

• 1) 正演计算

• 2) 梯度的反演计算

• 3) 将梯度应用于变量

例如，我们可以按以下方式 defun 正向传递和梯度计算：

1    tfe = tf.contrib.eager    

2    model.call = tfe.defun(model.call)    

3    model.compute_gradients = tfe.defun(model.compute_gradients)

要 defun 优化器的应用梯度步骤，我们需要将其包装在另一个函数内：

1    def apply_gradients(optimizer, gradients, variables, global_step=None):    

2            optimizer.apply_gradients(    

3                    zip(gradients, variables), global_step=global_step) 

4    apply_gradients = tfe.defun(apply_gradients)

tf.contrib.eager.defun 正处于积极开发中，将其加以应用是一项不断发展的技术。如需更多信息，请查看其文档字符串。

注：其文档字符串链接

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/eager/function.py#L1147

使用 tf.contrib.eager.defun 包装 Python 函数会使 TensorFlow API 在 Python 函数中进行调用，以构建图表，而不是立即执行运算，从而优化整个程序。并非所有 Python 函数都可成功转换为等效图，特别是带有动态控制流的函数（例如，Tensor contents 中的 if 或 while）。tf.contrib.autograph 是一种相关工具，可以增加能够转换为 TensorFlow 图表的 Python 代码的表面积。截至 2018 年 8 月，使用 defun 集成 Autograph 的工作仍在进行中。

注：tf.contrib.autograph 链接

https://www.tensorflow.org/guide/autograph

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

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

1    dataset = tf.data.TFRecordDataset(filename) 

2    dataset = dataset.repeat(epochs).map(parser).batch(batch_size)

为提升性能，我们还可使用预取函数并调整 num_parallel_calls。

由于数据集由图像和标签对组成，在 Eager Execution 中循环使用此数据集非常简单。在本例中，我们甚至不需要明确定义迭代器：

1    for image, label in dataset:    

2        logits = model(image, training=True)    

3        ... 

使用估算器包装 Keras 模型并以图表形式执行

由于 tf.keras API 也支持图表构建，因此使用 Eager Execution 构建的相同模型也可用作提供给估算器的图表构建函数，但代码稍有更改。要修改使用 Eager Execution 构建的 RevNet 示例，我们只需使用 model_fn 包装 Keras 模型，并按照 tf.estimator API 的指示使用此模型。

1    def model_fn(features, labels, mode, params):

2        model = RevNet(params["hyperparameters"])

3        if mode == tf.estimator.ModeKeys.TRAIN: 

4            optimizer = tf.train.MomentumOptimizer(learning_rate, momentum)

5            logits, saved_hidden = model(features, training=True)    

6            grads, loss = model.compute_gradients(saved_hidden, labels, training=True) 

7            with tf.control_dependencies(model.get_updates_for(features)):

8                train_op = optimizer.apply_gradients(zip(grads, model.trainable_variables)) 

9            return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op) 

您可以使用 the tf.data API 照常定义 tf.estimator API 所需的 input_fn，并从 TFRecords 中读取数据。

使用 TPU Estimator 包装 Keras 模型以进行 Cloud TPU 训练

使用 Estimator 包装模型和输入管道使模型可以在 Cloud TPU 上运行。

所需步骤如下：

• 设置 Cloud TPU 的特定配置

• 从 tf.estimator.Estimator 切换到 tf.contrib.tpu.TPUEstimator

• 使用 tf.contrib.tpu.CrossShardOptimizer 包装常用优化器

注：Cloud TPU 链接

https://github.com/tensorflow/tpu

配置链接

https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUConfig

tf.contrib.tpu.TPUEstimator 链接

https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/TPUEstimator

tf.contrib.tpu.CrossShardOptimizer 链接

https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/CrossShardOptimizer

如需了解具体说明，请查看 RevNet 示例文件夹中的 TPU 估算器脚本。我们希望日后可以使用 tf.contrib.tpu.keras_to_tpu_model 进一步简化使 Keras 模型在 TPU 上运行的流程。

注：TPU 估算器脚本链接

https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/eager/python/examples/revnet/main_estimator_tpu.py

tf.contrib.tpu.keras_to_tpu_model 链接

https://github.com/tensorflow/tpu/tree/master/models/experimental/keras

可选：模型性能

与一般反向传播相比，有了 tf.GradientTape，再加上无需额外正向传递的梯度计算可简化流程，我们能够仅以 25% 的计算开销来执行 RevNet 的可逆反向传播。

图中蓝色和橙色的曲线分别表示随着全局步骤的增加，一般反向传播和可逆反向传播的每秒采样率。该图来自在单个 Tesla P100 上使用批次大小为 32 的模拟 ImageNet 数据训练的 RevNet-104。

为了验证所节省的内存，我们在训练过程中绘制内存使用情况。蓝色和黑色曲线分别是一般和可逆反向传播。该图记录了使用批次大小为 128 的模拟 ImageNet 数据训练 RevNet-104 图表模式的 100 次迭代。该图是在 CPU 上进行训练时由 mprof 生成，以便我们使用一般反向传播以相同的批次大小进行训练。

结论

我们以 RevNet 为例，展示了如何使用 Eager Execution 和 tf.keras API 对机器学习模型快速进行原型设计。这不仅可简化模型构建体验，而且我们轻易就能将模型转换为估算器，并在 Cloud TPU 上进行部署，以获得高性能。您可以在此处找到本文的完整代码。此外，请务必查看使用 Eager Execution 的其他示例。

注：此处连接诶

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples/revnet

其他示例链接

https://github.com/tensorflow/tensorflow/tree/master/tensorflow/contrib/eager/python/examples

更多 AI 相关阅读：