Skip to content

Latest commit

 

History

History
72 lines (46 loc) · 2.92 KB

README.md

File metadata and controls

72 lines (46 loc) · 2.92 KB
Raw

Differentiable Gaussian Mixture Weight Sharing Network Quantization, ICML 2022 Spotlight

Finding the Task-Optimal Low-Bit Sub-Distribution in Deep Neural Networks, ICML 2022
Runpei Dong*, Zhanhong Tan*, Mengdi Wu, Linfeng Zhang, and Kaisheng Ma

Created by Runpei Dong*, Zhanhong Tan*, Mengdi Wu, Linfeng Zhang, and Kaisheng Ma.

PMLR | arXiv | Models

This repository contains the code release of the paper Finding the Task-Optimal Low-Bit Sub-Distribution in Deep Neural Networks (ICML 2022).

Installation

Our code works with Python 3.8.3. we recommend to use Anaconda, and you can install the dependencies by running:

$ python3 -m venv env
$ source env/bin/activate
(env) $ python3 -m pip install -r requirements.txt

How to Run

The main procedures are written in script main.py, please run the following command for instructions:

$ python main.py -h

Datasets

Before running the code, you can specify the path for datasets in config.py, or you can specify it by --train-dir and --val-dir.

Training on ImageNet

We have provided a simple SHELL script to train a 4-bit ResNet-18 with DGMS. Run:

$ sh tools/train_imgnet.sh

Inference on ImageNet

To inference compressed models on ImageNet, you only need to follow 2 steps:

  • Step-1: Download the checkpoints released on Google Drive.

  • Step-2: Run the inference SHELL script we provide:

    $ sh tools/validation.sh

Q-SIMD

The TVM based Q-SIMD codes can be download from Google Drive.

Citation

If you find our work useful in your research, please consider citing:

@inproceedings{dong2021finding,
  title={Finding the Task-Optimal Low-Bit Sub-Distribution in Deep Neural Networks},
  author={Dong, Runpei and Tan, Zhanhong and Wu, Mengdi and Zhang, Linfeng and Ma, Kaisheng},
  booktitle={Proceedings of the International Conference on Machine Learning (ICML)},
  year={2022}
}

License

DGMS is released under the Apache 2.0 license. See the LICENSE file for more details.