Dropout & BatchNorm
Theme
Palette
Download
Caption (tap to copy)
📖 Deep dive (full written explanation)
This is the hands-on post, and it is intentionally code-heavy. The previous posts built understanding from concept through mechanics; here you assemble a complete, runnable network that uses both BatchNorm and Dropout, train it, and — crucially — manage the train/eval switch the layers depend on.
Each snippet is real code that runs in order. The most valuable thing you can do is run it, then deliberately toggle the layers off or move them around to feel their effect on accuracy and stability. That experimentation is how the abstract behavior from earlier posts becomes intuition.
The model definition shows the canonical placement of both layers. A linear layer feeds BatchNorm1d, whose output goes through a ReLU, then Dropout, then a final linear classifier. In the forward method you can read the ordering directly: bn1 wraps fc1's output before the activation, and drop comes after the activation.
This ordering — Linear, BatchNorm, activation, Dropout — is the safe default established in earlier posts, and seeing it expressed as actual module attributes and a forward pass makes it concrete. Note that BatchNorm and Dropout are declared once in __init__ and reused in forward, exactly like any other layer.
The training loop's most important line is the easiest to overlook: model.train(). Calling it at the start of each epoch puts both layers into training mode — Dropout actively zeros neurons, and BatchNorm uses live batch statistics while updating its running averages. The rest of the loop is the standard forward, loss, zero_grad, backward, step sequence.
Without model.train(), a model that was previously in eval mode would silently skip Dropout and freeze BatchNorm's statistics, undermining the very regularization and normalization you added. Making the train() call an explicit, habitual part of the loop is the cheapest insurance against one of the most common bugs with these layers.
The validation loop mirrors the training loop but with two essential changes: model.eval() and torch.no_grad(). The eval() call flips both layers into inference mode — Dropout passes everything through, and BatchNorm switches to its stored running statistics — so your evaluation is deterministic and reflects how the model will actually behave in deployment.
The torch.no_grad() block is a separate optimization that disables gradient tracking to save memory and time, since you do not backpropagate during validation. It is important not to confuse the two: no_grad() does NOT put the layers into eval mode, and eval() does NOT stop gradient tracking. You generally want both, but for different reasons.
This cycle diagram captures the rhythm of a training run: switch to train mode and turn the noise on, fit a number of batches while updating both the weights and BatchNorm's running statistics, switch to eval mode and turn the noise off, then validate using the frozen running statistics. Then the cycle repeats for the next epoch.
Visualizing it as a loop emphasizes that the mode is something you flip back and forth many times over a training run, not once. A frequent bug is forgetting to switch back to train() after validating, which silently disables Dropout for all subsequent training — the cyclic view is a reminder that both transitions matter.
For image models, the layers have convolutional variants, and this block shows them. BatchNorm2d normalizes per channel across the batch and spatial dimensions, which is the correct behavior for feature maps. Dropout2d zeros entire feature-map channels rather than individual pixels, because adjacent pixels in a channel are highly correlated and dropping single pixels does little.
The pattern Conv2d, BatchNorm2d, ReLU, Dropout2d is the direct image-domain analogue of the MLP block. Matching the dimensionality of the normalization and dropout to your tensor rank is essential — using a 1d layer on 4d image tensors, or vice versa, is a common source of shape errors.
These shape-and-switch gotchas are where most practitioners lose time. BatchNorm1d expects 2d input of (N, features) while BatchNorm2d expects 4d (N, C, H, W); using the wrong one raises a shape error. You must call model.eval() before validating or your metrics are meaningless. Dropout2d zeros whole feature maps, not single pixels. BatchNorm needs more than one example per batch in training mode or it cannot compute a variance.
The final point is the most commonly conflated: torch.no_grad() and model.eval() do different things. The first disables gradient tracking; the second changes layer behavior. Wrapping validation in no_grad() while forgetting eval() leaves Dropout and BatchNorm in training mode, quietly corrupting your evaluation even though no error appears.
This snippet pulls back the curtain on what BatchNorm stores. running_mean and running_var hold the exponential moving averages of the per-feature statistics, each with one value per feature (256 here), and they are the numbers used at eval time. weight and bias are gamma and beta — the learned scale and shift.
Being able to inspect these tensors is genuinely useful for debugging: if running statistics look wildly off, your batches may be unrepresentative or too small; if gamma or beta have collapsed, something upstream may be wrong. It also reinforces the conceptual split from the mechanics post — running stats are tracked by averaging, while gamma and beta are learned by gradient descent.
The data-flow diagram traces tensor shapes through the block, which is the single most useful thing to hold in your head when wiring these layers. A batch of shape (N, 256) passes through the linear layer, then BatchNorm1d and ReLU and Dropout all preserving (N, 256), and finally the classifier maps it to logits of shape (N, 10).
Most wiring bugs with BatchNorm and Dropout are shape or rank bugs, so being able to recite this chain lets you catch a mismatch by reasoning rather than trial and error. When something breaks, mentally walking these shapes usually reveals whether you reached for the wrong BatchNorm dimensionality.
The order gotcha is worth a dedicated slide because it causes silent damage. If Dropout runs before BatchNorm, the zeros it injects distort the mean and variance that BatchNorm then computes, so BatchNorm normalizes against corrupted statistics — hurting both training stability and the train/test match. The safe, standard ordering is Linear or Conv, then BatchNorm, then activation, then Dropout.
The second half of the slide reiterates the eval() trap because it is the most damaging mistake of all: leaving Dropout active and BatchNorm on noisy batch statistics at test time silently tanks accuracy. Both of these are quiet failures — no exception, just worse results — which is exactly why they deserve explicit attention before the dedicated mistakes post.
This recap consolidates the build: place the layers as Linear/Conv, then BatchNorm, then activation, then Dropout; call model.train() before fitting and model.eval() before validating; match BatchNorm1d or 2d to your tensor rank; use Dropout2d for convolutional feature maps; and find BatchNorm's running statistics in running_mean and running_var.
With a working, correctly-switched model in hand, the only thing left is to learn the traps — the quiet mistakes that make these layers hurt instead of help. That is the focus of the final post.
The teaser points to the common-mistakes post. You now have a network that trains and validates correctly; the next step is hardening your instincts by learning the quiet failure modes — forgetting eval mode, over-dropping, wrong layer order, BatchNorm on tiny batches, and over-regularizing — along with the one-line fix for each.