A Pytorch implementation of "Matrix Capsules with EM routing"
This is a pytorch implementation of Matrix Capsules with EM routing
Capsules.py, there are two implemented classes:
ConvCaps. The ClassCapsules in the paper is actually a special case of
ConvCapswith whole receptive field, transformation matrix sharing and Coordinate Addition.
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.
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
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.
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.
foriteration in E-step of