PyTorch implementation for Towards Efficient Capsule Networks
Accepted at ICIP 2022, Special Session SCENA
Abstract: From the moment Neural Networks dominated the scene for image processing, the computational complexity needed to solve the targeted tasks skyrocketed: against such an unsustainable trend, many strategies have been developed, ambitiously targeting performance's preservation. Promoting sparse topologies, for example, allows the deployment of deep neural networks models on embedded resource-constrained devices. Recently, Capsule Networks were introduced to enhance explainability of a model, where each capsule is an explicit representation of an object or its parts. These models show promising results on toy datasets, but their low scalability prevents deployment on more complex tasks. In this work, we explore sparsity besides capsule representations to improve their computational efficiency by reducing the number of capsules. We show how pruning with Capsule Network achieves high generalization with less memory requirements, computational effort, and inference and training time.
Via Pip/Conda
pip install -r requirements.txt
Via Docker
ToDO
This repository uses Weight & Biases for experiment tracking. Set your account's project and entity names in the src/wandb_project.json file.
To create a Weight & Biases account see (creating a weights and biases account) and the associated quickstart guide.
If you want to train the networks on Tiny ImageNet, download the dataset from Google Drive and unzip in the data/tiny-imagenet-200 folder
Download checkpoints from EagleEye Google Drive and put them in the dump/models/EagleEye/ folder
Configuration files are stored in the in the src/configs/ folder.
For example, to finetune a CapsNet with a ResNet50 backbone with 25% of remaining FLOPs on cifar10 run
cd src/
python main.py --config configs/resnet50/cifar10/resnet50_flops25_cifar10_structured_finetuning.json --output wandb_runs.csv
or with the backbone freezed
python main.py --config configs/resnet50/cifar10/resnet50_flops25_cifar10_structured_freeze.json --output wandb_runs.csv
More arguments can be found in the corresponding config json files.
Once you trained all the networks, you can see the results running
python results.py --wandb_runs wandb_runs.csv
Use this bibtex if you enjoyed this repository and you want to cite it:
@misc{renzulliICIP22,
doi = {10.48550/ARXIV.2208.09203},
url = {https://arxiv.org/abs/2208.09203},
author = {Renzulli, Riccardo and Grangetto, Marco},
keywords = {Computer Vision and Pattern Recognition (cs.CV), FOS: Computer and information sciences, FOS: Computer and information sciences},
title = {Towards Efficient Capsule Networks},
publisher = {arXiv},
year = {2022},
copyright = {Creative Commons Attribution 4.0 International}
}
ICIP bibtex coming soon!
This code is based on the EagleEye implementation
Contact: [email protected]