内容简介:文 / 软件工程实习生 Xuechen Li来源:TensorFlow 公众号Eager Execution 可简化 TensorFlow 中的模型构建体验,而 Graph Execution 可提供优化,以加快模型运行速度及提高存储效率。本篇博文展示了如何编写 TensorFlow 代码,以便将借助 tf.keras API 并使用 Eager Execution 构建的模型转换为图表,最终借助 tf.estimator API 的支持,在 Cloud TPU 上部署此模型。
文 / 软件工程实习生 Xuechen Li
来源:TensorFlow 公众号
Eager Execution 可简化 TensorFlow 中的模型构建体验,而 Graph Execution 可提供优化,以加快模型运行速度及提高存储效率。本篇博文展示了如何编写 TensorFlow 代码,以便将借助 tf.keras API 并使用 Eager Execution 构建的模型转换为图表,最终借助 tf.estimator API 的支持,在 Cloud TPU 上部署此模型。
注:tf.keras 链接 www.tensorflow.org/guide/keras
tf.estimator 链接 www.tensorflow.org/guide/estim…
我们使用可逆残差网络(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 链接 www.tensorflow.org/api_docs/py…
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:
- 正演计算
- 梯度的反演计算
- 将梯度应用于变量
例如,我们可以按以下方式 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 正处于积极开发中,将其加以应用是一项不断发展的技术。如需更多信息,请查看其文档字符串。 注:其文档字符串链接 github.com/tensorflow/…
使用 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 链接 www.tensorflow.org/guide/autog…
使用 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 链接 github.com/tensorflow/… 配置链接 www.tensorflow.org/api_docs/py… tf.contrib.tpu.TPUEstimator 链接 www.tensorflow.org/api_docs/py… tf.contrib.tpu.CrossShardOptimizer 链接 www.tensorflow.org/api_docs/py…
如需了解具体说明,请查看 RevNet 示例文件夹中的 TPU 估算器脚本。我们希望日后可以使用 tf.contrib.tpu.keras_to_tpu_model 进一步简化使 Keras 模型在 TPU 上运行的流程。 注:TPU 估算器脚本链接 github.com/tensorflow/… tf.contrib.tpu.keras_to_tpu_model 链接 github.com/tensorflow/…
可选:模型性能
与一般反向传播相比,有了 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 的其他示例。 注:此处连接诶 github.com/tensorflow/… 其他示例链接 github.com/tensorflow/…
以上所述就是小编给大家介绍的《使用 Eager Execution 编码并运行图表:以通过 RevNet 优化代码为例》,希望对大家有所帮助,如果大家有任何疑问请给我留言,小编会及时回复大家的。在此也非常感谢大家对 码农网 的支持!
猜你喜欢:本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们。