Batch Normalization › Code Implementation
Language AI HandbookAll chapters Batch Normalization › Code ImplementationOne section · Sections in this chapter13 / 17 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
This section walks through implementing batch normalization from scratch in PyTorch, then shows how to use PyTorch's built-in nn.BatchNorm1d and nn.BatchNorm2d modules in a realistic network. The goal is to solidify your understanding of both the mechanics and the practical pitfalls.
Batch Norm from Scratch
Let's implement batch normalization from scratch to see exactly how the forward pass works, including the training/inference mode switch.
import torch
class BatchNorm1dManual:
"""
Manual batch normalization for 1D inputs.
Demonstrates the training vs inference modes.
"""
def __init__(self, num_features, eps=1e-5, momentum=0.1):
self.num_features = num_features
self.eps = eps
self.momentum = momentum
# Learnable parameters
self.gamma = torch.ones(num_features)
self.beta = torch.zeros(num_features)
# Running statistics (updated during training, used during inference)
self.running_mean = torch.zeros(num_features)
self.running_var = torch.ones(num_features)
self.training = True
def forward(self, x):
if self.training:
# Compute batch statistics
batch_mean = x.mean(dim=0)
batch_var = x.var(dim=0, unbiased=False) # Population variance
# Normalize using batch statistics
x_hat = (x - batch_mean) / torch.sqrt(batch_var + self.eps)
# Update running statistics with exponential moving average
self.running_mean = (
1 - self.momentum
) * self.running_mean + self.momentum * batch_mean
self.running_var = (
1 - self.momentum
) * self.running_var + self.momentum * batch_var
else:
# Use running statistics for inference
x_hat = (x - self.running_mean) / torch.sqrt(
self.running_var + self.eps
)
# Scale and shift with learned parameters
return self.gamma * x_hat + self.betaInput activations: [[2. 1. 3.] [4. 3. 1.] [3. 2. 4.] [1. 4. 2.]] Input mean per feature: [2.5, 2.5, 2.5] Input std per feature: [1.291, 1.291, 1.291] Output after batch norm (training mode): [[-0.4472 -1.3416 0.4472] [ 1.3416 0.4472 -1.3416] [ 0.4472 -0.4472 1.3416] [-1.3416 1.3416 -0.4472]] Output mean per feature: [0.0, 0.0, 0.0] Output std per feature: [1.1547, 1.1547, 1.1547]
The output shows that batch normalization brings the mean of each feature very close to zero and the standard deviation close to one. With the initial and , the output is simply the normalized input.
Using PyTorch's Built-in Batch Norm
PyTorch provides nn.BatchNorm1d (for fully connected layers) and nn.BatchNorm2d (for convolutional layers). Let's see how to use them and inspect their internal state.
# BatchNorm1d: for fully connected layers
# Input shape: (batch_size, num_features)
bn1d = nn.BatchNorm1d(num_features=128)
# BatchNorm2d: for convolutional layers
# Input shape: (batch_size, channels, height, width)
bn2d = nn.BatchNorm2d(num_features=64) # 64 channelsBatchNorm1d: gamma (weight) shape = torch.Size([128]), beta (bias) shape = torch.Size([128]) gamma initialized to all 1.0: True beta initialized to all 0.0: True BatchNorm2d: gamma (weight) shape = torch.Size([64]) (one per channel) Running mean (first 3): [0.0, 0.0, 0.0] Running var (first 3): [1.0, 1.0, 1.0]
PyTorch initializes gamma (weight) to ones and beta (bias) to zeros, matching our manual implementation. The running statistics start at zero mean and unit variance.
Building a Network with Batch Normalization
Let's build a small fully connected classifier that incorporates batch normalization, and observe the effect on training stability compared to a network without it.
from sklearn.datasets import make_classification
from sklearn.preprocessing import StandardScaler
from torch.utils.data import DataLoader, TensorDataset
torch.manual_seed(42)
np.random.seed(42)
# Generate a synthetic classification dataset
X_np, y_np = make_classification(
n_samples=1000,
n_features=20,
n_informative=10,
n_redundant=5,
n_classes=2,
class_sep=0.8,
random_state=42,
)
# Standardize inputs
scaler = StandardScaler()
X_np = scaler.fit_transform(X_np).astype(np.float32)
y_np = y_np.astype(np.int64)
# Split into train/test
split = 800
X_train = torch.tensor(X_np[:split])
y_train = torch.tensor(y_np[:split])
X_test = torch.tensor(X_np[split:])
y_test = torch.tensor(y_np[split:])
train_loader = DataLoader(
TensorDataset(X_train, y_train), batch_size=64, shuffle=True
)# Define two architectures: with and without batch normalization
class MLPWithBN(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.BatchNorm1d(hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim),
)
def forward(self, x):
return self.net(x)
class MLPWithoutBN(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super().__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, hidden_dim),
nn.ReLU(),
nn.Linear(hidden_dim, output_dim),
)
def forward(self, x):
return self.net(x)
def train_model(model, loader, n_epochs=30, lr=0.01):
optimizer = optim.SGD(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()
losses = []
for epoch in range(n_epochs):
model.train()
epoch_loss = 0.0
for X_batch, y_batch in loader:
optimizer.zero_grad()
logits = model(X_batch)
loss = criterion(logits, y_batch)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
losses.append(epoch_loss / len(loader))
return losses
# Train both models - batch norm allows a higher learning rate
model_bn = MLPWithBN(input_dim=20, hidden_dim=64, output_dim=2)
model_no_bn = MLPWithoutBN(input_dim=20, hidden_dim=64, output_dim=2)
losses_bn = train_model(model_bn, train_loader, n_epochs=30, lr=0.05)
losses_no_bn = train_model(model_no_bn, train_loader, n_epochs=30, lr=0.01)# Evaluate both models on the test set
def evaluate_model(model, X, y):
model.eval()
with torch.no_grad():
logits = model(X)
preds = logits.argmax(dim=1)
accuracy = (preds == y).float().mean().item()
return accuracy
acc_bn = evaluate_model(model_bn, X_test, y_test)
acc_no_bn = evaluate_model(model_no_bn, X_test, y_test)Test accuracy WITH batch normalization: 0.935 Test accuracy WITHOUT batch normalization: 0.715 Final training loss WITH batch norm: 0.0820 Final training loss WITHOUT batch norm: 0.6064 Batch norm model trained with lr=0.05; plain model trained with lr=0.01 (Higher lr without batch norm would diverge)
The model with batch normalization achieves better accuracy and lower training loss, even though it used a higher learning rate (0.05 vs 0.01). Without batch normalization, using the higher learning rate would likely cause training to diverge. This demonstrates one of the most practically important effects of batch normalization: it allows more aggressive learning rates, which directly translate to faster convergence and better final performance within a fixed training budget.
The training curves below show how batch normalization enables faster convergence.
The training=True vs training=False Distinction
One of the most important practical aspects of batch normalization is correctly switching between training and inference modes. Let's verify how model.eval() changes behavior and understand why it matters.
# Demonstrate the effect of switching to inference mode on batch norm behavior.
# We train briefly on a batch to warm up running statistics, then compare
# outputs between training mode and inference mode on the same single sample.
torch.manual_seed(42)
test_model = nn.Sequential(nn.Linear(5, 8), nn.BatchNorm1d(8), nn.ReLU())
# Pass a small training batch to warm up the running statistics
x_batch = torch.randn(16, 5)
test_model.train()
_ = test_model(x_batch) # Updates running_mean and running_var
# Now test with a single sample in BOTH modes
x_single = torch.randn(1, 5)
# Training mode on a single sample raises ValueError in PyTorch.
# Use a 2-sample batch to approximate what training mode would do.
x_two = torch.cat([x_single, torch.randn(1, 5)], dim=0)
test_model.train()
out_train_approx = test_model(x_two)[0:1] # Keep only the first sample's output
# Inference mode: uses stable running statistics, works with any batch size
test_model.eval()
out_inference = test_model(x_single)
outputs_match = torch.allclose(out_train_approx, out_inference, atol=1e-2)Training mode output (first 4 values): [1. 0. 0. 1.] Inference mode output (first 4 values): [0. 0. 0. 0.8766] Outputs match (within tolerance of 0.01): False Training mode uses live batch statistics, producing different outputs depending on which other samples appear in the batch. Inference mode uses fixed running statistics for consistent, deterministic predictions. Always switch to inference mode before generating predictions!
With a batch size of 1 in training mode, batch normalization computes the variance of a single value, which is always zero. This produces the wrong result and PyTorch will raise an error. Switching to inference mode before generating predictions ensures the stable running statistics are used, which is always the correct behavior.
Key Parameters
The key parameters for nn.BatchNorm1d and nn.BatchNorm2d are:
- num_features: The number of features (for
BatchNorm1d) or channels (forBatchNorm2d). This determines the shape of the learnable and parameters. - eps: The small constant added to the variance for numerical stability. Default is
1e-5. Rarely needs tuning, but you might increase it to1e-3if you experience NaN losses due to very small variance estimates. - momentum: The momentum for updating running statistics. Default is
0.1in PyTorch, meaning new batches contribute 10% to the running estimate. Use smaller values (e.g., 0.01) for very large datasets where you want the running statistics to update conservatively. - affine: If True (default), learnable and are included. If False, batch norm only normalizes without the scale and shift step. Setting
affine=Falseis occasionally useful when you want normalization without additional parameters. - track_running_stats: If True (default), maintains running mean and variance for inference. If False, always uses batch statistics. This can be useful for fine-tuning scenarios where you want the normalization to adapt to a new domain's statistics.