Add Support for Stochastic Depth

lucidrains / g-mlp-pytorch

Implementation of gMLP, an all-MLP replacement for Transformers, in Pytorch

MIT License

422 stars 56 forks source link

Add Support for Stochastic Depth #1

Closed mlw214 closed 3 years ago

mlw214 commented 3 years ago

This PR adds support for stochastic depth, which is used in the paper for the vision experiments. I went ahead an added it to gMLP as well for completeness.

I tried my best to match your style. Let me know if there are any problems, or if you want me to refactor anything.

lucidrains commented 3 years ago

@mlw214 Hi Miller! Thanks for the PR! I realized the way I had things setup didn't make it easy for people to change the survival probability on a schedule, so I rewrote a bunch of code. Thank you for bringing this to my attention anyways!

mlw214 commented 3 years ago

Roger that! The changes you added in addition to this are great! Looking forward to playing around with this more