Implement noise_shape keyword for Dropout layers

[around1991]

🚀 Feature

Tensorflow has a noise_shape keyword for tf.nn.dropout, which specifies which dimensions should have dropout masks calculated independently and which dimensions should have shared dropout masks. PyTorch should have a similar feature too.

Motivation

This is a very useful feature to have (for example, if we process the same example multiple times, we may want to tie dropout masks for each time we see that example), and while easy to implement independently, would be good to have common functionality for.

Pitch

Include extra kwarg to torch.nn.Dropout called noise_shape, with the same functionality as tf.nn.dropout.

Alternatives

Fairly easy to implement independently, but would save hassle for many people if there was a core implementation.

[zhangguanheng66]

@soumith I have a look at this feature request. If we decide to implement it, I'm happy to take this small task.

[soumith]

i think this is worth doing, but the first step to hook it up to is F.dropout in nn/functional.py
It ends up calling _VF.dropout with probability p. We should give a noise_shape to _VF.dropout which by default is the same shape as input, but otherwise it'll be a broadcastable shape.

For example:

inp = torch.randn(10, 3, 224, 224)
noise_shape = torch.Size([1, 3, 1, 1])

In here, the same random mask will be broadcasted to shape of 10, 3, 224, 224 where, masks are shared between all images across all pixels, but per-channel.

[visionscaper]

I was curious to know if anyone picked this up?

[soumith]

@visionscaper nope

[shihongzhi]

I'd like to pick this up.

[visionscaper]

Sounds great @shihongzhi!

[alanhdu]

Hi! This is a papercut that we would also like to see solved. Is there anything we can do to speed this along?

cc @cxhernandez