Training techniques

🎯 Goal

After reading this chapter:

• You will know techniques to efficiently train neural networks
• You will use tools to combat overfitting (Dropout, BatchNorm, regularization)
• You will be able to apply learning rate scheduling, gradient clipping, mixed precision
• You will get good results on small datasets too with transfer learning

What to learn

• Regularization: L1/L2 (weight decay), Dropout, BatchNorm, LayerNorm
• Initialization: Xavier (Glorot), He, Kaiming
• Optimizersin depth: SGD+momentum, Adam, AdamW, LAMB
• Learning rate scheduling: StepLR, CosineAnnealingLR, OneCycleLR, ReduceLROnPlateau
• Gradient clipping — protection from gradient explosion
• Mixed precision training(FP16/BF16) — faster + less memory
• Data augmentation — artificially expanding the dataset
• Transfer learning — reusing pretrained models
• Early stopping and checkpointing
• Weights & Biases / TensorBoard — experiment tracking

Libraries

pip install torch torchvision wandb tensorboard

Important topics

Regularization techniques

Dropout

Randomly “turning off” neurons during training — preventing overfitting.

import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.dropout = nn.Dropout(p=0.5)  # 50% of neurons turned off
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.dropout(x)
        return x

# In eval mode, dropout is automatically disabled (when `.eval()` is called)

Batch Normalization

Normalizing activations per batch — faster convergence + regularization effect.

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.bn1 = nn.BatchNorm1d(256)  # 1D BN (for MLP)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.bn1(x)
        x = torch.relu(x)
        return x

# For CNN: nn.BatchNorm2d # For Transformer: nn.LayerNorm (LayerNorm fits better)

Weight Decay (L2)

The weight_decay parameter in the optimizer.

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)

Learning Rate Scheduling

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

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

# Variant 1: Step decay (decays by gamma every N epochs)
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)

# Variant 2: Cosine annealing (smooth decay)
scheduler = CosineAnnealingLR(optimizer, T_max=EPOCHS)

# Variant 3: OneCycleLR (warmup + decay) — Karpathy's favorite
scheduler = OneCycleLR(optimizer, max_lr=1e-2, total_steps=EPOCHS * len(train_loader))

# Variant 4: ReduceLROnPlateau (when val loss doesn't improve)
scheduler = ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=3)

# In training loop
for epoch in range(EPOCHS):
    train_one_epoch(...)
    scheduler.step()         # At end of epoch (or for ReduceLROnPlateau: scheduler.step(val_loss))

Gradient Clipping

Training doesn’t “explode” when gradients become too large:

loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()

Especially needed in RNN/LSTMand Transformertraining.

Mixed Precision Training

Reduces GPU memory by 2x, increases speed by 2-3x.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for X, y in loader:
    X, y = X.cuda(), y.cuda()
    
    with autocast(dtype=torch.float16):
        logits = model(X)
        loss = criterion(logits, y)
    
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()
    optimizer.zero_grad()

Data Augmentation (for Images)

from torchvision import transforms

train_transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.RandomCrop(224),
    transforms.RandomHorizontalFlip(),
    transforms.ColorJitter(brightness=0.2, contrast=0.2),
    transforms.RandomRotation(15),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# NO augmentation for test
test_transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

Transfer Learning

import torchvision.models as models

# Pretrained ResNet-18
model = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)

# Variant 1: Retrain only the last layer (feature extraction)
for param in model.parameters():
    param.requires_grad = False  # freeze all

model.fc = nn.Linear(model.fc.in_features, num_classes)  # new classifier # Only model.fc.parameters() will be trained

# Variant 2: Fine-tuning (train all, with small LR)
optimizer = torch.optim.AdamW([
    {"params": model.layer1.parameters(), "lr": 1e-5},  # old layers — low LR
    {"params": model.layer4.parameters(), "lr": 1e-4},
    {"params": model.fc.parameters(), "lr": 1e-3},      # new layer — high LR
])

Code examples

Complete training pipeline (production-ready)

import torch
import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnnealingLR
from torch.cuda.amp import autocast, GradScaler

def train_model(
    model, train_loader, val_loader,
    epochs=20, lr=1e-3, weight_decay=1e-4,
    grad_clip=1.0, use_amp=True,
    save_path="best.pt",
):
    device = next(model.parameters()).device
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
    scheduler = CosineAnnealingLR(optimizer, T_max=epochs)
    scaler = GradScaler() if use_amp else None
    
    best_val_acc = 0
    
    for epoch in range(epochs):
        # Train
        model.train()
        train_loss = 0
        for X, y in train_loader:
            X, y = X.to(device), y.to(device)
            optimizer.zero_grad()
            
            if use_amp:
                with autocast():
                    logits = model(X)
                    loss = criterion(logits, y)
                scaler.scale(loss).backward()
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                scaler.step(optimizer)
                scaler.update()
            else:
                logits = model(X)
                loss = criterion(logits, y)
                loss.backward()
                torch.nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
                optimizer.step()
            
            train_loss += loss.item()
        
        scheduler.step()
        
        # Validate
        model.eval()
        val_correct = 0
        val_total = 0
        with torch.no_grad():
            for X, y in val_loader:
                X, y = X.to(device), y.to(device)
                logits = model(X)
                val_correct += (logits.argmax(dim=1) == y).sum().item()
                val_total += y.size(0)
        
        val_acc = val_correct / val_total
        print(f"Epoch{epoch+1}/{epochs}"
              f"train_loss={train_loss/len(train_loader):.4f}"
              f"val_acc={val_acc:.4f}"
              f"lr={optimizer.param_groups[0]['lr']:.6f}")
        
        # Save best
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), save_path)
    
    return best_val_acc

