Dropout Variants

Idea: Randomly zero parts of the network during training to prevent overfitting.

Equation

With keep mask $m\sim \mathrm{Bernoulli}(1-p)$, vanilla dropout gives $\mathbb{E}[m\odot x]=(1-p)x$, so inference must multiply by $(1-p)$. Inverted dropout (what every framework implements) rescales at train time instead, leaving inference as the identity:

$$y_{\text{train}}=\frac{m\odot x}{1-p},\mathbb{E}[y_{\text{train}}]=x,y_{\text{eval}}=x$$

DropPath / stochastic depth samples one mask per sample and applies it to a whole residual branch, so the block becomes $x+\mathrm{drop}(F(x))$ and sometimes reduces to $x$:

$$y=x+\frac{m_b}{1-p}F(x),m_b\sim \mathrm{Bernoulli}(1-p)\text{per batch element}$$

From scratch

def dropout(x, p=0.5, training=True):
    if not training or p == 0.0:
        return x
    mask = (torch.rand_like(x) >= p).to(x.dtype)
    return x * mask / (1.0 - p)                    # inverted
 
def drop_path(x, p=0.1, training=True):
    if not training or p == 0.0:
        return x
    shape = (x.shape[0],) + (1,) * (x.ndim - 1)    # one Bernoulli per sample
    mask = (torch.rand(shape, device=x.device, dtype=x.dtype) >= p)
    return x * mask / (1.0 - p)
 
def dropout2d(x, p=0.1, training=True):            # drop whole channels (N,C,H,W)
    if not training or p == 0.0:
        return x
    mask = (torch.rand(x.shape[:2], device=x.device) >= p)[:, :, None, None]
    return x * mask.to(x.dtype) / (1.0 - p)