[论文解读] Memory-Efficient Implementation of DenseNets

Geoff Pleiss, Danlu Chen+4 more|arXiv (Cornell University)|Jul 21, 2017

Advanced Neural Network Applications参考文献 12被引用 86

一句话总结

本文提出一种内存共享策略，结合重新计算，通过将特征图内存从二次方增长降至线性，实现可达数百层的DenseNets，时间开销适中。

ABSTRACT

The DenseNet architecture is highly computationally efficient as a result of feature reuse. However, a naive DenseNet implementation can require a significant amount of GPU memory: If not properly managed, pre-activation batch normalization and contiguous convolution operations can produce feature maps that grow quadratically with network depth. In this technical report, we introduce strategies to reduce the memory consumption of DenseNets during training. By strategically using shared memory allocations, we reduce the memory cost for storing feature maps from quadratic to linear. Without the GPU memory bottleneck, it is now possible to train extremely deep DenseNets. Networks with 14M parameters can be trained on a single GPU, up from 4M. A 264-layer DenseNet (73M parameters), which previously would have been infeasible to train, can now be trained on a single workstation with 8 NVIDIA Tesla M40 GPUs. On the ImageNet ILSVRC classification dataset, this large DenseNet obtains a state-of-the-art single-crop top-1 error of 20.26%.

研究动机与目标

• 动机：提升 DenseNet 在高容量模型中的参数效率与特征重用。
• 识别标准DenseNet训练中的二次内存瓶颈。
• 提出一种内存共享策略，将训练内存从二次方降低到线性。
• 在受限内存条件下演示极深 DenseNets 的训练以及在 ImageNet 上具有竞争力的性能。

提出的方法

• 识别 DenseNets 中导致二次内存的两个来源：预激活批量归一化和连续拼接。
• 引入共享内存存储1用于拼接输出，共享内存存储2用于批量归一化输出。
• 在反向传播期间重新计算拼接和批量归一化，以填充共享存储，而不是存储所有中间结果。
• 跨层共享梯度存储以避免梯度的二次增长。
• 衡量内存和时间开销，显示在显著内存节省的同时，训练时间额外约15-20%。

实验结果

研究问题

• RQ1 通过重复使用共享存储和重新计算，在降低的内存条件下能否有效训练DenseNets？
• RQ2 可以节省多少内存（从二次到线性），以及在何种计算开销下？
• RQ3 在内存高效的DenseNets上，ImageNet上可达到的实际深度和参数量的上限是什么？

主要发现

• 提出的共享内存策略使内存消耗随深度线性增长。
• 在 LuaTorch 中，160 层模型大约使用天真实现内存的 22%，在 12 GB 预算下能够训练约 340 层模型。
• 在 PyTorch 中，在单个 GPU 上训练近500层的极深DenseNets是可行的。
• 使用高效实现训练的DenseNets在ImageNet上达到264层（k=48，73M参数）的Top-1错误率为20.26%。
• 最深的cosine DenseNet达到了20.26%的Top-1错误率，超过了之前的最先进水平。
• 共享梯度存储有益，且不增加时间成本；增加共享的 BN/concat 存储会带来约15-20% 的时间开销。