小编Nem*_*ior的帖子

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

我考虑过使用OpenAI 的梯度检查点，希望这会有所帮助。然而，简单地按照 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 模式下使用梯度检查点？