Repository files navigation

A PyTorch >=1.0 implementation of DenseNets, optimized to save GPU memory.

1. 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, mostimplmementations (such as theoriginal) tend to be memory-hungry.In particular, the number of intermediate feature maps generated by batch normalization and concatenation operationsgrows 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 usecheckpointing 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, butreduces feature map consumption from quadratic to linear.

This implementation is inspired by thistechnical report, which outlines a strategy for efficient DenseNets via memory sharing.

• PyTorch >=1.0.0
• CUDA

In your existing project:There is one file in themodels folder.

• models/densenet.py is an implementation based off thetorchvision andproject killer implementations.

If you care about speed, and memory is not an option, pass theefficient=False argument into theDenseNet constructor.Otherwise, pass inefficient=True.

Options:

• All options are described inthe docstrings of the model files
• The depth is controlled byblock_config option
• efficient=True uses the memory-efficient version
• If you want to use the model for ImageNet, setsmall_inputs=False. For CIFAR or SVHN, setsmall_inputs=True.

Running the demo:

The only extra package you need to install ispython-fire:

pip install fire

• Single GPU:

CUDA_VISIBLE_DEVICES=0 python demo.py --efficient True --data<path_to_folder_with_cifar10> --save<path_to_save_dir>

• Multiple GPU:

CUDA_VISIBLE_DEVICES=0,1,2 python demo.py --efficient True --data<path_to_folder_with_cifar10> --save<path_to_save_dir>

Options:

• --depth (int) - depth of the network (number of convolution layers) (default 40)
• --growth_rate (int) - number of features added per DenseNet layer (default 12)
• --n_epochs (int) - number of epochs for training (default 300)
• --batch_size (int) - size of minibatch (default 256)
• --seed (int) - manually set the random seed (default None)

A comparison of the two implementations (each is a DenseNet-BC with 100 layers, batch size 64, tested on a NVIDIA Pascal Titan-X):

• LuaTorch (by Gao Huang)
• Tensorflow (by Joe Yearsley)
• Caffe (by Tongcheng Li)

@article{pleiss2017memory,  title={Memory-Efficient Implementation of DenseNets},  author={Pleiss, Geoff and Chen, Danlu and Huang, Gao and Li, Tongcheng and van der Maaten, Laurens and Weinberger, Kilian Q},  journal={arXiv preprint arXiv:1707.06990},  year={2017}}