Batch Normalization

import numpy as np
class BatchNorm1D:
    """Batch Normalization implementation from scratch"""
    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 = np.ones(num_features)  # Scale
        self.beta = np.zeros(num_features)  # Shift
        # Running statistics (for inference)
        self.running_mean = np.zeros(num_features)
        self.running_var = np.ones(num_features)
    def forward(self, x, training=True):
        """
        Args:
            x: Input of shape (batch_size, num_features)
            training: If True, use batch statistics; if False, use running statistics
        Returns:
            out: Normalized output
        """
        if training:
            # Step 1: Compute batch statistics
            batch_mean = x.mean(axis=0)
            batch_var = x.var(axis=0)
            # Step 2: Normalize
            x_normalized = (x - batch_mean) / np.sqrt(batch_var + self.eps)
            # Step 3: Scale and shift
            out = self.gamma * x_normalized + self.beta
            # Step 4: Update running statistics (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_normalized = (x - self.running_mean) / np.sqrt(self.running_var + self.eps)
            out = self.gamma * x_normalized + self.beta
        return out
# Example usage
bn = BatchNorm1D(num_features=3)
# Training data (batch_size=4, num_features=3)
x_train = np.random.randn(4, 3) * 10 + 5  # Mean ~5, high variance
print(f"Input mean: {x_train.mean(axis=0)}")
print(f"Input var: {x_train.var(axis=0)}")
# Forward pass in training mode
out_train = bn.forward(x_train, training=True)
print(f"\nOutput mean (training): {out_train.mean(axis=0)}")  # ~0 (normalized)
print(f"Output var (training): {out_train.var(axis=0)}")  # ~1 (normalized)
# Forward pass in inference mode
out_test = bn.forward(x_train, training=False)
print(f"\nOutput (inference): {out_test.mean(axis=0)}")

import torch
import torch.nn as nn
class CNNWithBatchNorm(nn.Module):
    """Standard CNN with Batch Normalization"""
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            # Block 1: Conv → BatchNorm → ReLU → Pool
            nn.Conv2d(3, 64, kernel_size=3, padding=1),
            nn.BatchNorm2d(64),  # BatchNorm2D for convolutional layers
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            # Block 2
            nn.Conv2d(64, 128, kernel_size=3, padding=1),
            nn.BatchNorm2d(128),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
            # Block 3
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(2, 2),
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256 * 4 * 4, 512),
            nn.BatchNorm1d(512),  # BatchNorm1D for fully connected layers
            nn.ReLU(inplace=True),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes)
        )
    def forward(self, x):
        x = self.features(x)
        x = self.classifier(x)
        return x
# Create model
model = CNNWithBatchNorm(num_classes=10)
print(model)
# Example input (batch of 8 RGB images, 32x32)
x = torch.randn(8, 3, 32, 32)
# CRITICAL: Set training mode
model.train()
output_train = model(x)
print(f"\nTraining mode output shape: {output_train.shape}")
# CRITICAL: Set eval mode for testing
model.eval()
output_eval = model(x)
print(f"Eval mode output shape: {output_eval.shape}")
# Inspect BatchNorm statistics
bn_layer = model.features[1]  # First BatchNorm2d layer
print(f"\nBatchNorm running mean: {bn_layer.running_mean[:5]}")
print(f"BatchNorm running var: {bn_layer.running_var[:5]}")
print(f"Learnable gamma (scale): {bn_layer.weight[:5]}")
print(f"Learnable beta (shift): {bn_layer.bias[:5]}")

import torch
import torch.nn as nn
# Sample input: (batch=2, channels=4, height=3, width=3)
x = torch.randn(2, 4, 3, 3)
print("Input shape:", x.shape)
print("="*60)
# 1. Batch Normalization
# Normalizes across batch dimension for each channel
bn = nn.BatchNorm2d(4)
out_bn = bn(x)
print("\n1. Batch Normalization (BatchNorm2d)")
print(f"Normalizes over: Batch dimension (N)")
print(f"Output shape: {out_bn.shape}")
print(f"For each channel, mean across batch ≈ 0:")
print(f"  Channel 0 mean: {out_bn[:, 0].mean():.6f}")
# 2. Layer Normalization
# Normalizes across channel dimension for each sample
ln = nn.LayerNorm([4, 3, 3])  # Normalize over C, H, W
out_ln = ln(x)
print("\n2. Layer Normalization (LayerNorm)")
print(f"Normalizes over: Channel, Height, Width dimensions")
print(f"Output shape: {out_ln.shape}")
print(f"For each sample in batch, mean ≈ 0:")
print(f"  Sample 0 mean: {out_ln[0].mean():.6f}")
# 3. Instance Normalization
# Normalizes across spatial dimensions for each channel and sample
in_norm = nn.InstanceNorm2d(4)
out_in = in_norm(x)
print("\n3. Instance Normalization (InstanceNorm2d)")
print(f"Normalizes over: Spatial dimensions (H, W) for each channel")
print(f"Output shape: {out_in.shape}")
print(f"For each sample and channel, mean across spatial dims ≈ 0:")
print(f"  Sample 0, Channel 0 mean: {out_in[0, 0].mean():.6f}")
# 4. Group Normalization
# Normalizes across groups of channels
gn = nn.GroupNorm(num_groups=2, num_channels=4)  # 4 channels, 2 groups
out_gn = gn(x)
print("\n4. Group Normalization (GroupNorm)")
print(f"Normalizes over: Groups of channels (here: 2 groups)")
print(f"Output shape: {out_gn.shape}")
# Comparison table
print("\n" + "="*60)
print("WHEN TO USE EACH:")
print("="*60)
print("BatchNorm:     Large batches, CNNs, standard training")
print("LayerNorm:     Transformers, RNNs, small batches, NLP")
print("InstanceNorm:  Style transfer, GANs")
print("GroupNorm:     Small batches, object detection, batch-independent")

