dgl/examples/pytorch/capsule

DGL implementation of Capsule Network

This repo implements Hinton and his team's Capsule Network. Only margin loss is implemented, for simplicity to understand the DGL.

Dependencies

• PyTorch 0.4.1+
• torchvision

pip install torch torchvision

Training & Evaluation

# Run with default config
python3 main.py
# Run with train and test batch size 128, and for 50 epochs
python3 main.py --batch-size 128 --test-batch-size 128 --epochs 50