dmlc--dgl
Hongzhi (Steve), Chen
23d09057e3
* [Misc] Black auto fix. * sort Co-authored-by: Steve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
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