Deep Learning With PyTorch

A neural network architecture only determines what the model could learn. Whether it actually learns depends on your training setup: the optimizer, learning rate schedule, regularization strategy, and normalization layers. This is where most deep learning time is spent, and where most failures originate.

Gradient Descent Variants

Batch gradient descent: Compute gradient over the entire dataset per step. Exact gradient, but slow - one parameter update per epoch.

Stochastic gradient descent (SGD): One random example per step. Fast but extremely noisy.

Mini-batch gradient descent: Gradient over a batch of 32–512 examples. The standard approach - balances noise (exploration) with signal (exploitation).

The noise in mini-batch SGD is actually beneficial: it acts as implicit regularization by preventing the optimizer from settling too deeply into sharp minima (which often don't generalize).

Adam and Its Variants

Adam (Adaptive Moment Estimation) maintains per-parameter estimates of both the first moment (mean gradient) and second moment (variance of gradient). It adapts the learning rate for each parameter:

pythonCopy
import torch.optim as optim

# Adam: fast convergence, good default
optimizer = optim.Adam(model.parameters(), lr=1e-3, betas=(0.9, 0.999), eps=1e-8)

# AdamW: Adam + decoupled weight decay - better generalization than Adam
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

# SGD with momentum: slower but sometimes better final performance than Adam
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9, weight_decay=1e-4)

Rule of thumb: Use AdamW for transformers and most modern architectures. Use SGD with momentum for CNNs when you need maximum final performance and can spend time on tuning. AdamW almost always outperforms vanilla Adam due to the decoupled weight decay.

Learning Rate: The Most Critical Hyperparameter

Too high: loss diverges. Too low: training stalls. The learning rate has a larger effect on final performance than almost any architectural choice.

pythonCopy
from torch.optim.lr_scheduler import OneCycleLR, CosineAnnealingLR, ReduceLROnPlateau

# OneCycleLR: warm up to peak LR, then cosine anneal - great for single-shot training
scheduler = OneCycleLR(
    optimizer, max_lr=1e-2,
    steps_per_epoch=len(loader),
    epochs=50
)
# Call scheduler.step() after each batch, not each epoch

# CosineAnnealing: gradually reduces LR following a cosine curve
scheduler = CosineAnnealingLR(optimizer, T_max=50, eta_min=1e-6)
# Call scheduler.step() after each epoch

# ReduceLROnPlateau: reduce LR when validation loss stops improving
scheduler = ReduceLROnPlateau(optimizer, patience=5, factor=0.5, min_lr=1e-6)
# scheduler.step(val_loss)  -  tracks metric

# Learning rate finder - find the steepest loss descent slope
lrs, losses = [], []
lr = 1e-7
for batch_x, batch_y in loader:
    optimizer.param_groups[0]['lr'] = lr
    optimizer.zero_grad()
    loss = criterion(model(batch_x), batch_y)
    loss.backward()
    optimizer.step()
    lrs.append(lr); losses.append(loss.item())
    lr *= 1.1
    if lr > 10 or loss.item() > 10 * losses[0]:
        break
# Plot losses vs. lrs - pick the LR just before the minimum

Batch Normalization

BatchNorm normalizes activations within each mini-batch, then scales and shifts with learned parameters γ and β. It dramatically stabilizes training by keeping activations in a healthy range.

pythonCopy
class BNBlock(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.linear = nn.Linear(dim, dim)
        self.bn = nn.BatchNorm1d(dim)  # after linear, before activation
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.bn(self.linear(x)))

BatchNorm has different behavior during training (uses batch statistics) and inference (uses running mean/variance accumulated during training). Always call model.train() during training and model.eval() during evaluation - forgetting this is a common source of production bugs where model performance is worse in deployment than in validation.

Layer normalization (used in transformers): normalizes across features instead of across the batch. Required when batch size is small or varies (e.g., sequences of different lengths).

pythonCopy
# For Transformers and RNNs
self.norm = nn.LayerNorm(d_model)
x = self.norm(x)  # no batch statistics - works at batch size 1

Dropout for Regularization

Dropout randomly zeros out neurons during training (with probability p), forcing the network to learn redundant representations. At inference, all neurons are active and outputs are scaled by (1-p).

