Skip to content

GuanhuaWang/sensAI

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

sensAI: ConvNets Decomposition via Class Parallelism for Fast Inference on Live Data

Environment

Linux, python 3.6+

Setup

pip install -r requirements.txt

Instruction

Supported CNN architectures and datasets:

Dataset Architecture(ARCH)
CIFAR-10 vgg19_bn, resnet110, resnet164, mobilenetv2, shufflenetv2
CIFAR-100 vgg19_bn, resnet110, resnet164
ImageNet-1K vgg19_bn, resnet50

1. Generate class groups

For CIFAR-10/CIFAR-100:

python3 group_selection.py \
        --arch $ARCH \
        --resume $pretrained_model \
        --dataset $DATASET \
        --ngroups $number_of_groups \
        --gpu_num $number_of_gpu 

For ImageNet-1K:

python3 group_selection.py \
        --arch $ARCH \
        --dataset imagenet \
        --ngroups $number_of_groups \
        --gpu_num $number_of_gpu \
        --data /{path_to_imagenet_dataset}/

Pruning candidate now stored in ./prune_candidate_logs/

2. Prune models

For CIFAR-10/CIFAR-100:

python3 prune_and_get_model.py \
        -a $ARCH \
        --dataset $DATASET \
        --resume $pretrained_model \
        -c ./prune_candidate_logs/ \
        -s ./{TO_SAVE_PRUNED_MODEL_DIR}/

For ImageNet-1K:

python3 prune_and_get_model.py \
        -a $ARCH \
        --dataset imagenet \
        -c ./prune_candidate_logs/ \
        -s ./{TO_SAVE_PRUNED_MODEL_DIR}/ \
        --pretrained

Pruned models are now saved in ./{TO_SAVE_PRUNED_MODEL_DIR}/$ARCH/

3. Retrain pruned models

For CIFAR-10/CIFAR-100:

python3 retrain_grouped_model.py \
        -a $ARCH \
        --dataset $DATASET \
        --resume ./{TO_SAVE_PRUNED_MODEL_DIR}/ \
        --train_batch $batch_size \
        --epochs $number_of_epochs \
        --num_gpus $number_of_gpus

For ImageNet-1K:

python3 retrain_grouped_model.py \
        -a $ARCH \
        --dataset imagenet \
        --resume ./{TO_SAVE_PRUNED_MODEL_DIR}/ \
        --epochs $number_of_epochs \
        --num_gpus $number_of_gpus \
        --train_batch $batch_size \
        --data /{path_to_imagenet_dataset}/

Retrained models now saved in ./{TO_SAVE_PRUNED_MODEL_DIR}_retrained/$ARCH/

4. Evaluate

For CIFAR-10/CIFAR-100:

python3 evaluate.py \
        -a $ARCH \
        --dataset=$DATASET \
        --retrained_dir ./{TO_SAVE_PRUNED_MODEL_DIR}_retrained/ \
        --test-batch $batch_size

For ImageNet-1K:

python3 evaluate.py \
        -d imagenet \
        -a $ARCH \
        --retrained_dir ./{TO_SAVE_PRUNED_MODEL_DIR}_retrained/ \
        --data /{path_to_imagenet_dataset}/

Contributors

Thanks for all the main contributors to this repository:

And many others Zihao Fan, Hank O'Brien , Yaoqing Yang, Adarsh Karnati, Jichan Chung, Yingxin Kang, Balaji Veeramani, Sahil Rao.

Citation

@inproceedings{wang2021sensAI,
 author = {Guanhua Wang and Zhuang Liu and Brandon Hsieh and Siyuan Zhuang and Joseph Gonzalez and Trevor Darrell and Ion Stoica},
 title = {{sensAI: ConvNets Decomposition via Class Parallelism for Fast Inference on Live Data}},
 booktitle = {Proceedings of Fourth Conference on Machine Learning and Systems (MLSys'21)},
 year = {2021}
}