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

栏目: 编程工具 · 发布时间: 6年前

内容简介:文 / 软件工程实习生 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:

  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 正处于积极开发中,将其加以应用是一项不断发展的技术。如需更多信息,请查看其文档字符串。 注:其文档字符串链接 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 优化代码为例》,希望对大家有所帮助,如果大家有任何疑问请给我留言,小编会及时回复大家的。在此也非常感谢大家对 码农网 的支持!

查看所有标签

猜你喜欢:

本站部分资源来源于网络,本站转载出于传递更多信息之目的,版权归原作者或者来源机构所有,如转载稿涉及版权问题,请联系我们

网飞传奇

网飞传奇

[美]吉娜·基廷 / 谭永乐 / 中信出版社 / 2014-1-1 / 42

飞的历史充满了传奇色彩,它的崛起伴随着复杂斗争、幸运转折、个人背叛……它自身的历史比它出租的那些电影还要更富有戏剧性。网飞在1997年建立,而建立的原因仅仅是因为创始人伦道夫和哈斯廷斯没有按时归还租借的DVD,还要缴纳因此而产生的滞纳金。 1999年,网飞公司摒弃了百视达的“每片付租”模式,转而采用了一种订阅模式:用户只需要支付固定费用,就能尽情租片观赏,免去了到期还片日、滞纳金、运费和手续......一起来看看 《网飞传奇》 这本书的介绍吧!

随机密码生成器
随机密码生成器

多种字符组合密码

RGB CMYK 转换工具
RGB CMYK 转换工具

RGB CMYK 互转工具

HSV CMYK 转换工具
HSV CMYK 转换工具

HSV CMYK互换工具