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).

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

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

$ python main.py -h

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.

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

$ sh tools/train_imgnet.sh

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

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

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

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