Home

Awesome

EagleEye: Fast Sub-net Evaluation for Efficient Neural Network Pruning

Python version support PyTorch version support

PyTorch implementation for EagleEye: Fast Sub-net Evaluation for Efficient Neural Network Pruning

Bailin Li, Bowen Wu, Jiang Su, Guangrun Wang, Liang Lin

Presented at ECCV 2020 (Oral)

Check slides about EagleEye: “High-performance AI on the Edge: from perspectives of model compression and hardware architecture design“, DMAI HiPerAIR, Aug. 2020.

pipeline

Citation

If you use EagleEye in your research, please consider citing:

@misc{li2020eagleeye,
    title={EagleEye: Fast Sub-net Evaluation for Efficient Neural Network Pruning},
    author={Bailin Li and Bowen Wu and Jiang Su and Guangrun Wang and Liang Lin},
    year={2020},
    eprint={2007.02491},
    archivePrefix={arXiv},
    primaryClass={cs.CV}
}

Update

Adaptive-BN-based Candidate Evaluation

For the ease of your own implementation, here we present the key code for proposed Adaptive-BN-based Candidate Evaluation. The official implementation will be released soon.

def eval_pruning_strategy(model, pruning_strategy, dataloader_train):
   # Apply filter pruning to trained model
   pruned_model = prune(model, pruning_strategy)

   # Adaptive-BN
   pruned_model.train()
   max_iter = 100
   with torch.no_grad():
      for iter_in_epoch, sample in enumerate(dataloader_train):
            pruned_model.forward(sample)
            if iter_in_epoch > max_iter:
                break

   # Eval top-1 accuracy for pruned model
   acc = pruned_model.get_val_acc()
   return acc

Baseline Model Training

The code used for training baseline models(MobileNetV1, ResNet50) will be released at CNNResearchToolkit. Welcome everyone to follow!

Setup

  1. Prepare Data

    Download ILSVRC2012 dataset from http://image-net.org/challenges/LSVRC/2012/index#introduction

  2. Download Pretrained Models

    We provide pretrained baseline models and reported pruned models in Google Drive. Please put the downloaded models in the dir of models/ckpt/.

  3. Prepare Runtime Environment

    Via pip/conda

    pip install -r requirements.txt
    

    Via Docker

    # Build Image
    docker build docker/ -t eagleeye:[tag]
    
    # launch docker container
    docker run -it --rm \
     -v [PATH-TO-EAGLEEYE]:/workspace/EagleEye \
     -v [PATH-TO-IMAGENET]:/data/imagenet \
     --ipc=host \
     eagleeye:[tag]
    

Usage

Our proposed EagleEye contains 3 steps:

  1. Adaptive-BN-based Searching for Pruning Strategy
  2. Candidate Selection
  3. Fine-tuning of Pruned Model

1. Adaptive-BN-based Searching for Pruning Strategy

On this step, pruning strategies are randomly generated. Then, Adaptive-BN-based evaluation are performed among these pruning strategies. Pruning strategies and their eval scores will be saved to search_results/pruning_strategies.txt.

If you do not want to perform searching by yourself, the provided search result could be found in search_results/.

Parameters involved in this steps:

NameDescription
--flops_targetThe remaining ratio of FLOPs of pruned model
--max_rate<br>--min_rateDefine the search space. The search space is [min_rate, max_rate]
--output_fileFile stores the searching results.

Sample scripts could refer to 1. Search of scripts/mbv1_50flops.sh.

Searching space for different models

ModelPruned FLOPs[min_rate, max_rate]
MobileNetV1-50%[0, 0.7]
ResNet50-25%[0, 0.4]
ResNet50-50%[0, 0.7]
ResNet50-75%[0, 0.8]

2. Candidate Selection

On this step, best pruning strategy is picked from output_file generated on step1.

The output looks like as following:

########## pruning_strategies.txt ##########
strategy index:84, score:0.143
strategy index:985, score:0.123

Sample scripts could refer to 2. Selection of scripts/mbv1_50flops.sh.

3. Fine-tuning of Pruned Model

This step take strategy index as input and perform fine-tuning on it.

Parameters involved in this steps:

NameDescription
--search_resultSearching results
--strategy_idIndex of best pruning strategy from step2
--lrLearning rate for fine-tuning
--weight_decayWeight decay while fine-tuning
--epochNumber of fine-tuning epoch

Sample scripts could refer to 3. Fine-tuning of scripts/mbv1_50flops.sh.

Inference of Pruned Model

For ResNet50:

python3 inference.py \
--model_name resnet50 \
--num_classes 1000 \
--checkpoint models/ckpt/{resnet50_25flops.pth|resnet50_50flops.pth|resnet50_72flops.pth} \
--gpu_ids 4 \
--batch_size 512 \
--dataset_path {PATH_TO_IMAGENET} \
--dataset_name imagenet \
--num_workers 20

For MobileNetV1:

python3 inference.py \
--model_name mobilenetv1 \
--num_classes 1000 \
--checkpoint models/ckpt/mobilenetv1_50flops.pth \
--gpu_ids 4 \
--batch_size 512 \
--dataset_path {PATH_TO_IMAGENET} \
--dataset_name imagenet \
--num_workers 20

After running above program, the output looks like below:

######### Report #########                                                                                                                                                  
Model:resnet50
Checkpoint:models/ckpt/resnet50_50flops_7637.pth
FLOPs of Original Model:4.089G;Params of Original Model:25.50M
FLOPs of Pruned   Model:2.057G;Params of Pruned   Model:14.37M
Top-1 Acc of Pruned Model on imagenet:0.76366
##########################

Results

Quantitative analysis of correlation

Correlation between evaluation and fine-tuning accuracy with different pruning ratios (MobileNet V1 on ImageNet classification Top-1 results)

corr

Results on ImageNet

ModelFLOPsTop-1 AccTop-5 AccCheckpoint
ResNet-503G<br>2G<br>1G77.1%<br>76.4%<br>74.2%93.37%<br>92.89%<br>91.77%resnet50_75flops.pth <br> resnet50_50flops.pth <br> resnet50_25flops.pth
MobileNetV1284M70.9%89.62%mobilenetv1_50flops.pth

Results on CIFAR-10

ModelFLOPsTop-1 Acc
ResNet-5062.23M94.66%
MobileNetV126.5M<br>12.1M<br>3.3M91.89% <br> 91.44% <br> 88.01%