低内存神经网络训练：技术报告
机器学习
2022-04-12 v2 机器学习
摘要
在训练神经网络模型时，内存日益频繁地成为瓶颈。尽管如此，与大量关于降低推理内存需求的文献相比，降低训练整体内存需求的技术研究较少。本文研究一个基本问题：训练一个神经网络究竟需要多少内存？为回答该问题，我们在两个代表性深度学习基准上分析了训练的整体内存使用——用于图像分类的 WideResNet 模型和用于机器翻译的 DynamicConv Transformer 模型——并全面评估了四种降低训练内存需求的标准技术：(1) 对模型施加稀疏性，(2) 使用低精度，(3) 微批处理（microbatching），以及 (4) 梯度检查点（gradient checkpointing）。我们探究了每种技术单独使用对训练峰值内存使用和最终模型质量的影响，并探讨了组合这些技术时所带来的内存、准确率和计算权衡。通过适当组合这些技术，我们表明：训练 CIFAR-10 上的 WideResNet-28-2 所需内存最多可降低 60.7 倍，准确率损失 0.4%；训练 IWSLT'14 德英翻译上的 DynamicConv 模型所需内存最多可降低 8.7 倍，BLEU 分数下降 0.15。
引用
@article{arxiv.1904.10631,
title = {Low-Memory Neural Network Training: A Technical Report},
author = {Nimit S. Sohoni and Christopher R. Aberger and Megan Leszczynski and Jian Zhang and Christopher Ré},
journal= {arXiv preprint arXiv:1904.10631},
year = {2022}
}
备注
Version notes: Copyedits and citation fixes