如何使用Tensorflow在单个GPU上训练具有相对较大批量的大型模型？

Question

我有一个非常大的模型，由于内存不足错误，无法在批量大小为64的单个GPU上进行训练。有人建议我使用较小的批量。但是，如果我减小批量，则精度会下降。其中一个解决方案是只喂一半当前批次，存储梯度然后喂入剩余的。这可以使用compute_gradients和apply_gradients明确完成。但它相对不方便（如果存在简洁的实现则可以）。所以我想知道这个问题是否有更好的解决方案（或简洁的实现）。

提前致谢。

Answer 1

您可以考虑调查一下：https://github.com/openai/gradient-checkpointing。

最近有很多关于以更多的前向传递为代价使反向提高内存效率的研究。这是TensorFlow的一个这样的方案的最近实现。