Pytorch version of Hinton's Capsule Theory paper: Dynamic Routing Between Capsules
Pytorch version of Hinton's paper: Dynamic Routing Between Capsules
Some implementations of CapsNet online have potential problems and it's uneasy to realize the bugs since MNIST is too simple to achieve satisfying accuracy.
Corresponding pipeline: Input > Conv1 > Caps(cnn inside) > Route > Loss
Caps
layer, by re-writing the function create_cell_fn
you can implement your own sub-network inside Caps
Layer def create_cell_fn(self):
"""
create sub-network inside a capsule.
:return:
"""
conv1 = nn.Conv2d(self.conv1_kernel_num, self.caps1_conv_kernel_num, kernel_size = self.caps1_conv_kernel_size, stride = self.caps1_conv1_stride)
#relu = nn.ReLU(inplace = True)
#net = nn.Sequential(conv1, relu)
return conv1
Route
, you can take use of Caps
Layer and Route
Layer to construct any type of networkRoute
layer.I'm a Research Assistant @ National University of Singapre, before joinging NUS, I was a first-year PhD candidate in Zhejiang University and then quitted. Contact me with email: [email protected] or wechat: dragen1860
conda install pytorch torchvision cuda80 -c soumith
git clone https://github.com/dragen1860/CapsNet-Pytorch.git cd CapsNet-Pytorch
glo_batch_size = 125
to appropriate size according to your GPU memory size.$ python main.py
$ tensorboard --logdir runs
OR you can comment the part of train code and test its performance with pretrained model mdl
file.
Model | Routing | Reconstruction | MNIST | |
---|---|---|---|---|
Baseline | - | - | 0.39 | |
Paper | 3 | no | 0.35 | |
Ours | 3 | no | 0.34 |
It takes about 150s per epoch for single GTX 970 4GB Card.
Keras:
TensorFlow:
MXNet:
Lasagne (Theano):
Chainer: