censomin / capsule-network

PyTorch implementation of the paper Dynamic Routing Between Capsules by Sara Sabour, Nicholas Frosst and Geoffrey Hinton

Capsule Network

PyTorch implementation of the following paper:

• Dynamic Routing Between Capsules by Sara Sabour, Nicholas Frosst and Geoffrey Hinton

Official implemenation

• Official implementation (TensorFlow) by Sara Sabour

Visual represenation

Image source: Mike Ross, A Visual Representation of Capsule Network Computations

Run the experiment

• For details, run python main.py --help

Example of reconstructed vs. original images

---

Requirements:

• PyTorch (http://www.pytorch.org)
• NumPy (http://www.numpy.org/)
• GPU

Default hyper-parameters (similar to the paper):

• Per-GPU batch_size = 128
• Initial learning_rate = 0.001
• Exponential lr_decay = 0.96
• Number of routing iteration (num_routing) = 3

Loss function hyper-parameters (see loss.py):

• Lambda for Margin Loss = 0.5
• Scaling factor for reconstruction loss = 0.0005

GPU Speed benchmarks:

(with above mentioned hyper-parameters)

• Single GeForce GTX 1080Ti - 35.6s per epoch
• Two GeForce GTX 1080Ti - 35.8s per epoch (twice the batch size -> half the iteration)