Towards Compact CNNs via Collaborative Compression

January 24, 2021 ยท View on GitHub

PyTorch implementation for Towards Compact CNNs via Collaborative Compression.

Running Code

In this code, you can run ResNet/DenseNet/VGGNet/GoogLeNet model on CIFAR10/ImageNet2012 dataset. The code has been tested by Python 3.6.8, Pytorch 1.7.1 and CUDA 10.0.

Running Example

Train

CIFAR-10

ResNet-56 (compression ratio = 50%)
python compress.py  --dataset cifar10 \
                    --net resnet56 \
                    --pretrained True \
                    --checkpoint pth/resnet56.pth \
                    --train_dir tmp/resnet56_CC_0.5 \
                    --train_batch_size 128 \
                    --com_ratio 0.5

or

sh compress.sh
DenseNet-40 (compression ratio = 50%)
python compress.py  --dataset cifar10 \
                    --net densenet40 \
                    --pretrained True \
                    --checkpoint pth/densenet40.pth \
                    --train_dir tmp/densenet40_CC_0.5 \
                    --train_batch_size 128 \
                    --com_ratio 0.5

ImageNet-2012

Setting ImageNet-2012 directory in dataset/imagenet.py

ResNet-50 (compression ratio = 50%)
python compress.py  --dataset imagenet \
                    --net resnet50 \
                    --pretrained True \
                    --checkpoint pth/resnet50.pth \ # download from torchvision
                    --train_dir tmp/resnet50_CC_0.5 \
                    --train_batch_size 256 \
                    --com_ratio 0.5