A PyTorch >=1.0 implementation of DenseNets, optimized to save GPU memory.
- Now works on PyTorch 1.0! It uses the checkpointing feature, which makes this code WAY more efficient!!!
While DenseNets are fairly easy to implement in deep learning frameworks, most implmementations (such as the original) tend to be memory-hungry. In particular, the number of intermediate feature maps generated by batch normalization and concatenation operations grows quadratically with network depth. It is worth emphasizing that this is not a property inherent to DenseNets, but rather to the implementation.
This implementation uses a new strategy to reduce the memory consumption of DenseNets. We use checkpointing to compute the Batch Norm and concatenation feature maps. These intermediate feature maps are discarded during the forward pass and recomputed for the backward pass. This adds 15-20% of time overhead for training, but reduces feature map consumption from quadratic to linear.
This implementation is inspired by this technical report, which outlines a strategy for efficient DenseNets via memory sharing.
- PyTorch >=1.0.0
- CUDA
A comparison of the two implementations (each is a DenseNet-BC with 100 layers, batch size 64, tested on a NVIDIA Pascal Titan-X):
| Implementation | Memory cosumption (GB/GPU) | Speed (sec/mini batch) |
|---|---|---|
| Naive | 2.863 | 0.165 |
| Efficient | 1.605 | 0.207 |
| Efficient (multi-GPU) | 0.985 | - |