PyTorch implementation of Stacked Capsule Auto-Encoders

Hi,
I have implemented Stacked Capsule Auto-Encoder (Kosiorek et al, 2019) in PyTorch. The original implementation by the authors of paper was created with TensorFlow v1 and DeepMind Sonnet. I could not find a good PyTorch implementation, so hope this provides an easier-to-understand alternative.

I mostly kept the architecture of model and hyper-parameters same with the original implementation. However, some parts are refactored for ease of use.

I am kind of new to PyTorch, so there might be some not-so-best practices :slight_smile:. Feel free to give any feedback.

3 Likes

Thanks for sharing the code, Baris! :slight_smile:

Thanks Baris! :pray: :pray: