An extension version of our CVPR 2020, oral: HRank: Filter Pruning using High-Rank Feature Map (Link).

Prior code version can be found here.

Tips

Any problem, please contact the authors via emails: [email protected] or [email protected]. Do not post issues with github as much as possible, just in case that I could not receive the emails from github thus ignore the posted issues.

Citation

If you find HRank useful in your research, please consider citing:

@inproceedings{lin2020hrank,
  title={HRank: Filter Pruning using High-Rank Feature Map},
  author={Lin, Mingbao and Ji, Rongrong and Wang, Yan and Zhang, Yichen and Zhang, Baochang and Tian, Yonghong and Shao, Ling},
  booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)},
  pages={1529--1538},
  year={2020}
}

Running Code

In this code, you can run our models on CIFAR-10 and ImageNet dataset. (Task on object detection is preparing)

For rank generation, the code has been tested by Python 3.6, Pytorch 1.0 and CUDA 9.0 on Ubuntu 16.04. And for evaluating(fine-tuning), we recommend you to run with PyTorch 1.6 for the usage of Thop.

Rank Generation

python rank_generation.py \
--dataset [dataset name] \
--data_dir [dataset dir] \
--pretrain_dir [pre-trained model dir] \
--arch [model arch name] \
--limit [batch numbers] \
--gpu [gpu_id]

For the ease of reproducibility, we provide the extracted ranks here:

Model Training

For model training, our code will automatically save the previous-epoch checkpoint file in checkpoint.pth.tar and best checkpoint in model_best.pth.tar after every epoch. And with the resume option, you can resume training from the previous checkpoint under the same job directory.

1. CIFAR-10

python evaluate_cifar.py \
--data_dir [CIFAR-10 dataset dir] \
--job_dir ./result/[model name]/[folder name] \
--arch [model name](vgg_16_bn, resnet_56, resnet_110, googlenet, densenet_40) \
--rank_conv_prefix [rank folder dir] \
--use_pretrain \
--pretrain_dir [pre-trained model dir] \
--compress_rate [compress rate] \
--gpu [gpu_id]

2. ImageNet

python evaluate.py \
--data_dir [ImageNet dataset dir] \
--use_dali \
--job_dir ./result/[model name]/[folder name] \
--arch [model name](resnet_50, mobilenet_v1, mobilenet_v2) \
--rank_conv_prefix [rank folder dir] \
--lr_type (cos or step) \
--use_pretrain \
--pretrain_dir [pre-trained model dir] \
--compress_rate [compress rate] \
--gpu [gpu_id]

3. Optional arguments(ImageNet)

optional arguments:
    --data_dir				Dataset directory.
    --job_dir				The directory where the summaries will be stored.
    --rank_conv_prefix 			Prefix directory of rank folder.
    --compress_rate 			Compress rate of each convolutional layer.
    --arch				Architecture of model. default: vgg_16_bn 
    					Optional: resnet_50, mobilenet_v2, mobilenet_v1, vgg_16_bn, resnet_56, resnet_110, densenet_40, googlenet
    --learning_rate			Initial learning rate. default: 0.1
    --lr_type 				Learning rate decay schedule. default: step. optional: step, cos
    --momentum   			Momentum for optimizer. default: 0.9
    --epochs				The number of epochs to train. default: 90
    --batch_size			Batch size for both training and validation. default: 64
    --weight_decay 			The weight decay of loss function. default: 1e-4
    --label_smooth			Label smooth parameter. default: 0.1
    
    --resume				If this parameter exists, model training by resuming from previous ckpt in the same directory.
    --use_dali 				If this parameter exists, use dali module to load ImageNet data.
    --use_pretrain			If this parameter exists, use pretrained model for finetuning.
    --pretrain_dir			Pretrained model directory.
    --test_only				If this parameter exists, only validate the model performance without training.
    --test_model_dir			Test model directory
    --gpu	 			Select gpu to use. default: 0

The following are the examples of compression rate setting for several models: (Please note that the following compression rates are only used to demonstrate the parameter format, which may not be used in our experiment. For the specific pruning rate, please refer to the configuration files of pruned model in the next section)

With the compression rate, our training module (evaluate_cifar.py and evaluate.py) can automatically calculate the params and FLOPs of that model and record them in the training logger.

After training, a total of four files can be found in the job_dir, including best model, final model, config file and logger file.

Experimental Results

We provide our pre-trained models, and pruned models. For your ease of reproducibility, the training loggers and configuration files are attached with the url link as well.

(The percentages in parentheses indicate the pruning rate)

CIFAR-10

ImageNet (Lacking GPUs... Coming as soon as Possible)

To verify our model performance, please use the command below (make sure you are using the corresponding compression rate in the configuration file of that model):

CIFAR-10:

python evaluate_cifar.py \
--data_dir [CIFAR-10 dataset dir] \
--job_dir ./result/[model name]/[folder name] \
--arch [model name](vgg_16_bn, resnet_56, resnet_110, googlenet, densenet_40) \
--test_only \
--test_model_dir [test model dir] \
--compress_rate [compress rate] \
--gpu [gpu_id]

ImageNet:

python evaluate.py \
--data_dir [ImageNet dataset dir] \
--job_dir ./result/[model name]/[folder name] \
--arch [model name](resnet_50, mobilenet_v1, mobilenet_v2) \
--test_only \
--test_model_dir [test model dir] \
--compress_rate [compress rate] \
--gpu [gpu_id]

Pre-trained Models

Additionally, we provide the pre-trained models used in our experiments.

CIFAR-10:

Vgg-16 | ResNet56 | ResNet110 | DenseNet-40 | GoogLeNet

ImageNet:

ResNet50