A PyTorch implementation of CapsNet based on Geoffrey Hinton's paper Dynamic Routing Between Capsules
- Anaconda
- PyTorch
conda install pytorch torchvision -c soumith
conda install pytorch torchvision cuda80 -c soumith # install it if you have installed cuda
- PyTorchNet
pip install git+https://github.com/pytorch/tnt.git@master
- tqdm
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 confusion matrix of the digit numbers are showed 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