This repository is the official implementation of the paper Pruning via Iterative Ranking of Sensitivity Statistics. Currently under review. Please use this preliminary BibTex entry when referring to our work:

@article{verdenius2020pruning,
       author = {{Verdenius}, Stijn and {Stol}, Maarten and {Forr{\'e}}, Patrick},
        title = "{Pruning via Iterative Ranking of Sensitivity Statistics}",
      journal = {arXiv e-prints},
     keywords = {Computer Science - Machine Learning, Statistics - Machine Learning},
         year = 2020,
        month = jun,
          eid = {arXiv:2006.00896},
        pages = {arXiv:2006.00896},
archivePrefix = {arXiv},
       eprint = {2006.00896},
 primaryClass = {cs.LG},
}

The repository implements novel pruning / compression algorithms for deep learning / neural networks. Additionally, it implements the shrinkage of actual tensors to really benefit from structured pruning without external hardware libraries. We implement:

• Structured (node) pruning before training
• Structured (node) pruning during training
• Unstructured (weight) pruning before training
• Unstructured (weight) pruning during training

• Install virtualenv

pip3 install virtualenv

• Create environment

virtualenv -p python3 ~/virtualenvs/SNIPIT

• Activate environment

source ~/virtualenvs/SNIPIT/bin/activate

• Install requirements:

pip install -r requirements.txt

• If you mean to run the 'Imagenette' dataset: download that from here and unpack in /gitignored/data/, then replace CIFAR10 with IMAGENETTE below to run. Additional datasets can be added in a similar way (Imagewoof, tiny-imagenet, etc.)

Some examples of training the models from the paper.

To run training for SNAP-it - our structured pruning before training algorithm - with a VGG16 on CIFAR10, run the following:

python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion SNAPit --pruning_limit 0.93 --epochs 80

To run training for SNIP-it - our unstructured pruning algorithm - with a ResNet18 on CIFAR10, run one of the following:

## during training
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion SNIPitDuring --pruning_limit 0.98 --outer_layer_pruning --epochs 80 --prune_delay 4 --prune_freq 4
## before training
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion SNIPit --pruning_limit 0.98 --outer_layer_pruning --epochs 80 

To evaluate a model on adversarial attacks (for now only supported on unstructured), run:

python main.py --eval --model MLP5 --data_set MNIST --checkpoint_name <see_results_folder> --checkpoint_model MLP5_finished --attack CarliniWagner

Results and saved models will be logged to the terminal, logfiles in result-folders and in tensorboard files in the /gitignored/results/ folder. To run tensorboard's interface run the following:

tensorboard --logdir ./gitignored/results/

The regular arguments for running are the following. Additionally, there are some more found in utils/config_utils.py.

Some notes:

• please note that as of now, residual connections (e.g. ResNets) and structured pruning are not supported together.
• please note that as of now, structured pruning and --outer_layer_pruning are not supported together.
• please note that as of now, if running unstructured pruning, you should also run with --outer_layer_pruning.

• Codebase is built modularly so that every criterion or model that is added to its designated folder, provided its filename is equal to its classname, can be ran via string argument immediately. This way its easily extendable.
• The same goes for training schemes; implemented here as classes and automatically loaded in by string reference. When you need new functionality concerning one aspect of training you can simply inheret the DefaultTrainer and then override only that function you need differently. Alternatively, you can make your own training scheme, the sky is the limit!
• All entry-points go through main.py, where the required models are loaded and thereafter redirected to the right training or testing scheme.
• All results show up at the path /gitignored/results/ in its own (date-stamped) folder. In here you find a copy of the codebase at the time of execution, its calling command, tensorboard output, saved models and logs.
• In the file utils/autoconfig.json certain automatic configurations get set to make it easier to run different models in sequence. You can disable this with --disable_autoconfig, but it is strongly recommended against.

## unpruned baselines
python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion EmptyCrit --epochs 80 --pruning_limit 0.0
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion EmptyCrit --epochs 80 --pruning_limit 0.0

## structured baselines
python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion StructuredRandom --pruning_limit 0.93 --epochs 80
python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion GateDecorators --pruning_limit 0.93 --epochs 70 --checkpoint_name <unpruned_results_folder> --checkpoint_model VGG16_finished 
python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion EfficientConvNets --pruning_limit 0.93 --epochs 80 --prune_delay 69 --prune_freq 1
python3 main.py --model VGG16 --data_set CIFAR10 --prune_criterion GroupHoyerSquare --hoyer_reg <REG> --epochs 80 --prune_delay 69 --prune_freq 1 --group_hoyer_square
python3 main.py --model VGG16 --data_set CIFAR10 --l0_reg <REG> --epochs 160 --l0

## unstructured baselines
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion UnstructuredRandom --pruning_rate 0.98 --pruning_limit 0.98 --outer_layer_pruning --epochs 80
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion <SNIP or GRASP> --pruning_limit 0.98 --outer_layer_pruning --epochs 80
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion HoyerSquare --hoyer_reg <REG> --outer_layer_pruning --epochs 80 --prune_delay 69 --prune_freq 1 --hoyer_square
python3 main.py --model ResNet18 --data_set CIFAR10 --prune_criterion IMP --pruning_limit 0.98 --outer_layer_pruning --epochs 80 --prune_delay 4 --prune_freq 4 --enable_rewinding --rewind_to 6

MIT Licence