gktejus / edmips Goto Github PK

by Zhaowei Cai, and Nuno Vasconcelos.

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

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.

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}
}

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/

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.

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]

The results are shown as following:

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.