Efficient-CapsNet on Pytorch

Efficient-CapsNet (Mazzia et al., 2021) is a novel architecture for Capsule Networks that improves the routing algorithm and reduces the number of parameters.

Schematic representation of the overall architecture of Efficient-CapsNet.

This project is a faithful PyTorch implementation of the paper with additional features such as logging and monitoring with tensorboard, and a configuration file for easy hyperparameter tuning. The code is based on authors' original implementation in Tensorflow here, and has been tested to match it numerically.

Installation

Python 3 dependencies

• pyyaml
• torch
• torchvision
• opencv-python
• pandas
• tensorboard

We recommend using a virtual environment to install the required packages, such as conda.

git clone [email protected]:akhdanfadh/efficient-capsnet-pytorch.git
cd efficient-capsnet-pytorch

conda create -n efficient-capsnet python=3.10
conda activate efficient-capsnet
conda install pytorch torchvision pytorch-cuda=12.1 -c pytorch -c nvidia
conda install tensorboard pyyaml pandas opencv

This project has been tested on WSL Ubuntu 22.04 with PyTorch 2.2 & CUDA 12.2 on a 3090.

Usage

Modify the config.yaml file to your needs and run the training script as follows. Please check the config file for the available options.

python train.py -c config.yaml -i run_id

For monitoring the training process, we use tensorboard. To start tensorboard, run the following command after training or in a separate terminal for live-monitoring.

tensorboard --logdir=saved

License

This project is licensed under the MIT License.

Acknowledgements

Code structure and training loop are based on the pytorch-template repository with lots of adjustments. The repo help us to understand organizing a deep learning project thoroughly.

Shoutout also to mCoding video on "Modern Python logging" that explains thoroughly how it works and best practices.

Citation

Kudos to the authors of the paper for their amazing work. If you find this code useful, please consider citing the original work:

@article{mazzia2021efficient,
    title={Efficient-CapsNet: capsule network with self-attention routing},
    author={Mazzia, Vittorio and Salvetti, Francesco and Chiaberge, Marcello},
    year={2021},
    journal={Scientific reports},
    publisher={Nature Publishing Group},
    volume={11}
}