pythonCopy
class RegularizedMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout=0.3):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(p=dropout),  # applied after activation
            nn.Linear(hidden_dim, hidden_dim // 2),
            nn.ReLU(),
            nn.Dropout(p=dropout),
            nn.Linear(hidden_dim // 2, output_dim)
        )

Typical dropout rates: 0.1–0.3 for hidden layers, 0.5 for the final layer in high-capacity networks. For transformers: 0.1 is standard.

Weight Decay (L2 Regularization)

Weight decay penalizes large parameter values, preventing overfitting. In AdamW it is decoupled from the adaptive learning rate:

pythonCopy
optimizer = optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)
# weight_decay = λ in L2 penalty: loss + λ||θ||²

Typical values: 1e-4 to 1e-2. Higher weight decay is needed for smaller datasets.

Gradient Clipping

Prevents exploding gradients in deep networks or RNNs:

pythonCopy
for epoch in range(n_epochs):
    for batch_x, batch_y in loader:
        optimizer.zero_grad()
        loss = criterion(model(batch_x), batch_y)
        loss.backward()
        # Clip gradient norm before stepping
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()

Use max_norm=1.0 as a default. If training is stable, you can increase it. If loss still diverges, reduce the learning rate.

The Complete Training Loop with Best Practices

pythonCopy
import torch
from torch.cuda.amp import autocast, GradScaler

device = 'cuda' if torch.cuda.is_available() else 'mps' if torch.backends.mps.is_available() else 'cpu'
model = model.to(device)
scaler = GradScaler()  # for mixed precision training (2x speedup on GPU)

best_val_loss = float('inf')
patience_counter = 0
PATIENCE = 10

for epoch in range(n_epochs):
    # Training phase
    model.train()
    train_loss = 0
    for batch_x, batch_y in train_loader:
        batch_x, batch_y = batch_x.to(device), batch_y.to(device)
        optimizer.zero_grad()

        with autocast():  # mixed precision: use float16 where safe
            logits = model(batch_x)
            loss = criterion(logits, batch_y)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        scaler.step(optimizer)
        scaler.update()
        train_loss += loss.item()

    # Validation phase
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch_x, batch_y in val_loader:
            batch_x, batch_y = batch_x.to(device), batch_y.to(device)
            logits = model(batch_x)
            val_loss += criterion(logits, batch_y).item()

    val_loss /= len(val_loader)
    scheduler.step(val_loss)  # ReduceLROnPlateau

    # Early stopping
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        torch.save(model.state_dict(), 'best_model.pt')
        patience_counter = 0
    else:
        patience_counter += 1
        if patience_counter >= PATIENCE:
            print(f"Early stopping at epoch {epoch}")
            break

# Load best checkpoint
model.load_state_dict(torch.load('best_model.pt'))

Diagnosing Training Failures

Loss diverges: LR too high. Reduce by 10x. Check for NaN in inputs (always validate data before training).

Loss decreases then plateaus on training set: LR too low, or vanishing gradients (add BatchNorm / check activations are not saturated).

Training loss low, validation loss high: Overfitting. Increase dropout, weight decay, or reduce model capacity. Get more training data.

Training loss barely decreases: Underfitting. Increase model capacity, reduce regularization, train longer.

Loss oscillates without decreasing: LR too high for fine-grained convergence. Use a schedule that decays LR over time.

Common Mistakes and Bad Instincts

Not calling model.eval() during validation. BatchNorm and Dropout behave differently in train vs. eval mode. Forgetting this makes validation metrics incorrect.

Zeroing gradients after the backward pass instead of before. Call optimizer.zero_grad() at the start of each iteration, not after optimizer.step(). The convention exists to allow gradient accumulation - once you understand it, you won't make this mistake.

Using a constant learning rate. Almost every deep learning paper uses a learning rate schedule. Cosine annealing or OneCycleLR consistently outperforms a fixed LR with negligible implementation cost.

Forgetting to move data to the same device as the model. RuntimeError: Expected all tensors to be on the same device is one of the most common PyTorch errors. Always batch_x.to(device) inside the loop.

Not using early stopping. Without early stopping, training runs until the epoch count - which is either too few (underfit) or too many (overfit). Monitor validation loss and stop when it stops improving.

Where to Go Next

• Module 18 (CNNs) applies these training techniques to convolutional architectures for image data.
• Module 20 (Transformers) shows how these same principles apply to attention-based sequence models.
• The post gradient-descent-from-scratch goes deeper on optimizer mathematics.