Repository files navigation

A Keras (branchtf2.2 supports TensorFlow 2) implementation of CapsNet in the paper:
Sara Sabour, Nicholas Frosst, Geoffrey E Hinton. Dynamic Routing Between Capsules. NIPS 2017
The currentaverage test error = 0.34% andbest test error = 0.30%.

Differences with the paper:

• We use the learning rate decay withdecay factor = 0.9 andstep = 1 epoch,
  while the paper did not give the detailed parameters (or they didn't use it?).
• We only report the test errors after50 epochs training.
  In the paper, I suppose they trained for1250 epochs according to Figure A.1?Sounds crazy, maybe I misunderstood.
• We use MSE (mean squared error) as the reconstruction loss andthe coefficient for the loss islam_recon=0.0005*784=0.392.
  This should beequivalent with using SSE (sum squared error) andlam_recon=0.0005 as in the paper.

Please use Keras==2.0.7 with TensorFlow==1.2 backend, or theK.batch_dot function may not work correctly.

However, if you useTensorflow>=2.0, then checkout branchtf2.2

Step 1. Clone this repository to local.

git clone https://github.com/XifengGuo/CapsNet-Keras.git capsnet-kerascd capsnet-kerasgit checkout tf2.2 # Only if use Tensorflow>=2.0

Step 2.Install Keras==2.0.7 with TensorFlow==1.2 backend.

pip install tensorflow-gpu==1.2pip install keras==2.0.7

or install Tensorflow>=2.0

pip install tensorflow==2.2

Step 3. Train a CapsNet on MNIST

Training with default settings:

python capsulenet.py

More detailed usage run for help:

python capsulenet.py -h

Step 4. Test a pre-trained CapsNet model

Suppose you have trained a model using the above command, then the trained model will besaved toresult/trained_model.h5. Now just launch the following command to get test results.

$ python capsulenet.py -t -w result/trained_model.h5

It will output the testing accuracy and show the reconstructed images.The testing data is same as the validation data. It will be easy to test on new data,just change the code as you want.

You can also justdownload a model I trained fromhttps://pan.baidu.com/s/1sldqQo1orhttps://drive.google.com/open?id=1A7pRxH7iWzYZekzr-O0nrwqdUUpUpkik

Step 5. Train on multi gpus

This requiresKeras>=2.0.9. After updating Keras:

python capsulenet-multi-gpu.py --gpus 2

It will automatically train on multi gpus for 50 epochs and then output the performance on test dataset.But during training, no validation accuracy is reported.

CapsNet classification testerror on MNIST. Average and standard deviation results arereported by 3 trials. The results can be reproduced by launching the following commands.

python capsulenet.py --routings 1 --lam_recon 0.0    #CapsNet-v1   python capsulenet.py --routings 1 --lam_recon 0.392  #CapsNet-v2python capsulenet.py --routings 3 --lam_recon 0.0    #CapsNet-v3 python capsulenet.py --routings 3 --lam_recon 0.392  #CapsNet-v4

Losses and accuracies:

About100s / epoch on a single GTX 1070 GPU.
About80s / epoch on a single GTX 1080Ti GPU.
About55s / epoch on two GTX 1080Ti GPU by usingcapsulenet-multi-gpu.py.

The result of CapsNet-v4 by launching

python capsulenet.py -t -w result/trained_model.h5

Digits at top 5 rows are real images from MNIST anddigits at bottom are corresponding reconstructed images.

python capsulenet.py -t --digit 5 -w result/trained_model.h5

For each digit, theith row corresponds to theith dimension of the capsule, and columns from left toright correspond to adding[-0.25, -0.2, -0.15, -0.1, -0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25] tothe value of one dimension of the capsule.

As we can see, each dimension has caught some characteristics of a digit. The same dimension ofdifferent digit capsules may represent different characteristics. This is because that differentdigits are reconstructed from different feature vectors (digit capsules). These vectors are mutuallyindependent during reconstruction.

• PyTorch:

  • XifengGuo/CapsNet-Pytorch
  • timomernick/pytorch-capsule
  • gram-ai/capsule-networks
  • nishnik/CapsNet-PyTorch
  • leftthomas/CapsNet

• TensorFlow:

  • naturomics/CapsNet-Tensorflow
    I referred to some functions in this repository.
  • InnerPeace-Wu/CapsNet-tensorflow
  • chrislybaer/capsules-tensorflow

• MXNet:

  • AaronLeong/CapsNet_Mxnet

• Chainer:

  • soskek/dynamic_routing_between_capsules

• Matlab:

  • yechengxi/LightCapsNet