深度网络训练中反向传播的GPU内存使用优化

在现代深度学习中，设计更大规模的深度神经网络（DNN）以执行更复杂的任务并获得更高精度已成为趋势。而卷积神经网络（CNN）则成为计算机视觉任务的标准方法。然而，卷积层中中间数据的内存分配会在模型训练期间造成严重的内存压力。为解决此问题，已提出许多解决方案。除硬件依赖方案外， rematerialization 作为一种通用方法，可通过高效地以计算换取内存来降低GPU内存使用。其核心思想是在前向阶段选择一组中间结果作为检查点，仅将其保留在内存中以降低内存占用。在后向阶段，按需从最近的内存检查点重新计算中间数据。这种重计算会增加执行时间，但通过不在前向阶段存储所有中间结果来节省内存。本文聚焦于高效寻找最优检查点子集，以实现模型训练期间峰值内存使用的最小化。我们首先描述使用数学方程训练神经网络的理论背景。通过这些方程识别在前向和后向阶段计算模型权重梯度所需的全部关键数据。我们首先确定检查点选择问题，并提出一种时间复杂度为 O(n^3) 的动态规划算法来求解寻找最优检查点子集的问题。通过大量实验，我们基于理论分析制定更准确的模型描述，并在追踪基础上修订目标函数，提出一种时间复杂度为 O(n) 的最优检查点子集寻找算法。