Home

Awesome

CapsNet

A PyTorch implementation of CapsNet based on NIPS 2017 paper Dynamic Routing Between Capsules.

Requirements

conda install pytorch torchvision -c soumith
conda install pytorch torchvision cuda80 -c soumith # install it if you have installed cuda

pip install git+https://github.com/pytorch/tnt.git@master

Usage

git clone https://github.com/leftthomas/CapsNet.git
cd CapsNet
python -m visdom.server & python main.py

Visdom now can be accessed by going to 127.0.0.1:8097 in your browser, or your own host address if specified.

Benchmarks

Highest accuracy was 99.57% after 30 epochs. The model may achieve a higher accuracy as shown by the trend of the loss/accuracy graphs below.

<table> <tr> <td> <img src="results/train_loss.png"/> </td> <td> <img src="results/test_loss.png"/> </td> </tr> </table> <table> <tr> <td> <img src="results/train_acc.png"/> </td> <td> <img src="results/test_acc.png"/> </td> </tr> </table>

The confusion matrix of the digit numbers are showed below. <img src="results/confusion_matrix.png"/>

The reconstructions of the digit numbers are showed at right and the ground truth at left.

<table> <tr> <td> <img src="results/ground_truth.jpg"/> </td> <td> <img src="results/reconstruction.jpg"/> </td> </tr> </table>

Default PyTorch Adam optimizer hyperparameters were used with no learning rate scheduling. Epochs with batch size of 100 takes ~2 minutes on a NVIDIA GTX 1070 GPU.

Other Implementations

Credits

Primarily referenced this implementation: PyTorch implementation by @Gram.AI