Fully Automated Model Compression with Intergrated Pruning, Quantizationand Distillation
This repository contains all the experiments of our paper "Fully Automated Model Compression with Intergrated Pruning, Quantizationand Distillation". It also includes some pretrain_models which we list in the paper.
Requirements
pip install -r requirements.txt
Running
feature extract
# [optional]cache imagenet dataset in RAM for accelerting I/O
code_path='/ITPruner/Imagenet/'
chmod +x ${code_path}/prep_imagenet.sh
cd ${code_path}
echo "preparing data"
bash ${code_path}/prep_imagenet.sh >> /dev/null
echo "preparing data finished"
python3 -m torch.distributed.launch --nproc_per_node=1 feature_extract.py \
--model "mobilenet" \
--path "Exp_base/mobilenet_base" \
--dataset "imagenet" \
--save_path '/home/data/imagenet' \
--test_batch_size 250 \
--target_bops 7000000000 \
--beta 243 \
--add_channel_rate 0.1 \
--a_bits 5 \
or
bash ./run_scripts/feature_extract/a5/extract_mobilenet_7G.sh
Train
Because of random seed, cfg obtained through feature extraction may have a little difference from ours. Our cfg are given in .sh files.
# [optional]cache imagenet dataset in RAM for accelerting I/O
code_path='/home/wmk/codes/new/ITPruner/Imagenet/'
chmod +x ${code_path}/prep_imagenet.sh
cd ${code_path}
echo "preparing data"
bash ${code_path}/prep_imagenet.sh >> /dev/null
echo "preparing data finished"
python3 -m torch.distributed.launch --nproc_per_node=4 train.py --train \
--model "mobilenet" \
--path "Exp_train/train_mobilenet_50m_${RANDOM}" \
--dataset "imagenet" \
--save_path '/home/data/imagenet' \
--base_path "Exp_base/mobilenet_base" \
--warm_epoch 1 \
--sync_bn \
--n_epochs 250 \
--train_batch_size 256 \
--test_batch_size 250 \
--label_smoothing 0.1 \
--a_bits 5 \
--q_type 0 \
--q_level 0 \
--weight_observer 0 \
--init_lr 0.1 \
--cfg "[32, 32, 64, 64, 151, 128, 280, 256, 256, 256, 256, 256, 512, 716]" \
--quant_cfg "[4, 5, 6, 4, 4, 5, 7, 6, 8, 4, 4, 4, 6, 4]" \
or
bash ./run_scripts/train/a5/train_mobilenet_7G.sh
Evaluate
We provide some pretrain_models which we list in the paper.
python3 -m torch.distributed.launch --nproc_per_node=1 evaluate.py \
--model "mobilenet" \
--path "pretrain_models/train_mobilenet_150m_31752" \
--dataset "imagenet" \
--save_path 'your_data_path' \
--cfg "[23, 42, 63, 67, 132, 134, 275, 194, 202, 202, 194, 277, 522, 687]"
or
bash ./run_scripts/evaluate/evaluate_mobilenet_150m.sh