PyTorch implementation of Stacked Capsule Auto-Encoders

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.


Thanks for sharing the code, Baris! :slight_smile:

Thanks Baris! :pray: :pray:

Beautiful code, like a poem. But the performances seems not comparable to the reported results in the paper.

Yes, you’re right. There must be a bug somewhere causing worse performance, but I haven’t had the opportunity to spend some time to find it yet.

I reimplement it according to the official code, but I only get .6 accuracy also.