Weights & Biases integration

import wandb

wandb.init(project="my-ml-project", config={
    "lr": 1e-3,
    "batch_size": 64,
    "epochs": 20,
    "architecture": "ResNet-18",
})

# Inside training loop
wandb.log({
    "train_loss": train_loss,
    "val_acc": val_acc,
    "lr": optimizer.param_groups[0]["lr"],
}, step=epoch)

wandb.finish()

TensorBoard integration

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter("runs/experiment_1")

for epoch in range(epochs):
    # ... training ...
    writer.add_scalar("Loss/train", train_loss, epoch)
    writer.add_scalar("Accuracy/val", val_acc, epoch)
    writer.add_histogram("fc.weights", model.fc.weight, epoch)

writer.close()

# $ tensorboard --logdir=runs

Transfer Learning full example

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 1. Pretrained ResNet
model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)

# Freeze backbone
for param in model.parameters():
    param.requires_grad = False

# New classifier (10 classes for diseases)
model.fc = nn.Sequential(
    nn.Linear(model.fc.in_features, 512),
    nn.ReLU(),
    nn.Dropout(0.3),
    nn.Linear(512, 10),
)

# 2. Optimizer for fc parameters only
optimizer = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)

# 3. Train (only fc)
train_model(model, train_loader, val_loader, epochs=5, lr=1e-3)

# 4. Unfreeze and fine-tune (small LR)
for param in model.parameters():
    param.requires_grad = True

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
train_model(model, train_loader, val_loader, epochs=10, lr=1e-5)

Backend integration

Training service (background job)

from celery import Celery
import torch

celery_app = Celery("training", broker="redis://localhost:6379")

@celery_app.task(bind=True)
def train_model_task(self, dataset_path, hyperparams):
    # Progress tracking
    def on_epoch_end(epoch, val_acc):
        self.update_state(
            state="PROGRESS",
            meta={"epoch": epoch, "val_acc": val_acc},
        )
    
    model = create_model()
    train_loader, val_loader = create_loaders(dataset_path, hyperparams["batch_size"])
    
    best_acc = train_model(model, train_loader, val_loader, **hyperparams, 
                            on_epoch_end=on_epoch_end)
    
    # Save to S3 or local
    model_path = f"models/run_{self.request.id}.pt"
    torch.save(model.state_dict(), model_path)
    
    return {"best_acc": best_acc, "model_path": model_path}

# FastAPI endpoint
@app.post("/train")
def start_training(dataset_path: str, epochs: int = 20):
    task = train_model_task.delay(dataset_path, {"epochs": epochs, "batch_size": 64, "lr": 1e-3})
    return {"task_id": task.id}

@app.get("/train/{task_id}")
def get_training_status(task_id: str):
    task = train_model_task.AsyncResult(task_id)
    return {
        "state": task.state,
        "info": task.info if task.info else {},
    }

Resources

• PyTorch tutorials — Training techniques(link)
• “Bag of Tricks for Image Classification with CNNs” — paper (training improvements)
• Andrej Karpathy — “A Recipe for Training Neural Networks”(blog)
• Weights & Biases — Best Practicescourses
• OneCycleLR — Leslie Smith paper

🏋️ Exercises

🟢 Easy

1. Add Dropout in MLP, observe the difference between train and val accuracy.
2. Compare Adam and SGD on the same model.
3. Add ReduceLROnPlateau, visualize the plateau.

🟡 Medium

1. Mixed precision: Run the same training with FP32 and AMP, observe time and memory difference.
2. Augmentation: Compare a simple CNN with and without augmentation (CIFAR-10).
3. Transfer learning: Get 90%+ accuracy on a small dataset of 100 images using pretrained ResNet.

🔴 Hard

1. Custom LR scheduler: Write a scheduler combining warmup + cosine annealing.
2. Hyperparameter sweep: 50 trials with Optuna or wandb sweeps, find the best configuration.
3. Production training service: Celery + FastAPI + S3 + W&B — complete pipeline.

Capstone

notebooks/month-03/04_training_techniques.ipynb:

• Compare 2 variants on the CIFAR-10 dataset:
• Baseline: simple CNN, Adam, no augmentation
• Improved: BatchNorm + Dropout + augmentation + OneCycleLR + AMP
• Logs in Wandb or TensorBoard
• Test accuracy: baseline ~70%, improved 85%+

✅ Checklist

• I know when to use Dropout, BatchNorm
• I know the difference between Adam vs AdamW (weight decay)
• I know learning rate scheduling types
• I know when gradient clipping is needed
• I can apply mixed precision training (AMP)
• I use data augmentation (vision)
• I get good results on small datasets with transfer learning
• I do experiment tracking with W&B or TensorBoard

Moving on to CNN — Convolutional Networks.