显存不够时，如何利用GPU训练数据

参考：https://github.com/openai/gradient-checkpointing

           https://www.jqr.com/article/000537

           https://www.infoq.cn/article/fRN0eVBNedTIrm5_Ps1M

使用gradient-checkpointing来节省存储

       训练深度神经网路的内存密集部分是通过反向传播计算损失梯度，通过在模型定义的计算图中使用检查节点，并在反向传播期间重新计算这些节点之间图的部分，可以在减低内存的情况下降低梯度。

        在训练N层的深度前馈神经网络时，我们可以将内存消耗降低到，所需要代价仅仅是增加一个额外的向前传播。在TensorFlow中实现该资源库提供的功能，需要使用TensorFlow图编辑去自动的重写返现传递的计算图。

工作原理

        对于一个N层的简单feed-forward neural network，用以获取梯度的计算图如下：

      神经网络层和的激活对应标有f的节点，在向前传递期间forward pass，按顺序评估这些节点。相对激活节点的损失梯度和这些层的参数标记为b节点。在向后传递期间，这些标记为b的节点以相反的顺序进行评估。b节点的计算需要f节点的结果，因此正想传递之后f节点也被保存在内存中。只有当反向传播的足够远，使得f节点的所有依赖关系、子节点、计算完成后，f节点才能从内存中擦除。这意味着简单的反向传播所需要的内存随着神经网络层数n呈现线性增长的趋势。接下来我们展示了节点计算的顺序，紫色阴影的圆代表了在任何给定的时刻需要讲那些节点保存在存储器中。

  

      如上述的简单反向传播在计算方面是最佳的：每个节点他只计算一次。然而，如果我们愿意重新计算这些节点，我们就可以节省大量的内存。我们可能只是在每次需要的它的时候从forward重新计算每个节点。执行顺序和内存使用如下图所示：

      使用这种策略，在我们的图中计算梯度所需要的内存在神经网络层数n是恒定的，这在内存方面是最佳的。但是请注意的是，节点评估的数量现在是按照n^2进行放大的，而之前的尺度是n：即n个节点中的每个节点都按照n次的顺序重新计算。因此计算图对于深度学习变得慢得多，这也造成了该方法在深度学习中变得不切实际。

这些检查节点在正向传递forward pass之后保存在内存中，而其余节点最多重新计算一次。重新计算后，这些非检查节点被保存在内存中直到不再需要它们为止。对于简单的前馈神经网络的情况，所有神经元激活节点是由向前通道forward pass定义的图形分离器或关节点。这意味着我们只需要在反向传播期间backprop重新计算b节点和b节点之前的最后一个检查节点间的其他节点。当反向传播backprop进展到足以到达检查节点的时候，在检查节点间被重新计算的节点，就可以从内存中擦除。计算和内存的使用可以参照下图：

对于我们示例中的简单前馈网络，最佳的选择是将每隔sqrt(n)的节点标记为检查点。这样，检查点个数和检查点之间节点的个数都是sqrt(n)的量级，这也意味着所需要的内存现在也与我们的网络中层数的平方更成比例。由于每个节点最多重新计算一次，因此该策略所需要的额外计算相当于通过网络的单个正向传递。

我们的软件包实现了检查点反向传播在上面的图中展现出来了，这是通过获取标准的反向传播图形并使用TensorFlow图形编辑器自动重写实现的。对于包含关节点（单节点图分频器）我们使用sqrt(n)策略自动选择检查点，为前馈神经网络提供sqrt(n)内存的使用。对于仅包含多节点图分隔的更一般图，我们的反向检查点试试还是有效的。但是我们当前需要用于手动选择检查点。

安装：确保CUDA环境变量配置好了

pip install tf-nightly-gpu
pip install toposort networkx pytest

使用：

from memory_saving_gradients import gradients

替换TensorFlow中的tf.gradients（前提是tf.gradients是直接使用的，而不是在tf.train.Optimizer中隐含的使用的）。

除了tf.gradients的常规参数外，我们的gradients函数还有一个额外的参数，即checkpoints。这些checkpoints参数告诉梯度函数在前向传播时那些节点是你希望设置为检查节点的。这些在checkpoints之间的节点在反向传播期间被重新计算。您可以为gradients(ys,xs,checkpoints=[tensor1, tensor2])提供张量列表，或者您也可以使用几个关键字之一：

• 'collection'(defualt)：通过ty.get_collection('checkpoint')返回checkpoints的所有张量，然后在模型定义的时候，您需要使用tf.add_to_collection('checkpoints',tensor)确保将tensors张量增加到了这个集合。
• 'memory'：这里使用启发式的方法自动的选择一组节点成为checkpoint，我们希望内存用量。该启发式的工作通过在图中自动识别关键节点articulation points，即当该articulation points关键节点被移除时，将图拆分成连个断开的部分张量，然后检查这些tensors的合适数量。这目前试用许多的模型，但不是全部模型。
• 'speed'：此选项通过检查通常计算成本高昂的所有操作的输出（即卷积核矩阵乘法）来尝试最大化运行速度。

覆盖tf.gradients

如果想要试用gradients方法直接替代生效，只需要覆盖就可以了，但前提是python中已经试用tf.gradinets名称的函数，替代方法如下

import tensorflow as tf
import memory_saving_gradients
# monkey patch tf.gradients to point to our custom version, with automatic checkpoint selection
tf.__dict__["gradients"] = memory_saving_gradients.gradients_memory

此后，对tf.gradients的所有调用都将试用memory saving version去替代。

同样的Keras时也可以这样去做：

from tensorflow.python.keras._impl.keras import backend as K
K.__dict__["gradients"] = memory_saving_gradients.gradients_memory

测试

测试文件夹包含了用于测试代码正确性和分析各模型内存使用情况的脚本，修改代码后，您可以从该文件夹中运行./run_all_tests.sh来执行测试方法。

限制

提供代码在运行模型之前在python中执行所有图形操作，对于大型图而言是缓慢的。当前用于自动选择检查点的算法纯粹是启发式的，因而在我们测试之外的某些模型上预计可能会失败。在这种情况下应当手动的选择检查点：将您选择的检查点添加到TensorFlow中明教checkpoints的集合，并在调用我们gradinets function时使用chechpoints=collection。