A PyTorch implementation of CapsNet based on NIPS 2017 paper Dynamic Routing Between Capsules.
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
pip install tqdm
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.
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.
The reconstructions of the digit numbers are showed at right and the ground truth at left.
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.
Primarily referenced this implementation: PyTorch implementation by @Gram.AI