Sharpness-aware Quantization for Deep Neural Networks

Recent Update

2021.11.23: We release the source code of SAQ.

Setup the environments

1. Clone the repository locally:

git clone https://github.com/zhuang-group/SAQ

2. Install pytorch 1.8+, tensorboard and prettytable

conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
pip install tensorboard
pip install prettytable

Data preparation

ImageNet

1. Download the ImageNet 2012 dataset from here, and prepare the dataset based on this script.

2. Change the dataset path in link_imagenet.py and link the ImageNet-100 by

python link_imagenet.py

CIFAR-100

Download the CIFAR-100 dataset from here.

After downloading ImageNet and CIFAR-100, the file structure should look like:

dataset
├── imagenet
    ├── train
    │   ├── class1
    │   │   ├── img1.jpeg
    │   │   ├── img2.jpeg
    │   │   └── ...
    │   ├── class2
    │   │   ├── img3.jpeg
    │   │   └── ...
    │   └── ...
    └── val
        ├── class1
        │   ├── img4.jpeg
        │   ├── img5.jpeg
        │   └── ...
        ├── class2
        │   ├── img6.jpeg
        │   └── ...
        └── ...
├── cifar100
    ├── cifar-100-python
    │   ├── meta
    │   ├── test
    │   ├── train
    │   └── ...
    └── ...

Training

Fixed-precision quantization

1. Download the pre-trained full-precision models from the model zoo.

2. Train low-precision models.

To train low-precision ResNet-20 on CIFAR-100, run:

sh script/train_qsam_cifar_r20.sh

To train low-precision ResNet-18 on ImageNet, run:

sh script/train_qsam_imagenet_r18.sh

Mixed-precision quantization

1. Download the pre-trained full-precision models from the model zoo.

2. Train the configuration generator.

To train the configuration generator of ResNet-20 on CIFAR-100, run:

sh script/train_generator_cifar_r20.sh

To train the configuration generator on ImageNet, run:

sh script/train_generator_imagenet_r18.sh

3. After training the configuration generator, run following commands to fine-tune the resulting models with the obtained bitwidth configurations on CIFAR-100 and ImageNet.

sh script/finetune_cifar_r20.sh

sh script/finetune_imagenet_r18.sh

Results on CIFAR-100

Results on ImageNet

License

This repository is released under the Apache 2.0 license as found in the LICENSE file.

Acknowledgement

This repository has adopted codes from SAM, ASAM and ESAM, we thank the authors for their open-sourced code.