mil-ad/snip

Pytorch implementation of the paper "SNIP: Single-shot Network Pruning based on Connection Sensitivity" by Lee et al.

SNIP: Single-shot network pruning based on connection sensitivity

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.

Environment

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.

Usage

$ conda activate snip-env
$ python train.py

Results

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%):

snip's Issues

code problem

Thank for you great work! will you release the code adjusting the snip to resnet?

When creating weight masks you turn gradient off weights instead of the masks

Line 36 in snip.py should be
layer.weight_mask.requires_grad = False
not as it is
layer.weight.requires_grad = False
which essentially ruins the entire model :P

dose weights real pruned?

Hi!
weights and its gradients are set to 0 in the code which i think is different from real pruning.

Batch layer comes after Relu in your vgg implementation fully connected layers. That is not common practice in deep nets.

model weight

Hello, my research direction is neural network hardware accelerators. I need model parameters to test. Do you have well-trained models, such as vgg, LeNet, AlexNet, etc. Any network model is fine, I need more models to test, please send me an email, [email protected], thank you very much.