Training & Optimization

Master the art and science of training deep learning models: hyperparameter tuning, optimizers, learning rate schedules, regularization, and hardware acceleration.

The Training Process

Training a deep learning model is an iterative process of minimizing a loss function. Each complete pass through the training data is called an epoch, and each epoch consists of many iterations (mini-batches):

1. Forward pass

   Pass a batch of data through the network to compute predictions.

2. Loss computation

   Compare predictions to true labels using a loss function.

3. Backward pass

   Compute gradients of the loss with respect to all parameters.

4. Parameter update

   Update weights using the optimizer.

5. Repeat

   Continue until the model converges or a stopping criterion is met.

Hyperparameter Tuning

Hyperparameters are settings you choose before training begins. The most important ones:

Optimizers

Optimizers determine how weights are updated based on gradients. Each has different properties:

• SGD (Stochastic Gradient Descent): The simplest optimizer. With momentum, it accumulates gradients from previous steps to smooth updates and escape local minima. Often yields the best final performance but requires careful tuning.
• Adam (Adaptive Moment Estimation): Maintains per-parameter learning rates using running averages of gradients and squared gradients. Converges faster than SGD and is less sensitive to the initial learning rate. The most popular default choice.
• AdamW: Adam with decoupled weight decay. Fixes a subtle issue in Adam's weight decay implementation. Preferred for training Transformers (used in BERT, GPT, etc.).

import torch.optim as optim

# SGD with momentum
optimizer = optim.SGD(model.parameters(), lr=0.01,
                      momentum=0.9, weight_decay=1e-4)

# Adam (most common default)
optimizer = optim.Adam(model.parameters(), lr=3e-4)

# AdamW (best for Transformers)
optimizer = optim.AdamW(model.parameters(), lr=3e-4,
                        weight_decay=0.01)

Learning Rate Schedulers

The learning rate often needs to change during training. Starting high and decreasing allows fast initial progress followed by fine-grained optimization:

from torch.optim.lr_scheduler import (
    StepLR, CosineAnnealingLR, OneCycleLR
)

# Step decay: multiply LR by 0.1 every 30 epochs
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)

# Cosine annealing: smooth decay following a cosine curve
scheduler = CosineAnnealingLR(optimizer, T_max=100)

# One-cycle: warm up then decay (great for fast training)
scheduler = OneCycleLR(optimizer, max_lr=0.01,
                       total_steps=1000)

# In training loop:
for epoch in range(num_epochs):
    train_one_epoch()
    scheduler.step()  # Update learning rate

Regularization Techniques

Regularization prevents overfitting - when the model memorizes training data instead of learning generalizable patterns:

• Dropout: Randomly sets a fraction of neurons to zero during training. Forces the network to learn redundant representations. Typically 0.1-0.5 for hidden layers.
• Weight decay (L2 regularization): Adds a penalty proportional to the squared magnitude of weights. Encourages smaller weights and simpler models.
• Batch normalization: Normalizes layer inputs to have zero mean and unit variance. Stabilizes training, allows higher learning rates, and acts as a mild regularizer.
• Layer normalization: Similar to batch norm but normalizes across features instead of across the batch. Preferred for Transformers and RNNs.

Data Augmentation

Data augmentation artificially increases the training set by applying random transformations to the input data. This is one of the most effective regularization techniques:

from torchvision import transforms

# Image augmentation pipeline
train_transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(15),
    transforms.RandomResizedCrop(224, scale=(0.8, 1.0)),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# No augmentation for validation!
val_transform = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

Early Stopping

Early stopping monitors validation loss during training and stops when it stops improving, preventing overfitting:

best_val_loss = float('inf')
patience = 5
patience_counter = 0

for epoch in range(100):
    train_loss = train_one_epoch()
    val_loss = evaluate()

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

GPU/TPU Training

Deep learning training benefits enormously from hardware acceleration:

• GPUs: NVIDIA GPUs (A100, H100, RTX 4090) provide massive parallelism for matrix operations. Moving your model and data to GPU can speed up training 10-100x.
• TPUs: Google's Tensor Processing Units, optimized for tensor operations. Available through Google Cloud and Colab.
• Multi-GPU training: Distribute training across multiple GPUs using data parallelism (split batches) or model parallelism (split the model).

Mixed Precision Training

Mixed precision uses both 16-bit (FP16) and 32-bit (FP32) floating point numbers during training. FP16 uses half the memory and is faster on modern GPUs, while critical operations remain in FP32 for numerical stability:

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for images, labels in train_loader:
    optimizer.zero_grad()

    # Forward pass in FP16
    with autocast():
        outputs = model(images.cuda())
        loss = criterion(outputs, labels.cuda())

    # Backward pass with gradient scaling
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()