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 方法来定义正向传递:
- 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 作为上下文管理器来跟踪梯度(仅在有需要时): 注: tf.GradientTape 链接 https://www.tensorflow.org/api_docs/python/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), 即可计算有关 x 的 y 梯度. 我们也可使用参数 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 14 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 等函数便可直接应用于 Tensor.
使用 tf.train.Checkpoint 存储和加载检查点
为确保能够使用 Eager Execution 和 Graph 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 性能
由于解读 Python 代码会产生开销, Eager Execution 有时会比执行等效图要慢. 通过使用 tf.contrib.eager.defun 将由 TensorFlow 运算组成的 Python 函数编译成可调用的 TensorFlow 图表, 可以弥补这种性能差距. 在训练深度学习模型时, 我们通常可以在三个主要位置应用 tf.contrib.eager.defun:
正演计算
梯度的反演计算
将梯度应用于变量
例如, 我们可以按以下方式 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 正处于积极开发中, 将其加以应用是一项不断发展的技术. 如需更多信息, 请查看其文档字符串. 注: 其文档字符串链接 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 文件:
- dataset = tf.data.TFRecordDataset(filename)
- dataset = dataset.repeat(epochs).map(parser).batch(batch_size)
复制代码
为提升性能, 我们还可使用预取函数并调整 num_parallel_calls.
由于数据集由图像和标签对组成, 在 Eager Execution 中循环使用此数据集非常简单. 在本例中, 我们甚至不需要明确定义迭代器:
- for image, label in dataset:
- logits = model(image, training=True)
- ...
复制代码
使用估算器包装 Keras 模型并以图表形式执行
由于 tf.keras API 也支持图表构建, 因此使用 Eager Execution 构建的相同模型也可用作提供给估算器的图表构建函数, 但代码稍有更改. 要修改使用 Eager Execution 构建的 RevNet 示例, 我们只需使用 model_fn 包装 Keras 模型, 并按照 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)
复制代码
您可以使用 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://juejin.im/post/5b87c6e3f265da432144bdfe