EdMIPS: Rethinking Differentiable Search for Mixed-Precision Neural Networks

EdMIPS: Rethinking Differentiable Search for Mixed-Precision Neural Networks

by Zhaowei Cai, and Nuno Vasconcelos.

This implementation is written by Zhaowei Cai at UC San Diego.

Introduction

EdMIPS is an efficient algorithm to search the optimal mixed-precision neural network directly without proxy task on ImageNet given computation budgets. It can be applied to many popular network architectures, including ResNet, GoogLeNet, and Inception-V3. More details can be found in the paper.

Citation

If you use our code/model/data, please cite our paper:

@inproceedings{cai20edmips,
  author = {Zhaowei Cai and Nuno Vasconcelos},
  Title = {Rethinking Differentiable Search for Mixed-Precision Neural Networks},
  booktitle = {CVPR},
  Year  = {2020}
}

Installation

1. Install PyTorch and ImageNet dataset following the official PyTorch ImageNet training code.

2. Clone the EdMIPS repository, and we'll call the directory that you cloned EdMIPS into EdMIPS_ROOT

   git clone https://github.com/zhaoweicai/EdMIPS.git
   cd EdMIPS_ROOT/

Searching the Mixed-precision Network with EdMIPS

You can start training EdMIPS. Take ResNet-18 for example.

python search.py \
  -a mixres18_w1234a234 --epochs 25 --step-epoch 10 --lr 0.1 --lra 0.01 --cd 0.00335 -j 16 \
  [your imagenet-folder with train and val folders]

The other network architectures are also available, including ResNet-50, GoogLeNet and Inception-V3.

Training the Searched Mixed-precision Network

After the EdMIPS searching is finished, with the checkpoint arch_checkpoint.pth.tar, you can start to train the classification model with the learned bit allocation.

python main.py \
  -a quantres18_cfg --epochs 95 --step-epoch 30 -j 16 \
  --ac arch_checkpoint.pth.tar \
  [your imagenet-folder with train and val folders]

Results

The results are shown as following:

Disclaimer

1. The training of EdMIPS has some variance. Tune --cd a little bit to get the optimal bit allocation you want.

2. The BitOps are counted only on the quantized layers. They are normalized to the bit space as in the above table.

3. Since some changes have been made after the paper submission, you may get slightly worse performances (0.1~0.2 points) than those in the paper.

If you encounter any issue when using our code/model, please let me know.