python之使用 Tensorflow Eager Execution 的 OpenAI 梯度检查点_编程设计

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

我考虑过使用 Gradient Checkpointing来自 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 模式下使用梯度检查点？

请您参考如下方法:

openai的gradient checkpointing代码是基于graph重写的，所以不支持eager execution。

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