[Paper Review] Memory-Efficient Implementation of DenseNets

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

Advanced Neural Network Applications12 references86 citations

TL;DR

The paper introduces a memory-sharing strategy with recomputation to train extremely deep DenseNets by reducing feature-map memory from quadratic to linear, enabling models up to hundreds of layers with modest time overhead.

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%.

Motivation & Objective

• Motivate DenseNet's parameter efficiency and feature reuse for high-capacity models.
• Identify the quadratic memory bottleneck in standard DenseNet training.
• Propose a memory-sharing strategy to reduce training memory from quadratic to linear.
• Demonstrate training of very deep DenseNets and competitive ImageNet performance under memory budget.

Proposed method

• Identify the two sources of quadratic memory in DenseNets: pre-activation batch normalization and contiguous concatenation.
• Introduce Shared Memory Storage 1 for concatenation outputs and Shared Memory Storage 2 for batch normalization outputs.
• Recompute concatenation and batch normalization during back-propagation to populate shared storages instead of storing all intermediates.
• Share gradient storage across layers to avoid quadratic growth in gradients.
• Measure memory and time overhead, showing ~15-20% extra training time with significant memory savings.

Experimental results

Research questions

• RQ1 Can DenseNets be trained effectively under reduced memory by reusing shared storage and recomputation?
• RQ2 How much memory can be saved (quadratic to linear) and at what computational cost?
• RQ3 What are the practical limits of depth and parameter count achievable on ImageNet with memory-efficient DenseNets?

Key findings

• Memory consumption becomes linear in depth with the proposed shared memory strategy.
• In LuaTorch, a 160-layer model uses about 22% of the memory of the naïve implementation, enabling training of ~340-layer models within a 12 GB budget.
• In PyTorch, training very deep DenseNets is feasible with nearly 500 layers on a single GPU.
• DenseNets trained with the efficient implementation achieve top-1 error of 20.26% on ImageNet with 264 layers (k=48, 73M parameters).
• The deepest cosine DenseNet achieves a top-1 error of 20.26%, outperforming previous state-of-the-art.
• Sharing gradient storage is beneficial with no time cost; adding shared BN/concat storage adds ~15-20% time overhead.