Importance Estimation for Neural Network Pruning
This repo contains required scripts to reproduce results from paper:
Importance Estimation for Neural Network Pruning
Pavlo Molchanov, Arun Mallya, Stephen Tyree, Iuri Frosio, Jan Kautz .
In CVPR 2019.
License
Copyright (C) 2019 NVIDIA Corporation.
All rights reserved. Licensed under the CC BY-NC-SA 4.0 (Attribution-NonCommercial-ShareAlike 4.0 International)
The code is released for academic research use only. For commercial use, please contact [email protected].
Installation
Clone this repo.
git clone https://github.com/NVlabs/Taylor_pruning.git
cd Taylor_pruning/This code requires PyTorch 1.0 and python 3+. Please install dependencies by
pip install -r requirements.txtFor the best reproducibility of results you will need NVIDIA DGX1 server with 8 V100. During pruning and finetuning we at most use 4 GPUs.
The code was tested with python3.6 the following software versions:
| Software | version |
|---|---|
| cuDNN | v7500 |
| Pytorch | 1.0.1.post2 |
| CUDA | v10.0 |
Preparation
Dataset preparation
Pruning examples use ImageNet 1k dataset which needs to be downloaded beforehand. Use standard instructions to setup ImageNet 1k for Pytorch, e.g. from here.
Models preparation
We use pretrained models provided with Pytorch, they need to be downloaded:
wget https://download.pytorch.org/models/resnet50-19c8e357.pth
wget https://download.pytorch.org/models/resnet101-5d3b4d8f.pth
mkdir ./models/pretrained/
mv {resnet50-19c8e357.pth,resnet101-5d3b4d8f.pth} ./models/pretrained/
Pruning examples
Results in Table 3 can be reproduced the following script running on 4 V100 GPUs:
python main.py --name=runs/resnet50/resnet50_prune72 --dataset=Imagenet \
--lr=0.001 --lr-decay-every=10 --momentum=0.9 --epochs=25 --batch-size=256 \
--pruning=True --seed=0 --model=resnet50 --load_model=./models/pretrained/resnet50-19c8e357.pth \
--mgpu=True --group_wd_coeff=1e-8 --wd=0.0 --tensorboard=True --pruning-method=22 \
--data=/imagenet/ --no_grad_clip=True --pruning_config=./configs/imagenet_resnet50_prune72.json
Note: we run finetuning for 25 epochs and do not use weight decay. Better model can be obtained by training longer, increasing learning and using weight decay.
Example of config file ./configs/imagenet_resnet50_prune72.json:
{
"method": 22,
"frequency" : 30,
"prune_per_iteration" : 100,
"maximum_pruning_iterations" : 32,
"starting_neuron" : 0,
"fixed_layer" : -1,
"l2_normalization_per_layer": false,
"rank_neurons_equally": false,
"prune_neurons_max": 3200,
"use_momentum": true,
"pruning_silent": false,
"pruning_threshold" : 100.0,
"start_pruning_after_n_iterations" : 0,
"push_group_down" : false,
"do_iterative_pruning" : true,
"fixed_criteria" : false,
"seed" : 0,
"pruning_momentum" : 0.9
}
We provide config files for pruning to 56%, 72%, 81%, 91% of original ResNet-50 and to 40%, 50%, 55%, 75% of ResNet-101 models. Percentage means the ratio of gates to be active after pruning.
Parameter description
Pruning methods (different criteria) are encoded with integer as:
| method id | name | description | comment |
|---|---|---|---|
| 22 | Taylor_gate | Gate after BN in Table 2, Taylor FO in Table 1, Taylor-FO-BN in Table 3 | Best method |
| 0 | Taylor_weight | Conv weight/conv/linear weight with Taylor FO In Table 2 and Table 1 | |
| 1 | Random | Random | |
| 2 | Weight norm | Weight magnitude/ weight | |
| 3 | Weight_abs | Not used | |
| 6 | Taylor_output | Taylor-output as in [27] | |
| 10 | OBD | Optimal Brain Damage | |
| 11 | Taylor_gate_SO | Taylor SO | |
| 23 | Taylor_gate_FG | uses gradient per example to compute Taylor FO, Taylor FO- FG in Table 1, Gate after BN - FG in Table 2 | |
| 30 | BN_weight | BN scale in Table 2 | |
| 31 | BN_Taylor | BN scale Taylor FO in Table 2 |
Citation
If you use this code for your research, please cite our papers.
@inproceedings{molchanov2019taylor,
title={Importance Estimation for Neural Network Pruning},
author={Molchanov, Pavlo and Mallya, Arun and Tyree, Stephen and Frosio, Iuri and Kautz, Jan},
booktitle={Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition},
year={2019}
}
To come
- Config files for DenseNet and VGG16 pruning on ImageNet
- Examples of CIFAR experiments
- Getting shrinked model after pruning (for ResNet-101 only)
- Oracle estimates of true importance, code to compute correlation with Oracle
Acknowledgments
This code heavily reuses Pytorch example for ImageNet training provided here.