OMPQ: Orthogonal Mixed Precision Quantization

This repository contains all the experiments of our paper "OMPQ: Orthogonal Mixed Precision Quantization". It also includes some base models and pretrain models which we list in the paper.

Requirements

DALI (for accelerating data processing)
Apex (for distributed running)
other requirements, running requirements.txt

pip install -r requirements.txt

Running

To start running our code to get the optimal bit configuration, you need to download the basemodel, and copy the path of base model to "--path".

Bit Configuration

#!/usr/bin/env bash
python3 -m torch.distributed.launch --nproc_per_node=1 feature_extract.py \
 --model "resnet18" \
 --path "/Path/to/Base_model" \      # pretrained base model
 --dataset "imagenet" \
 --save_path '/Path/to/Dataset/' \   # Dataset path
 --beta 10.0 \                       # Hyper-parameter for bit difference
 --model_size 6.7 \                  # Target model size
 --quant_type "QAT"                  # Post-Training Quantization(PTQ) or Quantization-Aware Training(QAT)

bash ./mixed_bit/run_scripts/quant_resnet18.sh

QAT

Because of random seed, bit configuration obtained through feature extraction may have a little difference from ours. Our bit configurations are given in bit_config.py. Our quantized models and logs are also given in this link.

#!/usr/bin/env bash
python quant_train.py \
 -a resnet18 \
 --epochs 90 \
 --lr 0.0001 \
 --batch_size 128 \
 --data /Path/to/Dataset/ \
 --save_path /Path/to/Save_quant_model/ \
 --act_range_momentum=0.99 \
 --wd 1e-4 \
 --data_percentage 1 \
 --pretrained \
 --fix_BN \
 --checkpoint_iter -1 \
 --quant_scheme modelsize_6.7_a6_75B

bash ./QAT/run_scripts/train_resnet18.sh

PTQ

For the post-training quantization, we only require a few GPU hours to get the quantization model by running codes as follows:

python main_imagenet.py --data_path /Path/to/Dataset/ --arch resnet18 --n_bits_w 2 --channel_wise --n_bits_a 8 --act_quant --test_before_calibration --bit_cfg "[4, 3, 3, 4, 4, 4, 4, 4, 4, 4, 3, 3, 4, 4, 3, 3, 3, 3]"

bash ./PTQ/run_scripts/train_resnet18.sh

Experimental Results

Table 1 and Table 2 in "OMPQ: Orthogonal Mixed Precision Quantization".