Capsule Routing via Variational Bayes (AAAI 2020)
[Official Pytorch implementation]
Examplary code for our AAAI 2020 paper on capsule networks.
Author: Fabio De Sousa Ribeiro
E-mail: [email protected]
Overview
Modular vb-routing and conv capsule layers so you can stack them to build your own capsnet to play around with.
self.Conv_1 = nn.Conv2d(in_channels=2, out_channels=64,
kernel_size=5, stride=2)
self.PrimaryCaps = PrimaryCapsules2d(in_channels=64, out_caps=16,
kernel_size=3, stride=2, pose_dim=4)
self.ConvCaps = ConvCapsules2d(in_caps=16, out_caps=5,
kernel_size=3, stride=1, pose_dim=4)
self.Routing = VariationalBayesRouting2d(in_caps=16, out_caps=5,
cov='diag', pose_dim=4, iter=3,
alpha0=1., # Dirichlet(pi | alpha0) prior
m0=torch.zeros(4*4), kappa0=1., # Gaussian(mu_j | m0, (kappa0 * Lambda_j)**-1) prior
Psi0=torch.eye(4*4), nu0=4*4+1) # Wishart(Lambda_j | Psi0, nu0) prior97.1% test acc on smallNORB with just 1 caps layer.
98.7% with 3 caps layers (as in paper).
For more see the poster in images/Poster_AAAI2020.pdf.
Run
python src/main.py
Dataset Download
- You can download smallNORB in .npy format and already resized to 48x48 for convenience.
Citation
@inproceedings{ribeiro2020capsule,
title={Capsule Routing via Variational Bayes.},
author={Ribeiro, Fabio De Sousa and Leontidis, Georgios and Kollias, Stefanos D},
booktitle={AAAI},
pages={3749--3756},
year={2020}
}