This is an unofficial PyTorch implementation of the paper SNIP: Single-shot Network Pruning based on Connection Sensitivity by Namhoon Lee, Thalaiyasingam Ajanthan and Philip H. S. Torr.
It doesn not cover all the experiment in the paper but it does include the main ones:
- LeNet5-Caffe on MNIST
- VGG-D on CIFAR-10
I haven't had the time to add an argparser yet the network type and pruning level should be changed directly in the code.
This has been tested with Python 3.7.1 and PyTorch 1.0.0. The exact environment can be replicated by:
$ conda env create -f environment.yml
This would create a conda environment called snip-env
.
$ conda activate snip-env
$ python train.py
Three runs with different seeds with LeNet5-Caffe on MNIST (sparsity level of 98%):
Two runs with different seeds with VGG-D on CIFAR-10 (sparsity level of 95%):