Matrix-Capsules-pytorch

by YuxianMeng

A Pytorch implementation of "Matrix Capsules with EM routing"

207 Stars 40 Forks Last release: Not found 53 Commits 0 Releases

Available items

No Items, yet!

The developer of this repository has not created any items for sale yet. Need a bug fixed? Help with integration? A different license? Create a request here:

Matrix-Capsules-pytorch

This is a pytorch implementation of Matrix Capsules with EM routing

In

Capsules.py
, there are two implemented classes:
PrimaryCaps
and
ConvCaps
. The ClassCapsules in the paper is actually a special case of
ConvCaps
with whole receptive field, transformation matrix sharing and Coordinate Addition.

In

train.py
, I define a CapsNet in the paper using classes in
Capsules.py
, and could be used to train a model for MNIST dataset.

Train a small CapsNet on MNIST

python train.py -batch_size=64 -lr=2e-2 -num_epochs=5 -r=1 -print_freq=5
.

Note: * more args can be found in

utils.py
, and if you want to change A,B,C,D, go to
line 62
of
train.py
* m and lambda schedule need to be changed if you want to train a capsnet with r=2 or 3. The default schedule make capsnet does not converge in those cases.

Results

The test accuracy is around 97.6% after 1 epoch and 98.7% after 2 epochs of training with a small Capsule of A,B,C,D,r = 64,8,16,16,1. After 30 epochs of training, the best acc is around 99.3%. More results on different configurations are welcomed.

TODO

  • using more matrix operation rather than
    for
    iteration in E-step of
    Capsules.py
    .
  • make capsules work when heightin != widthin
  • find better lambda/m schedule to speed up the convergence.

We use cookies. If you continue to browse the site, you agree to the use of cookies. For more information on our use of cookies please see our Privacy Policy.