EFFICIENT-CAPSNET: CAPSULE NETWORK WITH SELF-ATTENTION ROUTING

Tensorflow 2.x (with Keras API) Implementation of the Efficient-CapsNet (Mazzia et al., 2021)

Training and Testing (MNIST)

Use

python main.py --checkpoint_dir=checkpoint --num_epochs=15

• checkpoint: path to save trained weights of the model.
• num_epochs: Number of epochs. You can resume any previous training.

You can also simplify

python main.py --flagfile=./flags.txt

Environment (Docker)

If you use Docker, the code has been verified to work on tensorflow/tensorflow:2.6.0-gpu.

• Pull Image & Run container & Get into shell

docker run -it --rm --runtime=nvidia tensorflow/tensorflow:2.6.0-gpu bash

• (In container) Install and Run code

# Install git
apt-get update
apt-get install git
git clone https://github.com/kaparoo/Efficient-CapsNet.git
cd Efficient-CapsNet

# Install matplotlib
pip install matplotlib

# Training
python main.py --checkpoint_dir=checkpoint --num_epochs=15

References

• Efficient-CapsNet: Capsule Network with Self-Attention Routing (arXiv)
• Official Code (GitHub)