import torch
import torch.nn as nn
# Create a simple model with BatchNorm
model = nn.Sequential(
    nn.Linear(10, 20),
    nn.BatchNorm1d(20),
    nn.ReLU()
)
# Sample input
x = torch.randn(32, 10)  # Batch of 32
print("="*60)
print("TRAINING MODE (model.train())")
print("="*60)
model.train()
output1 = model(x)
output2 = model(x)  # Same input
print(f"Output 1 mean: {output1.mean():.6f}")
print(f"Output 2 mean: {output2.mean():.6f}")
print(f"Outputs are identical: {torch.allclose(output1, output2)}")
print("❌ Different outputs! (uses different batch statistics each time)")
print("\n" + "="*60)
print("EVAL MODE (model.eval())")
print("="*60)
model.eval()
output3 = model(x)
output4 = model(x)  # Same input
print(f"Output 3 mean: {output3.mean():.6f}")
print(f"Output 4 mean: {output4.mean():.6f}")
print(f"Outputs are identical: {torch.allclose(output3, output4)}")
print("✓ Identical outputs! (uses fixed running statistics)")
print("\n" + "="*60)
print("WHY THIS MATTERS")
print("="*60)
print("""
Training Mode (model.train()):
  - Uses batch statistics (mean/var of current batch)
  - Different batches → different statistics → different outputs
  - Updates running statistics
Eval Mode (model.eval()):
  - Uses running statistics (accumulated during training)
  - Deterministic behavior
  - Essential for reproducible inference
COMMON BUG: Forgetting model.eval() during testing
→ Results in inconsistent predictions!
""")
# Demonstrate the critical importance
bn_layer = model[1]
print(f"Running mean: {bn_layer.running_mean[:5]}")
print(f"Running var: {bn_layer.running_var[:5]}")

import torch
import torch.nn as nn
import torch.optim as optim
import time
# Synthetic dataset
torch.manual_seed(42)
X_train = torch.randn(1000, 100)
y_train = (X_train.sum(dim=1) > 0).long()
# Model WITHOUT BatchNorm
class ModelWithoutBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 2)
        )
    def forward(self, x):
        return self.network(x)
# Model WITH BatchNorm
class ModelWithBN(nn.Module):
    def __init__(self):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(100, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Linear(256, 2)
        )
    def forward(self, x):
        return self.network(x)
def train_model(model, lr, epochs=50):
    optimizer = optim.SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()
    losses = []
    start_time = time.time()
    for epoch in range(epochs):
        model.train()
        optimizer.zero_grad()
        outputs = model(X_train)
        loss = criterion(outputs, y_train)
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    train_time = time.time() - start_time
    return losses, train_time
# Train without BatchNorm (requires low learning rate)
print("Training WITHOUT BatchNorm (lr=0.01)...")
model_no_bn = ModelWithoutBN()
losses_no_bn, time_no_bn = train_model(model_no_bn, lr=0.01)
# Train with BatchNorm (can use much higher learning rate)
print("Training WITH BatchNorm (lr=0.1 - 10x higher!)...")
model_with_bn = ModelWithBN()
losses_with_bn, time_with_bn = train_model(model_with_bn, lr=0.1)
# Results
print("\n" + "="*60)
print("RESULTS")
print("="*60)
print(f"Without BatchNorm:")
print(f"  Final loss: {losses_no_bn[-1]:.4f}")
print(f"  Training time: {time_no_bn:.2f}s")
print(f"\nWith BatchNorm:")
print(f"  Final loss: {losses_with_bn[-1]:.4f}")
print(f"  Training time: {time_with_bn:.2f}s")
print(f"\nSpeedup: {time_no_bn / time_with_bn:.2f}x faster convergence")
print("\n" + "="*60)
print("KEY INSIGHTS")
print("="*60)
print("""
1. BatchNorm allows 10x higher learning rate (0.1 vs 0.01)
2. Converges faster (fewer epochs to reach same loss)
3. More stable training
4. Less sensitive to weight initialization
Try training without BatchNorm at lr=0.1 → training will diverge!
""")