paganpasta / eqxvision

A Python package of computer vision models for the Equinox ecosystem.
https://eqxvision.readthedocs.io
MIT License
97 stars 11 forks source link

scale activations in `DropPath` #66

Closed eringrant closed 1 year ago

eringrant commented 1 year ago

Great project! An additional test + a small fix to align DropPath with standard dropout and torchvision.stochastic_depth.

paganpasta commented 1 year ago

Hey, thanks for the PR.

I was thinking maybe the mode values change to 'batch, rows' will also be beneficial. What do you think?

eringrant commented 1 year ago

For the longer term, I think it makes sense to make equinox.nn.dropout upstream more expressive to handle things like mode = global (as here), or tiled dropout (as handled here, for example). (But of course, this repo is not the place to propose changes to another repo. :)

For the shorter term, I prefer an API that refers to axes by index rather than by their expected semantics. So for example, rather than expecting rows to be meaningful for every input value to eqxvision.layers.DropPath.__call__, I prefer to identify by index the axes over which to broadcast the dropout mask (similarly to the tiled dropout implementation above).

paganpasta commented 1 year ago

That makes sense. It might be beneficial to raise an issue at Equinox for this extension.

For the shorter term, I prefer an API that refers to axes by index rather than by their expected semantics...

Agreed. I'd be happy to accept a PR on this.