Pytorch
Dropout
Neural Networks
Machine Learning
Deep Learning

How to implement dropout in Pytorch, and where to apply it

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

In PyTorch, dropout is implemented as a layer that randomly zeroes activations during training and automatically turns itself off during evaluation. The practical question is not just how to call nn.Dropout, but where it belongs in the model so it regularizes useful hidden representations without destroying the signal you actually want to learn.

Basic dropout in PyTorch

The standard module is nn.Dropout(p), where p is the probability of dropping each activation:

python
1import torch
2import torch.nn as nn
3
4class MLP(nn.Module):
5    def __init__(self):
6        super().__init__()
7        self.net = nn.Sequential(
8            nn.Linear(100, 256),
9            nn.ReLU(),
10            nn.Dropout(p=0.5),
11            nn.Linear(256, 64),
12            nn.ReLU(),
13            nn.Dropout(p=0.3),
14            nn.Linear(64, 10),
15        )
16
17    def forward(self, x):
18        return self.net(x)

This is the usual pattern for fully connected networks: apply dropout after an activation in hidden layers.

Train mode versus eval mode matters

PyTorch enables dropout only when the model is in training mode:

python
1model = MLP()
2model.train()   # dropout active
3
4model.eval()    # dropout disabled

This is critical. If you evaluate with model.train() still active, predictions become noisy and metrics are misleading. Conversely, if you forget to set training mode during training after custom evaluation steps, dropout will not regularize anything.

Where dropout usually belongs

The common rule is:

  • use dropout on hidden layers
  • use little or none on the input
  • usually do not put dropout on the final output layer

Why? Hidden layers learn internal co-adaptations that dropout is meant to break up. The output layer already has a direct task-specific meaning, so randomly deleting output logits usually hurts more than it helps.

For dense networks, the standard placement is:

text
Linear -> ReLU -> Dropout

or the equivalent with another activation.

Dropout in convolutional models

For convolutional networks, you can still use ordinary dropout, but spatial variants such as nn.Dropout2d or nn.Dropout3d are often more appropriate because they drop whole feature maps or channels in a way that better matches convolutional structure.

python
1class ConvNet(nn.Module):
2    def __init__(self):
3        super().__init__()
4        self.features = nn.Sequential(
5            nn.Conv2d(3, 32, kernel_size=3, padding=1),
6            nn.ReLU(),
7            nn.Dropout2d(p=0.2),
8            nn.Conv2d(32, 64, kernel_size=3, padding=1),
9            nn.ReLU(),
10            nn.Dropout2d(p=0.3),
11        )
12        self.classifier = nn.Linear(64 * 32 * 32, 10)
13
14    def forward(self, x):
15        x = self.features(x)
16        x = torch.flatten(x, 1)
17        return self.classifier(x)

In many CNNs, dropout is used more sparingly than in old fully connected networks, especially when batch normalization and data augmentation already regularize the model.

How much dropout to use

Typical starting values are:

  • '0.1 to 0.3 for mild regularization'
  • around 0.5 for stronger regularization in dense layers

There is no universal best number. Too little may have no effect. Too much can cause underfitting by destroying too much information every step. Validation performance should decide.

Common Pitfalls

The most common mistake is applying dropout during evaluation by forgetting model.eval(). That makes inference inconsistent and harder to reproduce.

Another mistake is placing dropout on the final output layer. In most classification and regression models, that weakens the task signal instead of regularizing useful hidden features.

Developers also overuse dropout everywhere, including tiny models or layers that already have limited capacity. If the model begins to underfit badly, dropout may be too aggressive.

Finally, do not expect dropout to replace all other regularization. Weight decay, data augmentation, architecture choice, and batch normalization still matter.

Summary

  • In PyTorch, use nn.Dropout or its spatial variants as regularization layers.
  • Apply dropout mainly to hidden representations, often after activations.
  • Keep model.train() and model.eval() correct, because dropout behavior depends on mode.
  • Use moderate dropout probabilities and tune them with validation results.
  • Avoid placing dropout blindly on inputs or final outputs unless you have a specific reason.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.