PyTorch Batch Normalization and Dropout
Batch Normalization and Dropout are two of the most core regularization and stabilization techniques in deep neural network training. The former solves the internal covariate shift problem during training, making deeper networks trainable; the latter prevents overfitting by randomly dropping neurons, enhancing model generalization. The two are usually used together and are standard components of modern neural networks.
1. Batch Normalization
1.1 Basic Principles
During training of deep networks, small changes in the parameters of the previous layer are continuously amplified as layers deepen, causing the input distribution of subsequent layers to keep changing — this phenomenon is calledInternal Covariate Shift. Batch normalization solves this problem by standardizing the output of each layer, forcibly pulling the activation values back to a stable distribution.
Normalization formula:
\[\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}\]
\[y_i = \gamma \hat{x}_i + \beta\]
where:
- \(\mu_B\), \(\sigma_B^2\) are the mean and variance of the current batch
- \(\epsilon\) is a small constant to prevent division by zero (default
1e-5） - \(\gamma\) and \(\beta\) are learnable scaling and shifting parameters that allow the network to decide the final distribution shape on its own
Benefits of batch normalization:
- Allows larger learning rates, speeding up training
- Reduces sensitivity to initialization, making training more stable
- Has a certain regularization effect, reducing reliance on Dropout
- Alleviates vanishing/exploding gradients, making deeper networks trainable
1.2 BatchNorm1d / 2d / 3d
PyTorch provides three versions based on input dimensions, with identical usage, differing only in the shape of data they process.
BatchNorm1d: for fully connected layers / sequence data
Example
import torch
import torch.nn as nn
# Input shape: (N, C) or (N, C, L)
# N=batch_size, C=number of features/channels, L=sequence length
bn1d = nn.BatchNorm1d(
num_features=128, # number of features/channels
eps=1e-5, # small constant to prevent division by zero (default 1e-5)
momentum=0.1, # momentum for moving average (default 0.1)
affine=True, # whether to learn gamma and beta (default True)
track_running_stats=True, # whether to track running mean/variance (default True)
)
# used after fully connected layers
x = torch.randn(32, 128) # (batch=32, features=128)
out = bn1d(x)
print(out.shape) # torch.Size([32, 128])
# sequence data (e.g., after 1D convolution)
x_seq = torch.randn(32, 128, 50) # (batch, channels, seq_len)
out_seq = bn1d(x_seq)
print(out_seq.shape) # torch.Size([32, 128, 50])
import torch.nn as nn
# Input shape: (N, C) or (N, C, L)
# N=batch_size, C=number of features/channels, L=sequence length
bn1d = nn.BatchNorm1d(
num_features=128, # number of features/channels
eps=1e-5, # small constant to prevent division by zero (default 1e-5)
momentum=0.1, # momentum for moving average (default 0.1)
affine=True, # whether to learn gamma and beta (default True)
track_running_stats=True, # whether to track running mean/variance (default True)
)
# used after fully connected layers
x = torch.randn(32, 128) # (batch=32, features=128)
out = bn1d(x)
print(out.shape) # torch.Size([32, 128])
# sequence data (e.g., after 1D convolution)
x_seq = torch.randn(32, 128, 50) # (batch, channels, seq_len)
out_seq = bn1d(x_seq)
print(out_seq.shape) # torch.Size([32, 128, 50])
BatchNorm2d: for CNN image feature maps
Example
# Input shape: (N, C, H, W)
# normalize independently for each channel C
bn2d = nn.BatchNorm2d(num_features=64) # 64 channels
x = torch.randn(32, 64, 28, 28) # (batch, channels, height, width)
out = bn2d(x)
print(out.shape) # torch.Size([32, 64, 28, 28])
# normalize independently for each channel C
bn2d = nn.BatchNorm2d(num_features=64) # 64 channels
x = torch.randn(32, 64, 28, 28) # (batch, channels, height, width)
out = bn2d(x)
print(out.shape) # torch.Size([32, 64, 28, 28])
BatchNorm3d: for 3D convolution (video / medical imaging)
Example
# Input shape: (N, C, D, H, W)
bn3d = nn.BatchNorm3d(num_features=32)
x = torch.randn(4, 32, 16, 32, 32) # (batch, channels, depth, h, w)
out = bn3d(x)
print(out.shape) # torch.Size([4, 32, 16, 32, 32])
bn3d = nn.BatchNorm3d(num_features=32)
x = torch.randn(4, 32, 16, 32, 32) # (batch, channels, depth, h, w)
out = bn3d(x)
print(out.shape) # torch.Size([4, 32, 16, 32, 32])
View learnable parameters
Example
bn = nn.BatchNorm2d(64)
print(f"weight (gamma) shape: {bn.weight.shape}") # torch.Size([64])
print(f"bias (beta) shape: {bn.bias.shape}") # torch.Size([64])
print(f"running_mean shape: {bn.running_mean.shape}") # torch.Size([64])
print(f"running_var shape: {bn.running_var.shape}") # torch.Size([64])
# running_mean / running_var are not trainable parameters
# They are moving-average statistics used during inference
print(f"trainable parameter count: {sum(p.numel() for p in bn.parameters())}") # 128（64×2）
print(f"weight (gamma) shape: {bn.weight.shape}") # torch.Size([64])
print(f"bias (beta) shape: {bn.bias.shape}") # torch.Size([64])
print(f"running_mean shape: {bn.running_mean.shape}") # torch.Size([64])
print(f"running_var shape: {bn.running_var.shape}") # torch.Size([64])
# running_mean / running_var are not trainable parameters
# They are moving-average statistics used during inference
print(f"trainable parameter count: {sum(p.numel() for p in bn.parameters())}") # 128（64×2）
Meaning of the momentum parameter
Example
# PyTorch's momentum is the opposite of the usual understanding:
# running_mean = (1 - momentum) × running_mean + momentum × batch_mean
# i.e., the larger the momentum, the greater the weight of the current batch's statistics
# default 0.1: usually appropriate
# for small batch size, consider tuning it down: momentum=0.01, to avoid unstable statistics
bn = nn.BatchNorm2d(64, momentum=0.01)
# momentum=None: uses cumulative moving average (CMA), suitable for very small batches
bn = nn.BatchNorm2d(64, momentum=None)
# running_mean = (1 - momentum) × running_mean + momentum × batch_mean
# i.e., the larger the momentum, the greater the weight of the current batch's statistics
# default 0.1: usually appropriate
# for small batch size, consider tuning it down: momentum=0.01, to avoid unstable statistics
bn = nn.BatchNorm2d(64, momentum=0.01)
# momentum=None: uses cumulative moving average (CMA), suitable for very small batches
bn = nn.BatchNorm2d(64, momentum=None)
1.3 LayerNorm Layer Normalization
PairNormalize all features of a single sampleIt does not depend on batch, suitable for scenarios with small or variable batch sizes.The standard normalization method for Transformers。
BatchNorm：在 batch 维度计算均值/方差（跨样本，同通道） LayerNorm：在特征维度计算均值/方差（同样本，跨特征）
Example
import torch
import torch.nn as nn
# ── For fully connected / Transformer ────────────────────────
# normalized_shape: normalize over the last few dimensions
ln = nn.LayerNorm(normalized_shape=512)
x = torch.randn(32, 10, 512) # (batch, seq_len, embed_dim)
out = ln(x)
print(out.shape) # torch.Size([32, 10, 512])
# ── Normalize over multiple dimensions ───────────────────────────────
ln_2d = nn.LayerNorm([64, 28, 28]) # normalize over the three dimensions C, H, W
x = torch.randn(8, 64, 28, 28)
out = ln_2d(x)
# ── Typical usage in Transformer ──────────────────────
class TransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads, ff_dim):
super().__init__()
self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(embed_dim, ff_dim),
nn.GELU(),
nn.Linear(ff_dim, embed_dim),
)
self.norm1 = nn.LayerNorm(embed_dim) # ← normalize after each sublayer
self.norm2 = nn.LayerNorm(embed_dim)
def forward(self, x):
attn_out, _ = self.attn(x, x, x)
x = self.norm1(x + attn_out) # residual + normalization
x = self.norm2(x + self.ff(x))
return x
import torch.nn as nn
# ── For fully connected / Transformer ────────────────────────
# normalized_shape: normalize over the last few dimensions
ln = nn.LayerNorm(normalized_shape=512)
x = torch.randn(32, 10, 512) # (batch, seq_len, embed_dim)
out = ln(x)
print(out.shape) # torch.Size([32, 10, 512])
# ── Normalize over multiple dimensions ───────────────────────────────
ln_2d = nn.LayerNorm([64, 28, 28]) # normalize over the three dimensions C, H, W
x = torch.randn(8, 64, 28, 28)
out = ln_2d(x)
# ── Typical usage in Transformer ──────────────────────
class TransformerBlock(nn.Module):
def __init__(self, embed_dim, num_heads, ff_dim):
super().__init__()
self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
self.ff = nn.Sequential(
nn.Linear(embed_dim, ff_dim),
nn.GELU(),
nn.Linear(ff_dim, embed_dim),
)
self.norm1 = nn.LayerNorm(embed_dim) # ← normalize after each sublayer
self.norm2 = nn.LayerNorm(embed_dim)
def forward(self, x):
attn_out, _ = self.attn(x, x, x)
x = self.norm1(x + attn_out) # residual + normalization
x = self.norm2(x + self.ff(x))
return x
1.4 GroupNorm Group Normalization
Divide the channels into several groupsgroup, and normalize within each group. It does not depend on batch size and outperforms BatchNorm in small-batch scenarios (e.g., object detection, medical image segmentation).
Example
# num_groups: how many groups to divide the channels into (must divide num_channels evenly)
# Special case: num_groups=1 is equivalent to LayerNorm
# num_groups=num_channels is equivalent to InstanceNorm
gn = nn.GroupNorm(
num_groups=32, # number of groups (commonly: 8, 16, 32)
num_channels=256, # number of channels
eps=1e-5,
affine=True,
)
x = torch.randn(4, 256, 56, 56) # Small batch (4 images), 256 channels
out = gn(x)
print(out.shape) # torch.Size([4, 256, 56, 56])
# GroupNorm is completely unaffected when batch=1 (e.g., inference on a single image)
x_single = torch.randn(1, 256, 56, 56)
out_single = gn(x_single) # Works normally
# Special case: num_groups=1 is equivalent to LayerNorm
# num_groups=num_channels is equivalent to InstanceNorm
gn = nn.GroupNorm(
num_groups=32, # number of groups (commonly: 8, 16, 32)
num_channels=256, # number of channels
eps=1e-5,
affine=True,
)
x = torch.randn(4, 256, 56, 56) # Small batch (4 images), 256 channels
out = gn(x)
print(out.shape) # torch.Size([4, 256, 56, 56])
# GroupNorm is completely unaffected when batch=1 (e.g., inference on a single image)
x_single = torch.randn(1, 256, 56, 56)
out_single = gn(x_single) # Works normally
1.5 InstanceNorm Instance Normalization
PairEach channel of each sampleIndependently normalized, not involving the batch dimension. Commonly used inimage style transfer— preserves each image's own style information, preventing statistics from other images in the batch from interfering.
Example
# InstanceNorm2d: normalizes H×W for each (N, C)
in2d = nn.InstanceNorm2d(
num_features=64,
affine=False, # Default False (does not learn gamma/beta)
)
x = torch.randn(8, 64, 128, 128)
out = in2d(x)
print(out.shape) # torch.Size([8, 64, 128, 128])
in2d = nn.InstanceNorm2d(
num_features=64,
affine=False, # Default False (does not learn gamma/beta)
)
x = torch.randn(8, 64, 128, 128)
out = in2d(x)
print(out.shape) # torch.Size([8, 64, 128, 128])
1.6 Comparison of Normalization Methods
| Method | Normalization dimension | Batch dependency | Applicable scenarios |
|---|---|---|---|
BatchNorm |
Across batch (same channel) | Strong dependency | CNN image classification (large batch) |
LayerNorm |
Across features (same sample) | No dependency | Transformer、NLP、RNN |
GroupNorm |
Within-group channels (same sample) | No dependency | Small-batch object detection, segmentation |
InstanceNorm |
Spatial dimension (same sample and same channel) | No dependency | Image style transfer, generative models |
Quick Reference for Selection Suggestions
# batch_size >= 16, CNN image tasks → BatchNorm2d
# Transformer / BERT / GPT → LayerNorm
# batch_size < 8, object detection/segmentation → GroupNorm(num_groups=32)
# Image style transfer / CycleGAN → InstanceNorm2d
# Transformer / BERT / GPT → LayerNorm
# batch_size < 8, object detection/segmentation → GroupNorm(num_groups=32)
# Image style transfer / CycleGAN → InstanceNorm2d
2. Dropout
2.1 Basic Principles
Dropout randomly sets the outputs of some neurons to 0 during training (with probabilityp), forcing the network not to rely on any single neuron, thereby learning more robust, distributed feature representations and effectively preventing overfitting.
During training:
\[y_i = \begin{cases} 0 & \text{with probability } p \\ \dfrac{x_i}{1-p} & \text{with probability } 1-p \end{cases}\]
Retained neurons are scaled by \(\frac{1}{1-p}\) to keep the expected value unchanged (i.e., Inverted Dropout).
During inference: Dropout Automatically disabled, all neurons participate normally in computation, equivalent to the ensemble average of multiple sub-networks.
训练时示意（p=0.5）： 输入: [1.0, 2.0, 3.0, 4.0, 5.0] 掩码: [ 1, 0, 1, 0, 1 ] ← 随机生成 输出: [2.0, 0.0, 6.0, 0.0, 10.0] ← 保留的值乘以 1/(1-0.5)=2 推理时： 输出: [1.0, 2.0, 3.0, 4.0, 5.0] ← 原样输出
2.2 Dropout / Dropout2d / Dropout3d
Dropout: used for fully connected layers
Example
import torch
import torch.nn as nn
dropout = nn.Dropout(p=0.5) # p: probability of zeroing out
x = torch.ones(2, 10)
print("Training mode:")
dropout.train()
print(dropout(x)) # About 50% of values are 0, retained values are 2.0
print("\n"Evaluation mode:")
dropout.eval()
print(dropout(x)) # All 1.0, Dropout disabled
import torch.nn as nn
dropout = nn.Dropout(p=0.5) # p: probability of zeroing out
x = torch.ones(2, 10)
print("Training mode:")
dropout.train()
print(dropout(x)) # About 50% of values are 0, retained values are 2.0
print("\n"Evaluation mode:")
dropout.eval()
print(dropout(x)) # All 1.0, Dropout disabled
Example output:
训练模式:
tensor([[2., 0., 2., 0., 2., 0., 0., 2., 2., 0.],
[0., 2., 0., 2., 0., 2., 2., 0., 0., 2.]])
评估模式:
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.],
[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
Dropout2d: for CNN feature maps (entire channel drop)
Dropout2dRandomly drops entire channels, i.e., all spatial positions of a channel are zeroed out together. This is more suitable for convolutional features than pointwise Dropout, because activations of adjacent pixels are highly correlated, making pointwise dropping less effective.
Example
dropout2d = nn.Dropout2d(p=0.3)
# Input: (N, C, H, W)
x = torch.ones(4, 16, 8, 8) # 4 images, 16 channels
out = dropout2d(x)
print(out.shape) # torch.Size([4, 16, 8, 8])
# Verify: entire channel drop
# Channel drop status for the 0th image (non-zero = kept channel, all zero = dropped channel)
kept_channels = (out[0].sum(dim=(1, 2)) != 0).sum().item()
print(f"Image 0 kept {kept_channels}/16 channels")
# Input: (N, C, H, W)
x = torch.ones(4, 16, 8, 8) # 4 images, 16 channels
out = dropout2d(x)
print(out.shape) # torch.Size([4, 16, 8, 8])
# Verify: entire channel drop
# Channel drop status for the 0th image (non-zero = kept channel, all zero = dropped channel)
kept_channels = (out[0].sum(dim=(1, 2)) != 0).sum().item()
print(f"Image 0 kept {kept_channels}/16 channels")
Dropout3d: used for 3D convolution
Example
# Input: (N, C, D, H, W), the entire (D, H, W) feature volume is dropped
dropout3d = nn.Dropout3d(p=0.2)
x = torch.ones(2, 8, 16, 16, 16)
out = dropout3d(x)
print(out.shape) # torch.Size([2, 8, 16, 16, 16])
dropout3d = nn.Dropout3d(p=0.2)
x = torch.ones(2, 8, 16, 16, 16)
out = dropout3d(x)
print(out.shape) # torch.Size([2, 8, 16, 16, 16])
Recommendations for choosing the p value
Example
# Fully connected layers (wide hidden layers): higher Dropout
nn.Dropout(p=0.5) # Classic value, standard for hidden layers
# Fully connected layers (narrow hidden layers) or small networks: lower Dropout
nn.Dropout(p=0.2) # Avoid underfitting
# Dropout2d for convolutional layers: generally lower
nn.Dropout2d(p=0.1) # Convolutional layers already have some regularization effect
# Dropout before the classification head (e.g., EfficientNet / ViT)
nn.Dropout(p=0.3) # Usually between 0.2 and 0.5
# Transformer attention / FFN Dropout
nn.Dropout(p=0.1) # BERT default value, Transformers usually use smaller values
nn.Dropout(p=0.5) # Classic value, standard for hidden layers
# Fully connected layers (narrow hidden layers) or small networks: lower Dropout
nn.Dropout(p=0.2) # Avoid underfitting
# Dropout2d for convolutional layers: generally lower
nn.Dropout2d(p=0.1) # Convolutional layers already have some regularization effect
# Dropout before the classification head (e.g., EfficientNet / ViT)
nn.Dropout(p=0.3) # Usually between 0.2 and 0.5
# Transformer attention / FFN Dropout
nn.Dropout(p=0.1) # BERT default value, Transformers usually use smaller values
2.3 AlphaDropout
Specifically designed forSELU activation functionthe Dropout variant automatically adjusts the mean and variance after dropping neurons, preserving the self-normalizing property.
Example
# Used with SELU (self-normalizing network SELU + AlphaDropout)
class SelfNormalizingNet(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.SELU(),
nn.AlphaDropout(p=0.1), # ← Used with SELU
nn.Linear(hidden_dim, hidden_dim),
nn.SELU(),
nn.AlphaDropout(p=0.1),
nn.Linear(hidden_dim, num_classes),
)
def forward(self, x):
return self.net(x)
class SelfNormalizingNet(nn.Module):
def __init__(self, in_dim, hidden_dim, num_classes):
super().__init__()
self.net = nn.Sequential(
nn.Linear(in_dim, hidden_dim),
nn.SELU(),
nn.AlphaDropout(p=0.1), # ← Used with SELU
nn.Linear(hidden_dim, hidden_dim),
nn.SELU(),
nn.AlphaDropout(p=0.1),
nn.Linear(hidden_dim, num_classes),
)
def forward(self, x):
return self.net(x)
Note:AlphaDropout is only meaningful when using the SELU activation function; with other activation functions, use regular Dropout directly.
3. Training Mode and Evaluation Mode
The behavior of BatchNorm and Dropout during training and evaluationis completely different, and they must be switched correctly; this is one of the most common mistakes for beginners.
Comparison of Behavioral Differences
| Component | model.train() |
model.eval() |
|---|---|---|
| BatchNorm | Normalizes using the current batch's mean/variance;Updates running_mean/var | Normalizes using running_mean/var;No update |
| Dropout | Randomly zeroes out (with probability p) | Disabled, all neurons output normally |
Correct Usage
Example
import torch
import torch.nn as nn
model = nn.Sequential(
nn.Linear(128, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 10),
)
# ── Training phase ──────────────────────────────────────
model.train() # ← Switch to training mode (default state)
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs) # BN uses batch statistics, Dropout randomly drops
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# ── Evaluation/validation phase ─────────────────────────────────
model.eval() # ← Switch to evaluation mode (very important!)
with torch.no_grad(): # Also disable gradient computation
for inputs, labels in val_loader:
outputs = model(inputs) # BN uses running statistics, Dropout is disabled
# ...
# ── Inference on a single image ──────────────────────────────────
model.eval()
with torch.inference_mode(): # Faster than no_grad, for inference only
output = model(single_input.unsqueeze(0))
pred = output.argmax(1).item()
import torch.nn as nn
model = nn.Sequential(
nn.Linear(128, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(256, 10),
)
# ── Training phase ──────────────────────────────────────
model.train() # ← Switch to training mode (default state)
for inputs, labels in train_loader:
optimizer.zero_grad()
outputs = model(inputs) # BN uses batch statistics, Dropout randomly drops
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# ── Evaluation/validation phase ─────────────────────────────────
model.eval() # ← Switch to evaluation mode (very important!)
with torch.no_grad(): # Also disable gradient computation
for inputs, labels in val_loader:
outputs = model(inputs) # BN uses running statistics, Dropout is disabled
# ...
# ── Inference on a single image ──────────────────────────────────
model.eval()
with torch.inference_mode(): # Faster than no_grad, for inference only
output = model(single_input.unsqueeze(0))
pred = output.argmax(1).item()
Verify Whether the Mode Switch Is Correct
Example
# Use this code to verify that BN and Dropout work in the correct modes
x = torch.randn(8, 128)
model.train()
out_train_1 = model(x)
out_train_2 = model(x)
print("Training mode: same output twice?", torch.allclose(out_train_1, out_train_2))
# False ← Dropout mask differs each time, outputs differ (normal)
model.eval()
out_eval_1 = model(x)
out_eval_2 = model(x)
print("Evaluation mode: same output twice?", torch.allclose(out_eval_1, out_eval_2))
# True ← Dropout is disabled, BN uses fixed statistics, output is deterministic (normal)
x = torch.randn(8, 128)
model.train()
out_train_1 = model(x)
out_train_2 = model(x)
print("Training mode: same output twice?", torch.allclose(out_train_1, out_train_2))
# False ← Dropout mask differs each time, outputs differ (normal)
model.eval()
out_eval_1 = model(x)
out_eval_2 = model(x)
print("Evaluation mode: same output twice?", torch.allclose(out_eval_1, out_eval_2))
# True ← Dropout is disabled, BN uses fixed statistics, output is deterministic (normal)
4. Placement in the Network
Placement of BatchNorm
There are two schools of thought on the placement of BatchNorm, both of which are used in practice:
School 1: Before the activation function (original paper)
# Conv → BN → ReLU
nn.Sequential(
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64), # BN before activation
nn.ReLU(inplace=True),
)
nn.Sequential(
nn.Conv2d(32, 64, 3, padding=1),
nn.BatchNorm2d(64), # BN before activation
nn.ReLU(inplace=True),
)
School 2: After the activation function (some research suggests better performance)
# Conv → ReLU → BN
nn.Sequential(
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64), # BN after activation
)
nn.Sequential(
nn.Conv2d(32, 64, 3, padding=1),
nn.ReLU(inplace=True),
nn.BatchNorm2d(64), # BN after activation
)
Practical advice:Follow the design of the paper you are referencing; if there is no special requirement, Conv → BN → ReLU is more mainstream.
Placement of Dropout
Fully connected networks: after the activation function
nn.Sequential(
nn.Linear(512, 256),
nn.ReLU(),
nn.Dropout(0.5), # After ReLU
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(128, 10), # Usually no Dropout before the output layer
)
nn.Linear(512, 256),
nn.ReLU(),
nn.Dropout(0.5), # After ReLU
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.5),
nn.Linear(128, 10), # Usually no Dropout before the output layer
)
CNN: after BN (if BN is used, Dropout can be omitted)
# Note: BN itself has a regularization effect; if BN is already present in a CNN, Dropout is not needed in convolutional layers
# Dropout is usually only added before fully connected layers or the classification head
nn.Sequential(
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
# Usually no Dropout here (BN already provides regularization)
)
# Dropout is usually only added before fully connected layers or the classification head
nn.Sequential(
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
# Usually no Dropout here (BN already provides regularization)
)
Transformer: inside FFN and attention layers
class FeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Dropout(dropout), # Add Dropout inside FFN
nn.Linear(d_ff, d_model),
nn.Dropout(dropout), # Add once more before the output
)
def forward(self, x):
return self.net(x)
def __init__(self, d_model, d_ff, dropout=0.1):
super().__init__()
self.net = nn.Sequential(
nn.Linear(d_model, d_ff),
nn.GELU(),
nn.Dropout(dropout), # Add Dropout inside FFN
nn.Linear(d_ff, d_model),
nn.Dropout(dropout), # Add once more before the output
)
def forward(self, x):
return self.net(x)
Should BN and Dropout be used together?
Fully connected networks: can be used together
nn.Sequential(
nn.Linear(512, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.3), # BN + Dropout coexist, common and effective
nn.Linear(256, 10),
)
nn.Linear(512, 256),
nn.BatchNorm1d(256),
nn.ReLU(),
nn.Dropout(0.3), # BN + Dropout coexist, common and effective
nn.Linear(256, 10),
)
Convolutional networks: usually choose one or the other, or use Dropout only at the head.
Reason: BN's normalization operation and Dropout may cause a "variance shift" problem. That is, during training, Dropout changes the variance of activation values, while during inference, Dropout is turned off, causing the statistics seen by BN during inference to be inconsistent with those during training.
Recommended practice:
Option A: Use only BN in convolutional layers, and only Dropout in fully connected/classification heads
Example
class ConvNet(nn.Module):
def __init__(self, num_classes):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64), # Convolutional part: only BN
nn.ReLU(inplace=True),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d(1),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128, 256),
nn.ReLU(),
nn.Dropout(0.5), # Fully connected part: only Dropout
nn.Linear(256, num_classes),
)
def forward(self, x):
return self.classifier(self.features(x))
def __init__(self, num_classes):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(3, 64, 3, padding=1),
nn.BatchNorm2d(64), # Convolutional part: only BN
nn.ReLU(inplace=True),
nn.Conv2d(64, 128, 3, padding=1),
nn.BatchNorm2d(128),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d(1),
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(128, 256),
nn.ReLU(),
nn.Dropout(0.5), # Fully connected part: only Dropout
nn.Linear(256, num_classes),
)
def forward(self, x):
return self.classifier(self.features(x))
5. Complete Network Examples
Example 1: Fully Connected Classification Network (BN + Dropout)
Example
import torch
import torch.nn as nn
class MLPClassifier(nn.Module):
"""
Fully connected classifier with BatchNorm1d and Dropout
Suitable for tabular data classification tasks
"""
def __init__(self, in_dim, hidden_dims, num_classes, dropout=0.5):
super().__init__()
layers = []
prev_dim = in_dim
for hidden_dim in hidden_dims:
layers += [
nn.Linear(prev_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
]
prev_dim = hidden_dim
layers.append(nn.Linear(prev_dim, num_classes))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
model = MLPClassifier(
in_dim=784,
hidden_dims=[512, 256, 128],
num_classes=10,
dropout=0.4
)
# Verify output shape
x = torch.randn(32, 784)
model.eval()
print(model(x).shape) # torch.Size([32, 10])
import torch.nn as nn
class MLPClassifier(nn.Module):
"""
Fully connected classifier with BatchNorm1d and Dropout
Suitable for tabular data classification tasks
"""
def __init__(self, in_dim, hidden_dims, num_classes, dropout=0.5):
super().__init__()
layers = []
prev_dim = in_dim
for hidden_dim in hidden_dims:
layers += [
nn.Linear(prev_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(inplace=True),
nn.Dropout(dropout),
]
prev_dim = hidden_dim
layers.append(nn.Linear(prev_dim, num_classes))
self.net = nn.Sequential(*layers)
def forward(self, x):
return self.net(x)
model = MLPClassifier(
in_dim=784,
hidden_dims=[512, 256, 128],
num_classes=10,
dropout=0.4
)
# Verify output shape
x = torch.randn(32, 784)
model.eval()
print(model(x).shape) # torch.Size([32, 10])
Example 2: CNN Image Classification (mainly BN)
Example
class ConvBNBlock(nn.Module):
"""Standard convolutional block: Conv → BN → ReLU"""
def __init__(self, in_ch, out_ch, stride=1):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False),
nn.BatchNorm2d(out_ch), # bias=False: BN's beta already serves as the bias
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.block(x)
class SmallCNN(nn.Module):
"""Small CNN for CIFAR-10"""
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
ConvBNBlock(3, 32), # 32×32 → 32×32
ConvBNBlock(32, 64, stride=2), # 32×32 → 16×16
ConvBNBlock(64, 128, stride=2), # 16×16 → 8×8
ConvBNBlock(128, 256, stride=2), # 8×8 → 4×4
nn.AdaptiveAvgPool2d(1), # 4×4 → 1×1
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.5), # Use Dropout only in fully connected layers
nn.Linear(128, num_classes),
)
def forward(self, x):
return self.classifier(self.features(x))
model = SmallCNN(num_classes=10)
x = torch.randn(16, 3, 32, 32)
model.eval()
print(model(x).shape) # torch.Size([16, 10])
"""Standard convolutional block: Conv → BN → ReLU"""
def __init__(self, in_ch, out_ch, stride=1):
super().__init__()
self.block = nn.Sequential(
nn.Conv2d(in_ch, out_ch, 3, stride=stride, padding=1, bias=False),
nn.BatchNorm2d(out_ch), # bias=False: BN's beta already serves as the bias
nn.ReLU(inplace=True),
)
def forward(self, x):
return self.block(x)
class SmallCNN(nn.Module):
"""Small CNN for CIFAR-10"""
def __init__(self, num_classes=10):
super().__init__()
self.features = nn.Sequential(
ConvBNBlock(3, 32), # 32×32 → 32×32
ConvBNBlock(32, 64, stride=2), # 32×32 → 16×16
ConvBNBlock(64, 128, stride=2), # 16×16 → 8×8
ConvBNBlock(128, 256, stride=2), # 8×8 → 4×4
nn.AdaptiveAvgPool2d(1), # 4×4 → 1×1
)
self.classifier = nn.Sequential(
nn.Flatten(),
nn.Linear(256, 128),
nn.ReLU(),
nn.Dropout(0.5), # Use Dropout only in fully connected layers
nn.Linear(128, num_classes),
)
def forward(self, x):
return self.classifier(self.features(x))
model = SmallCNN(num_classes=10)
x = torch.randn(16, 3, 32, 32)
model.eval()
print(model(x).shape) # torch.Size([16, 10])
Example 3: Complete Training Loop
Example
import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
EPOCHS = 50
model = SmallCNN(num_classes=10).to(DEVICE)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)
def run_epoch(model, loader, optimizer=None, train=True):
if train:
model.train() # ← BN uses batch statistics, Dropout activates
else:
model.eval() # ← BN uses running statistics, Dropout is turned off
total_loss, correct = 0.0, 0
ctx = torch.enable_grad() if train else torch.no_grad()
with ctx:
for imgs, labels in loader:
imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
if train:
optimizer.zero_grad()
outputs = model(imgs)
loss = criterion(outputs, labels)
if train:
loss.backward()
optimizer.step()
total_loss += loss.item() * imgs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
best_acc = 0.0
for epoch in range(1, EPOCHS + 1):
train_loss, train_acc = run_epoch(model, train_loader, optimizer, train=True)
val_loss, val_acc = run_epoch(model, val_loader, train=False)
scheduler.step()
print(f"Epoch {epoch:2d}/{EPOCHS} | "
f"Train {train_loss:.4f}/{train_acc:.4f} | "
f"Val {val_loss:.4f}/{val_acc:.4f}")
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'best.pth')
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
EPOCHS = 50
model = SmallCNN(num_classes=10).to(DEVICE)
criterion = nn.CrossEntropyLoss(label_smoothing=0.1)
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)
def run_epoch(model, loader, optimizer=None, train=True):
if train:
model.train() # ← BN uses batch statistics, Dropout activates
else:
model.eval() # ← BN uses running statistics, Dropout is turned off
total_loss, correct = 0.0, 0
ctx = torch.enable_grad() if train else torch.no_grad()
with ctx:
for imgs, labels in loader:
imgs, labels = imgs.to(DEVICE), labels.to(DEVICE)
if train:
optimizer.zero_grad()
outputs = model(imgs)
loss = criterion(outputs, labels)
if train:
loss.backward()
optimizer.step()
total_loss += loss.item() * imgs.size(0)
correct += (outputs.argmax(1) == labels).sum().item()
n = len(loader.dataset)
return total_loss / n, correct / n
best_acc = 0.0
for epoch in range(1, EPOCHS + 1):
train_loss, train_acc = run_epoch(model, train_loader, optimizer, train=True)
val_loss, val_acc = run_epoch(model, val_loader, train=False)
scheduler.step()
print(f"Epoch {epoch:2d}/{EPOCHS} | "
f"Train {train_loss:.4f}/{train_acc:.4f} | "
f"Val {val_loss:.4f}/{val_acc:.4f}")
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'best.pth')
6. Hyperparameter Tuning Guide
BatchNorm Parameter Tuning
Example
# ── momentum (the parameter most often needing adjustment) ─────────────────
# Default value 0.1 works in most cases
# batch_size is very small (< 8) → statistics are unstable → decrease momentum
nn.BatchNorm2d(64, momentum=0.01)
# Many training steps, stable data distribution → can appropriately increase momentum
nn.BatchNorm2d(64, momentum=0.2)
# ── affine=False (do not learn gamma/beta) ───────────────
# Rarely used, usually keep default True
# Unless you explicitly do not want BN to change the scaling and offset of features
nn.BatchNorm2d(64, affine=False)
# ── track_running_stats=False ──────────────────────
# Each inference also uses current batch statistics (training and inference behavior are consistent)
# Suitable for special scenarios such as online learning and streaming data
nn.BatchNorm2d(64, track_running_stats=False)
# Default value 0.1 works in most cases
# batch_size is very small (< 8) → statistics are unstable → decrease momentum
nn.BatchNorm2d(64, momentum=0.01)
# Many training steps, stable data distribution → can appropriately increase momentum
nn.BatchNorm2d(64, momentum=0.2)
# ── affine=False (do not learn gamma/beta) ───────────────
# Rarely used, usually keep default True
# Unless you explicitly do not want BN to change the scaling and offset of features
nn.BatchNorm2d(64, affine=False)
# ── track_running_stats=False ──────────────────────
# Each inference also uses current batch statistics (training and inference behavior are consistent)
# Suitable for special scenarios such as online learning and streaming data
nn.BatchNorm2d(64, track_running_stats=False)
Dropout Tuning Strategies
Example
# ── Adjust p according to the degree of overfitting ───────────────────────────
# Training acc much higher than validation acc (severe overfitting) → increase p
# Training acc and validation acc close but both low (underfitting) → decrease p or remove Dropout
# ── Recommended p value reference ───────────────────────────────
p_values = {
'Large fully connected hidden layers (> 1024)': 0.5,
'Medium fully connected hidden layers (256~1024)': 0.3,
'Small fully connected hidden layers (< 256)': 0.2,
'CNN convolutional layers (Dropout2d)': 0.1,
'Before classification head': 0.3,
'Transformer FFN': 0.1,
'Transformer attention': 0.1,
}
# ── Dynamic Dropout (adjusted as training progresses) ──────────────────
# In some practices, p is small early in training and increased later (strengthen regularization as the model becomes more complex)
class DynamicDropoutNet(nn.Module):
def __init__(self, initial_p=0.1):
super().__init__()
self.dropout = nn.Dropout(initial_p)
self.fc = nn.Linear(256, 10)
def set_dropout(self, p):
self.dropout.p = p
def forward(self, x):
return self.fc(self.dropout(x))
model = DynamicDropoutNet(initial_p=0.1)
# Mid-training, increase Dropout
model.set_dropout(0.4)
# Training acc much higher than validation acc (severe overfitting) → increase p
# Training acc and validation acc close but both low (underfitting) → decrease p or remove Dropout
# ── Recommended p value reference ───────────────────────────────
p_values = {
'Large fully connected hidden layers (> 1024)': 0.5,
'Medium fully connected hidden layers (256~1024)': 0.3,
'Small fully connected hidden layers (< 256)': 0.2,
'CNN convolutional layers (Dropout2d)': 0.1,
'Before classification head': 0.3,
'Transformer FFN': 0.1,
'Transformer attention': 0.1,
}
# ── Dynamic Dropout (adjusted as training progresses) ──────────────────
# In some practices, p is small early in training and increased later (strengthen regularization as the model becomes more complex)
class DynamicDropoutNet(nn.Module):
def __init__(self, initial_p=0.1):
super().__init__()
self.dropout = nn.Dropout(initial_p)
self.fc = nn.Linear(256, 10)
def set_dropout(self, p):
self.dropout.p = p
def forward(self, x):
return self.fc(self.dropout(x))
model = DynamicDropoutNet(initial_p=0.1)
# Mid-training, increase Dropout
model.set_dropout(0.4)
7. Common Errors and Cautions
Error 1: Forgetting to switch to eval() mode
Incorrect code
# Error: model.eval() not called during evaluation
model.train() # Default state
for imgs, labels in val_loader:
with torch.no_grad():
outputs = model(imgs) # Dropout is still randomly dropping! BN is still updating!
# Results are non-deterministic, accuracy is artificially low, and running_mean/var is polluted
model.train() # Default state
for imgs, labels in val_loader:
with torch.no_grad():
outputs = model(imgs) # Dropout is still randomly dropping! BN is still updating!
# Results are non-deterministic, accuracy is artificially low, and running_mean/var is polluted
Correct code
# Correct: explicitly switch modes before evaluation
model.eval()
with torch.no_grad():
outputs = model(imgs) # Deterministic output
model.eval()
with torch.no_grad():
outputs = model(imgs) # Deterministic output
Error 2: Adding bias to a Conv layer and then using BN
Incorrect code
# Redundant: Conv's bias will be canceled by BN's beta, wasting parameters
nn.Conv2d(64, 128, 3, padding=1, bias=True), # Default True
nn.BatchNorm2d(128),
nn.Conv2d(64, 128, 3, padding=1, bias=True), # Default True
nn.BatchNorm2d(128),
Correct code
# Correct: disable bias when Conv is followed by BN
nn.Conv2d(64, 128, 3, padding=1, bias=False), # ← bias=False
nn.BatchNorm2d(128),
nn.Conv2d(64, 128, 3, padding=1, bias=False), # ← bias=False
nn.BatchNorm2d(128),
Error 3: Using BatchNorm with batch size=1
Incorrect code
# Error: BN cannot compute variance when batch=1, training will produce NaN or inf
# Single-sample inference is fine (uses running stats), but single-sample training will crash
model.train()
x = torch.randn(1, 64, 32, 32)
out = bn2d(x) # ← batch=1 in training mode will error or produce NaN
# Single-sample inference is fine (uses running stats), but single-sample training will crash
model.train()
x = torch.randn(1, 64, 32, 32)
out = bn2d(x) # ← batch=1 in training mode will error or produce NaN
Solution
# Solution 1: Use GroupNorm instead
nn.GroupNorm(num_groups=32, num_channels=64)
# Solution 2: Use LayerNorm
nn.LayerNorm([64, 32, 32])
# Solution 3: Switch to eval() during inference, does not affect BN
model.eval()
out = bn2d(x) # batch=1 in inference mode is completely normal
nn.GroupNorm(num_groups=32, num_channels=64)
# Solution 2: Use LayerNorm
nn.LayerNorm([64, 32, 32])
# Solution 3: Switch to eval() during inference, does not affect BN
model.eval()
out = bn2d(x) # batch=1 in inference mode is completely normal
Error 4: Dropout causes variance shift in BN
Problem
# Problem: Dropout changes the variance of activation values during training
# During inference, Dropout is turned off, and BN uses the training-time running_var (containing Dropout's effect)
# Statistics do not match, inference results are biased
nn.Sequential(
nn.Conv2d(64, 128, 3),
nn.Dropout2d(0.5), # ← Use Dropout in convolutional layers
nn.BatchNorm2d(128), # ← The variance seen by BN is changed by Dropout
)
# During inference, Dropout is turned off, and BN uses the training-time running_var (containing Dropout's effect)
# Statistics do not match, inference results are biased
nn.Sequential(
nn.Conv2d(64, 128, 3),
nn.Dropout2d(0.5), # ← Use Dropout in convolutional layers
nn.BatchNorm2d(128), # ← The variance seen by BN is changed by Dropout
)
Solution
# Solution 1: BN first, Dropout after
nn.Sequential(
nn.Conv2d(64, 128, 3),
nn.BatchNorm2d(128), # BN first (not affected by Dropout)
nn.ReLU(),
nn.Dropout2d(0.1), # Dropout after
)
# Solution 2: Use only BN in convolutional layers and only Dropout in fully connected layers (most recommended)
nn.Sequential(
nn.Conv2d(64, 128, 3),
nn.BatchNorm2d(128), # BN first (not affected by Dropout)
nn.ReLU(),
nn.Dropout2d(0.1), # Dropout after
)
# Solution 2: Use only BN in convolutional layers and only Dropout in fully connected layers (most recommended)
Error 5: Forgetting to switch frozen BN layers to eval mode
Incorrect code
# In transfer learning, when freezing the backbone, the pretrained BN running stats should not be updated
# Otherwise, the new data distribution will pollute the pretrained statistics
# Error: parameters are frozen, but BN's running_mean/var are still being updated
model = models.resnet50(weights='IMAGENET1K_V2')
for param in model.parameters():
param.requires_grad = False # Only learnable parameters are frozen; BN statistics are still updated
# Otherwise, the new data distribution will pollute the pretrained statistics
# Error: parameters are frozen, but BN's running_mean/var are still being updated
model = models.resnet50(weights='IMAGENET1K_V2')
for param in model.parameters():
param.requires_grad = False # Only learnable parameters are frozen; BN statistics are still updated
Correct code
# Correct solution: also switch the backbone BN layers to eval mode
def freeze_bn(model):
for module in model.modules():
if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
module.eval() # Fix running_mean/var
for param in module.parameters():
param.requires_grad = False
model = models.resnet50(weights='IMAGENET1K_V2')
freeze_bn(model) # Freeze all BN
model.fc = nn.Linear(model.fc.in_features, 10) # Only the new FC layer is trainable
def freeze_bn(model):
for module in model.modules():
if isinstance(module, (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d)):
module.eval() # Fix running_mean/var
for param in module.parameters():
param.requires_grad = False
model = models.resnet50(weights='IMAGENET1K_V2')
freeze_bn(model) # Freeze all BN
model.fc = nn.Linear(model.fc.in_features, 10) # Only the new FC layer is trainable
Quick Diagnostic Checklist
| Phenomenon | Possible cause | Check item |
|---|---|---|
| Loss fluctuates unpredictably | Forgot [model.eval()]model.eval() |
Confirm calling [model.eval()] before evaluationeval() |
| Large gap between training/validation acc | Dropout pDropout p too large |
Try decreasing ppor remove Dropout |
| Training NaN/inf | BatchNorm used with batch=1 | Switch to GroupNorm or LayerNorm |
| Poor transfer learning performance | BN running stats polluted | Call [.eval()] on backbone BN layerseval() |
| Unstable model inference results | Didn't switch to eval() during inferenceeval() |
Must call [model.eval()] before inferencemodel.eval() |
| Excessive Conv parameter count | Conv+BN did not disable bias | Set [bias=False]bias=False |