OpenAI Gradient Checkpointing with Tensorflow Eager Execution答案

【问题标题】：OpenAI Gradient Checkpointing with Tensorflow Eager ExecutionOpenAI Gradient Checkpointing with Tensorflow Eager Execution
【发布时间】：2018-06-08 21:41:55
【问题描述】：

我最近切换到了 Tensorflow Eager（目前正在使用 TF 1.8.0）并且非常喜欢它。但是，我现在有一个相当大的模型，当使用计算 TF 中的梯度所需的梯度磁带运行时，它不适合我的 GPU 内存（GTX 1080Ti，12GB VRAM）。正向传递（即不使用渐变胶带）工作正常。

我考虑过使用来自 OpenAI 的 Gradient Checkpointing，希望这会有所帮助。然而，简单地按照他们的 Git 中描述的方式使用它似乎对 Eager Execution 没有帮助，即

import tensorflow as tf
import tensorflow.contrib.eager as tfe
import memory_saving_gradients
tf.__dict__["gradients"] = memory_saving_gradients.gradients_memory
# using gradients_memory or gradients_speed does not change anything
# tf.__dict__["gradients"] = memory_saving_gradients.gradients_speed

[...]
with tfe.GradientTape() as g:
    output = run_large_model()
    loss = calculate_loss_on_output(output)
grads = g.gradient(full, model.variables)
optimizer.apply_gradients(zip(grads, model.variables))

内存不足，与是否使用梯度检查点无关。

我的猜测是梯度磁带仍然存储所有变量和向后传递所需的信息，梯度检查点没有效果，因为 Eager 模式下的 TF 实际上并没有构建图形（据我了解 - 或至少这是一个不同的图表）。

您是否有任何经验或想法如何解决这个问题，或者我需要做什么才能在 TF Eager 模式下使用梯度检查点？

【问题讨论】：

标签： python tensorflow gradient-descent tensorflow-gradient

【解决方案1】：

openai 的梯度检查点代码基于图形重写，因此不支持 Eager Execution。

tensorflow.contrib.layers 库有一个 recompute_grad 装饰器，它是等效的，但在图形和急切执行中都受支持。

【讨论】：

感谢您的回复。有没有办法轻松地将这个 recompute_grad 装饰器与 keras 模型结合起来？假设我有一个 Keras 模型 class BigNeuralNet(tf.keras.Model)，它在其 __init__() 方法中初始化了许多层（例如 self.layer1 = tf.layers.Conv2D(...)，然后在 call() 方法中使用这些层。模型的调用方法将张量作为输入，还有一些其他值（例如，用于控制 Batch Normalization 的训练或推理），它们在前向传递期间使用。仅在模型的调用方法上使用 recompute_graddecorator 就足够了吗？
我认为应该是。如果不是，请打开一个 github 问题并抄送我，我们将在那里讨论。
@AlexandrePassos 你是如何让recompute_grad 工作的？我正在使用 Keras，它对我不起作用。自从 Joey Yearsley 将它用于 DenseNet 后，我就联系了他，但他还没有让它与 keras 一起使用
你可以试试 2.0 nightly build 中的 tf.recompute_grad 吗？我们有测试表明它可以与 keras 一起使用。如果它不适合你，你可以用一个简短的复制示例提交一个 github 问题吗？