SWAP Network Pruning

Official Code of SWAP: Sparse Entropic Wasserstein Regression for Robust Network Pruning[ICLR2024]

1. Setup the enviroment:

   conda create -n cbs python=3.6.13
   conda activate cbs
   conda install pytorch==1.5.0 torchvision==0.6.0 cudatoolkit=10.2 -c pytorch
   sh setup.sh
   
   

2. Prepare data

   mkdir prob_regressor_data
   mkdir prob_regressor_results
   mkdir checkpoints
   

   Please download checkpoints from here and put them in the "checkpoints" folder.

3. Run Entropic Wasserstein Pruning:

   3.1 on MLPNet

   sh scripts/sweep_mnist_mlpnet_ot.sh
   

   3.2 on ResNet20

   sh scripts/sweep_cifar10_resnet20_ot.sh
   

   3.3 on MobileNetV1

   sh scripts/sweep_imagenet_mobilenet_ot.sh
   

Citaton

We thank Singh & Alistarh for sharing their code of WoodFisher. We thank also Yu, Xin and Serra et. al. for the CBS code CBS, from which the repository is forked. Our implementation is based on their code.

If our work is helpful to your research/project, we appreciate you to cite.