DropIT:丢弃中间张量以实现内存高效的深度神经网络训练
计算机视觉与模式识别
2023-03-03 v3
摘要
训练深度神经网络时的一个标准硬件瓶颈是GPU内存。内存的大部分被缓存用于反向传播中梯度计算的中间张量所占据。我们提出了一种减少该占用的新方法——丢弃中间张量(DropIT)。DropIT丢弃中间张量的min-k元素,并从稀疏化后的张量近似计算反向传播中的梯度。理论上,DropIT降低了估计梯度上的噪声,因此比vanilla-SGD具有更高的收敛速率。实验表明,我们可以在全连接层和卷积层中丢弃多达90%的中间张量元素,同时在各类任务(如分类、目标检测、实例分割)上针对Visual Transformer和卷积神经网络取得更高的测试精度。我们的代码和模型可在 https://github.com/chenjoya/dropit 获取。
引用
@article{arxiv.2202.13808,
title = {DropIT: Dropping Intermediate Tensors for Memory-Efficient DNN Training},
author = {Joya Chen and Kai Xu and Yuhui Wang and Yifei Cheng and Angela Yao},
journal= {arXiv preprint arXiv:2202.13808},
year = {2023}
}
备注
ICLR 2023