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