/Winning-The-Lottery-Ahead-of-Time

Code for Winning the Lottery Ahead of Time: Efficient Early Network Pruning (ICML 2022)

Primary LanguagePythonMIT LicenseMIT

Compression via Gradient Flow Preservation

Python 3.7 PyTorch 1.4 MIT

This repository is the official implementation of Winning the Lottery Ahead of Time: Efficient Early Network Pruning published at ICML 2022.

Setup

  • Install virtualenv

pip3 install virtualenv

  • Create environment

virtualenv -p python3 ~/virtualenvs/EarlyCroP

  • Activate environment

source ~/virtualenvs/EarlyCroP/bin/activate

  • Install requirements:

pip install -r requirements.txt

  • If you mean to run the 'Tiny-Imagenet' dataset: download and unpack in /gitignored/data/, then replace CIFAR10 with TINYIMAGENET below to run. Additional datasets can be added in a similar way (Imagewoof, imagenette, etc.)

Image Classification

To reproduce Image Classification results refer to Image Classification

NLP

To reproduce NLP results on the PSMM network refer to the folder NLP

Cite

Please cite our paper if you use our code in your own work:

@InProceedings{earlycrop,
  title = 	 {Winning the Lottery Ahead of Time: Efficient Early Network Pruning},
  author =       {Rachwan, John and Z{\"u}gner, Daniel and Charpentier, Bertrand and Geisler, Simon and Ayle, Morgane and G{\"u}nnemann, Stephan},
  booktitle = 	 {Proceedings of the 39th International Conference on Machine Learning},
  year = 	 {2022},
  volume = 	 {162},
  series = 	 {Proceedings of Machine Learning Research},
  publisher =    {PMLR},
  }

Licence

MIT Licence