mit-han-lab / haq

[CVPR 2019, Oral] HAQ: Hardware-Aware Automated Quantization with Mixed Precision
https://hanlab.mit.edu/projects/haq/
MIT License
369 stars 86 forks source link
automl efficient-model mixed-precision quantization

HAQ: Hardware-Aware Automated Quantization with Mixed Precision

Introduction

This repo contains PyTorch implementation for paper HAQ: Hardware-Aware Automated Quantization with Mixed Precision (CVPR2019, oral)

overview

@inproceedings{haq,
author = {Wang, Kuan and Liu, Zhijian and Lin, Yujun and Lin, Ji and Han, Song},
title = {HAQ: Hardware-Aware Automated Quantization With Mixed Precision},
booktitle = {IEEE Conference on Computer Vision and Pattern Recognition (CVPR)},
year = {2019}
}

Other papers related to automated model design:

Dependencies

We evaluate this code with Pytorch 1.1 (cuda10) and torchvision 0.3.0, you can install pytorch with conda:

# install pytorch
conda install -y pytorch torchvision cudatoolkit=10.0 -c pytorch

And you can use the following command to set up the environment:

# install packages and download the pretrained model
bash run/setup.sh

(If the server is down, you can download the pretrained model from google drive: mobilenetv2-150.pth.tar)

Current code base is tested under following environment:

  1. Python 3.7.3
  2. PyTorch 1.1
  3. torchvision 0.3.0
  4. numpy 1.14
  5. matplotlib 3.0.1
  6. scikit-learn 0.21.0
  7. easydict 1.8
  8. progress 1.4
  9. tensorboardX 1.7

Dataset

If you already have the ImageNet dataset for pytorch, you could create a link to data folder and use it:

# prepare dataset, change the path to your own
ln -s /path/to/imagenet/ data/

If you do not have the ImageNet yet, you can download the ImageNet dataset and move validation images to labeled subfolders. To do this, you can use the following script: https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh

We use a subset of ImageNet in the linear quantizaiton search phase to save the training time, to create the link of the subset, you can use the following tool:

# prepare imagenet100 dataset
python lib/utils/make_data.py

Reinforcement learning search

Finetune Policy

You can evaluate the K-Means quantized model like this:

# evaluate K-Means quantization
bash run/run_kmeans_quantize_eval.sh
Models preserve ratio Top1 Acc (%) Top5 Acc (%)
resnet50 (original) 1.0 76.15 92.87
resnet50 (10x compress) 0.1 75.48 92.42

You can evaluate the linear quantized model like this:

# evaluate linear quantization
bash run/run_linear_quantize_eval.sh
Models preserve ratio Top1 Acc (%) Top5 Acc (%)
mobilenetv2 (original) 1.0 72.05 90.49
mobilenetv2 (0.6x latency) 0.6 71.23 90.00