降低 PyTorch 选择性微分的内存消耗
机器学习
2024-08-22 v2
摘要
内存是许多深度学习任务的限制性资源。除了神经网络权重之外，主要的内存消耗之一是自动微分（AD）为反向传播构建的计算图。我们观察到，PyTorch 当前的 AD 实现在存储计算图时忽略了关于参数可微性的信息。然而，当只需要参数子集的梯度时（如在许多现代微调任务中），这些信息对于减少内存消耗是有用的。具体而言，每当参数被标记为不可微时，在其参数上呈线性作用的层（密集层、卷积层或归一化层）的输入就可以被丢弃。我们提供了此类层的直接替换、与可微性无关的实现，并证明了其在不影响运行时间的情况下减少内存消耗的能力。
引用
@article{arxiv.2404.12406,
title = {Lowering PyTorch's Memory Consumption for Selective Differentiation},
author = {Samarth Bhatia and Felix Dangel},
journal= {arXiv preprint arXiv:2404.12406},
year = {2024}
}
备注
The code is available at https://github.com/plutonium-239/memsave_torch . This paper was accepted to WANT@ICML'24