Question

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

我考虑过使用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模式下使用渐变检查点需要做些什么？

Answer 1

openai的渐变检查点代码基于图形重写，因此不支持急切执行。

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

具有Tensorflow Eager执行的OpenAI梯度检查点

1 个答案: