Source-linked AI summary
Memory-Efficient Implementation of DenseNets
Geoff Pleiss, Danlu Chen, Gao Huang, Tongcheng Li, Laurens van der Maaten, Kilian Q. Weinberger
TL;DR
Naïve DenseNet training can require quadratically growing feature-map memory, limiting model depth despite the architecture’s efficient feature reuse. The report uses shared memory allocations and recomputation to reduce this cost, enabling a 264-layer DenseNet with a 20.26% ImageNet single-crop top-1 error.
Problem
Quadratic feature-map memory in naïve DenseNet implementations can make training large models expensive or infeasible, limiting model size despite DenseNet’s feature-reuse efficiency.
Method
Shared memory allocations store intermediate concatenation and normalization outputs across layers, while cheap intermediate results are recomputed during back-propagation.
Results
20.26% single-crop top-1 error: a 264-layer DenseNet with 73M parameters achieves this result on ImageNet.
Takeaways & Limitations
Memory-efficient implementation removes feature-map memory as the main obstacle to training extremely deep DenseNets and supports a measurable ImageNet top-1 error reduction.
Takeaways & Limitations
The most efficient implementation’s total memory still does not grow linearly with depth because DenseNet parameter count is inherently quadratic.
Abstract
from arXiv · showhide
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%.
1 Introduction
DenseNets reuse features efficiently, but naïve implementations can make feature-map memory grow quadratically with depth. Shared memory allocations address this bottleneck, enabling substantially larger models and a 264-layer ImageNet result.
- Motivation: 20M vs 44M parameters and 80B/image vs 155B/image: a 201-layer DenseNet matches roughly the top-1 classification error of a 101-layer ResNet.DenseNet efficiency comes from connecting each layer to all previous layers within a pooling region, promoting feature reuse.
- Motivation: 22.2% top-1 single-crop error: a 161-layer DenseNet with k = 48 features per layer and 29M parameters on ImageNet.The report argues that larger networks could perform better, but existing implementations are limited by GPU memory.
- Memory bottleneck: O(m^2) feature-map memory: naïve DenseNet implementations store intermediate maps from batch normalization and concatenation during training.This quadratic growth is an implementation issue rather than an inherent property of the architecture.
- Contribution: 15−20% additional training time: shared memory allocations reduce feature-map memory consumption from quadratic to linear by overwriting and recomputing intermediate results.The strategy enables a 264-layer DenseNet with k = 48 and 73M parameters.
- Contribution: 20.26% single-crop top-1 error: the 264-layer, 73M-parameter DenseNet achieves this result on ImageNet.The report describes this as state-of-the-art to the authors’ knowledge.
2 The DenseNet Architecture
DenseNet layers access all preceding features within a dense block, but pre-activation normalization and contiguous concatenation create substantial intermediate-memory costs. These costs arise from implementation choices around repeated feature copies and storage.
- Dense connectivity: All preceding features: each DenseNet layer receives the outputs of earlier layers within its dense block, rather than only the most recent features.Dense blocks are typically followed by pooling or the final classifier.
- Layer operations: Batch normalization, ReLU, and convolution: DenseNet layers apply these operations in that order, with each commonly producing intermediate feature maps.The layer transformation is represented as Hℓ, while [·] denotes concatenation.
- Memory growth: O(m^2) memory: storing intermediate features for a dense block with m layers can create quadratic usage, even though convolutional features themselves grow linearly.Frameworks often retain these maps for back-propagation.
- Memory growth: Up to m normalized copies: pre-activation batch normalization duplicates each layer’s features with distinct scaling and bias, allocating up to (m −1)(m −2)/2 maps.These copies arise because each layer applies its own normalization parameters to previous features.
- Memory growth: 30−50% computation-time overhead: non-contiguous memory is slower for convolution, while forcing contiguous inputs requires copying previous features and can create quadratic memory cost.Contiguous operations are preferred but generate multiple copies of each feature.
- Implementation constraint: Feature tensors store minibatch as the outer dimension, so adjacent allocations concatenate along the minibatch dimension rather than the intended feature-map dimension.Consequently, filter outputs cannot simply be assigned to a pre-allocated contiguous feature block.
3 Naïve Implementation
The naïve DenseNet layer builds contiguous inputs and intermediate tensors through separate allocations, then retains additional forward and backward-pass data. This repeated storage explains its rapid memory growth.
- Forward pass: New memory allocations: the original implementation separately allocates tensors for concatenation, batch normalization, and subsequent operations.Figure 3 contrasts these allocations with temporary shared buffers in the efficient implementation.
- Forward pass: Each naïve layer first copies previous-layer features into a contiguous block before batch normalization and convolution.If each previous layer produces k features, the block accommodates ℓ × k feature maps.
- Backward pass: ℓ × k feature maps: back-propagation can additionally allocate storage for output gradients, normalized forward maps, and concatenated gradients.These allocations add to the memory required by the forward computation graph.
4 Memory-Efficient Implementation
The memory-efficient implementation replaces per-layer intermediate allocations with shared storage and recomputes cheap concatenation and normalization outputs during back-propagation. This trades modest computation for much lower memory use.
- Core strategy: Two pre-allocated Shared Memory Storage locations: intermediate concatenation and normalization outputs are written into shared blocks during the forward pass.During back-propagation, the concatenated and normalized features are recomputed as needed.
- Shared storage: O(ℓk) memory: shared storage holds concatenation outputs across layers instead of allocating a separate block for every concatenation.Because the storage is overwritten, required feature maps are recomputed during back-propagation.
- Shared storage: O(m) memory: batch-normalization outputs are assigned to a second shared allocation, which the convolution reads through pointers.The data is overwritten by later layers and therefore recomputed during the backward pass.
- Gradient storage: Single shared memory allocation: gradient tensors from concatenation, batch normalization, and convolution can share storage during back-propagation.PyTorch and MxNet provide gradient sharing out of the box, while the LuaTorch implementation uses an explicit sharing scheme.
- Backward pass: One additional backward-pass step: the implementation recomputes concatenation and batch normalization to repopulate shared storage before gradients are calculated.The forward pass remains similar to the naïve implementation except for where intermediate maps are stored.
5 Results
Memory-sharing implementations substantially reduce DenseNet training memory, enabling much deeper models with modest computation-time overhead and improved ImageNet error.
- 12 GB supports a 340-layer model, 2.5× as deep and with 6× as many parameters as the best naïve implementation model.The 160-layer efficient model uses 22% of the naïve implementation’s memory.
- The most efficient implementation’s total memory still grows nonlinearly because DenseNet parameter count is quadratic in depth.Parameter storage remains much smaller than feature-map storage and does not impede model depth.
- Nearly 500 layers and 13M parameters can be trained on a single GPU with the efficient PyTorch implementation.PyTorch is more memory efficient than LuaTorch, likely partly because autograd performs training-time memory optimizations.
- 15% LuaTorch and 20% PyTorch training-time overhead result from sharing batch-normalization and concatenation storage.Shared gradient storage adds no time cost; the additional overhead comes from recomputation during back-propagation.
- A 264-layer DenseNet with 73M parameters can be trained using the efficient LuaTorch implementation on 8 NVIDIA Tesla M40 GPUs.The prior deepest model trained with the original LuaTorch implementation had 161 layers, 29M parameters, and k = 48 features per layer.
- 20.26% top-1 error is achieved by the deepest cosine DenseNet, outperforming the previous state-of-the-art model.The new standard-training DenseNets also achieve nearly one percentage point better error than the deepest ResNet while using fewer parameters.
6 Conclusion
The report replaces naïve intermediate-feature storage with shared memory and recomputation, substantially reducing memory use with only a small computation-time increase. This removes memory as a practical barrier to extremely deep DenseNets and accompanies measurable ImageNet error reductions.
- Shared memory buffers and recomputation replace storing all intermediate feature maps during DenseNet training.The strategy uses shared storage and recomputes cheap transformations during back-propagation.
- Significantly lower memory use comes with only a small increase in computation time.
- The implementation enables models at roughly twice the depth of prior models and produces a measurable drop in ImageNet top-1 error.