Dropout2d

Contents

Dropout2d#

class brainstate.nn.Dropout2d(prob=0.5, channel_axis=-1, name=None)#

Randomly zero out entire channels (a channel is a 2D feature map).

Each channel will be zeroed out independently on every forward call with probability using samples from a Bernoulli distribution. The channel is a 2D feature map, e.g., the \(j\)-th channel of the \(i\)-th sample in the batched input is a 2D tensor \(\text{input}[i, j]\).

Usually the input comes from Conv2d modules.

As described in the paper [1], if adjacent pixels within feature maps are strongly correlated (as is normally the case in early convolution layers) then i.i.d. dropout will not regularize the activations and will otherwise just result in an effective learning rate decrease.

In this case, Dropout2d will help promote independence between feature maps and should be used instead.

Parameters:

  • prob (float) – Probability of an element to be kept. Default is 0.5.

  • channel_axis (int) – The axis representing the channel dimension. Default is -1.

  • name (str | None) – The name of the dynamic system.

Notes

Input shape: \((N, C, H, W)\) or \((C, H, W)\).

Output shape: \((N, C, H, W)\) or \((C, H, W)\) (same shape as input).

References

Examples

>>> import brainstate
>>> m = brainstate.nn.Dropout2d(prob=0.8)
>>> x = brainstate.random.randn(20, 32, 32, 16)
>>> with brainstate.environ.context(fit=True):
...     output = m(x)
>>> output.shape
(20, 32, 32, 16)
Contents