Repository files navigation

Code used for paper: [ICLR 2022]Learning Pruning-Friendly Networks via Frank-Wolfe: One-Shot, Any-Sparsity, and No Retraining.

Lu Miao*, Xiaolong Luo*, Tianlong Chen, Wuyang Chen, Dong Liu, Zhangyang Wang

A novel framework to train a large deep neural network (DNN) for onlyonce, which can then be pruned toany sparsity ratio to preserve competitive accuracywithout any re-training. We propose a sparsity-aware one-shot pruning method based on K-sparse polytope constraint and Stochastic Frank-Wolfe (SFW) optimizer. We also present the firstlearning-based initialization scheme specifically for boosting SFW-based DNN training.

Required environment

• pytorch >= 1.5.0
• torchvision

The following codes can reproduce the experiments involved in the paper.

SFW training for one-shot pruning

The following code is the training step in SFW-pruning framework.

python-utrain_prune.py--datacifar10--archResNet18--optimizerSFW--constraintk_sparse_constraints--lr1.0--lr_schemedynamic_change--momentum0.9--weight_decay0--k_sparseness10--k_frac0.05--tau15--modeinitialization--rescalegradient--sfw_init0--train_batchsize128--test_batchsize128--epoch_num180--color_channel3--gpu-1

The log file is saved in/saved_logs/SFW_one_shot_prune/. The trained model is saved in/saved_models/.

Test pruning performance

The following code conducts (unstructured) pruning and tests the performance of the pruned DNN. Pruning ratios are 10%, 20%, 30%, 40%, 50%, 60%, 70%, 80%, 90%, 91%, 92%, 93%, 94%, 95%, 96%, 97%, 98%, 99%.

python-utest_prune.py--datacifar10--archResNet18--optimizerSFW--constraintk_sparse_constraints--lr1.0--lr_schemedynamic_change--momentum0.9--weight_decay0--k_sparseness10--k_frac0.05--tau15--modeinitialization--rescalegradient--sfw_init0--train_batchsize128--test_batchsize128--epoch_num180--color_channel3--gpu-1

The argument choices are parallel with those oftrain_prune.py. The log file is saved in/saved_logs/SFW_prune_test/.

Check DNN weight distribution

The following code checks out the weight distribution of the DNN.

python-utest_weight_distribution.py--datacifar10--archResNet18--optimizerSFW--constraintk_sparse_constraints--lr1.0--lr_schemedynamic_change--momentum0.9--weight_decay0--k_sparseness10--k_frac0.05--tau15--modeinitialization--rescalegradient--sfw_init0--train_batchsize128--test_batchsize128--epoch_num180--color_channel3--gpu-1

The argument choices are also parallel with those oftrain_prune.py. The log file is saved in/saved_logs/weight_distribution/.

Optional argument choices

Some optional argument choices are as follows.

optionalarguments:--datacifar10|cifar100|mnist|svhn|tiny--archResNet18|VGG16|Mlp--optimizerSFW|SGD--constraintk_sparse_constraints|l2_constraints|unconstraints--lr1.0 (floatbetween0and1)--lr_schemedynamic_change|decrease_3|keep--momentum0.9 (recommand)--weight_decay0 (recommand)--k_sparseness10 (equalstothenumberoflabels)--k_frac0.05|0.01|0.1--tau15|5|10|20--modeinitialization|diameter|radius|None--rescalegradient|diameter|None--sfw_init0|1--train_batchsize128--test_batchsize128--epoch_num180--color_channel3--gpu-1 (GPUidtouse)

If use the dataset 'Tiny-Imagenet', please download the dataset to/data/tiny_imagenet_200/.

TBD