A standard hardware bottleneck when training deep neural networks is GPU memory. The bulk of memory is occupied by caching intermediate tensors for gradient computation in the backward pass. We propose a novel method to reduce this footprint - Dropping Intermediate Tensors (DropIT). DropIT drops min-k elements of the intermediate tensors and approximates gradients from the sparsified tensors in the backward pass. Theoretically, DropIT reduces noise on estimated gradients and therefore has a higher rate of convergence than vanilla-SGD. Experiments show that we can drop up to 90\% of the intermediate tensor elements in fully-connected and convolutional layers while achieving higher testing accuracy for Visual Transformers and Convolutional Neural Networks on various tasks (e.g., classification, object detection, instance segmentation). Our code and models are available at https://github.com/chenjoya/dropit.
翻译:深度神经网络训练中的一个标准硬件瓶颈是GPU内存。大部分内存被用于缓存反向传播中梯度计算所需的中间张量。我们提出了一种减少这种内存占用的新方法——丢弃中间张量(DropIT)。DropIT会丢弃中间张量的最小k个元素,并在反向传播中从稀疏化的张量近似计算梯度。理论上,DropIT能降低估计梯度中的噪声,因此比普通随机梯度下降法具有更高的收敛速率。实验表明,在全连接层和卷积层中,我们可以丢弃高达90%的中间张量元素,同时在各种任务(如分类、目标检测、实例分割)中,视觉Transformer和卷积神经网络都能达到更高的测试精度。我们的代码和模型可在https://github.com/chenjoya/dropit获取。