First clone this repo, then install all dependencies
pip install -r requirements.txt
The code was tested with Python 3.6.0.
Below is a description of the major sections of the code base. Run python main.py --help for a complete description of flags and hyperparameters.
This code base supports the following datasets: MNIST, CIFAR-10, CIFAR-100, Tiny ImageNet, ImageNet. All datasets except ImageNet will download automatically. For ImageNet setup locally in the Data folder.
There are four model classes each defining a variety of model architectures:
- Default models support basic dense and convolutional model.
- Lottery ticket models support VGG/ResNet architectures based on OpenLTH.
- Tiny ImageNet models support VGG/ResNet architectures based on this Github repository.
- ImageNet models supports VGG/ResNet architectures from torchvision.
Custom dense, convolutional, batchnorm, and residual layers implementing masked parameters can be found in the Layers folder.
All pruning algorithms are implemented in the Pruners folder.
Below is a list and description of the experiment files found in the Experiment folder:
singleshot.py: used to make figure 1, 2, and 6.multishot.py: used to make figure 5a.unit-conservation.py: used to make figure 3.layer-conservation.py: used to make figure 4.lottery-layer-conservation.py: used to make figure 5b.synaptic-flow-ratio.py: used to make figure 7.
All data used to generate the figures in our paper can be found in the Results/data folder. Run the notebook figures.ipynb to generate the figures.
Due to an error in multishop.py (which has since been fixed), IMP did not reset the parameters to their original values between iterations. All benchmarks in the paper are not affected as they are run in singleshot.py.
If you use this code for your research, please cite our paper, "Pruning neural networks without any data by iteratively conserving synaptic flow".