/stacked_capsule_autoencode_pl

implement stacked capsule autoencode with pytorch lightning

Primary LanguagePython

stacked_capsule_autoencode_pl

This code refers to the official google-research repository.

This is a pytorch_lightning (pytorch) implemention of the Stacked Capsule Autoencoder (SCAE), which was introduced in the following paper:A. R. Kosiorek, Sara Sabour, Y. W. Teh, and Geoffrey E. Hinton, "Stacked Capsule Autoencoders".

  • Note that this code still exist bugs, we will fix it in the future.

Our environment

  • 'torch'==1.16
  • 'pytorch_lightning'==0.9.0
  • 'tensorboard'==1.15.2
  • 'numpy'==1.17

Understanding the Code

  • The training loop is defined in main.py.
  • The model is built in model/SCAE.py.
  • Using tensorboard --logdir=./lightning_log to monitor the procedure of training.