• Fully Automated Model Compression with Intergrated Pruning, Quantizationand Distillation
• Requirements
• Running

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

• DALI
• Apex
• torchprofile
• other requirements, running requirements.txt

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