Custom Training Loops

Custom Training Loops

Writing custom PyTorch training loops beyond the standard Trainer.

Basic Custom Loop

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from tqdm import tqdm

def train_epoch(model, dataloader, optimizer, criterion, device, clip_grad=1.0):
    model.train()
    total_loss = 0

    for batch in tqdm(dataloader, desc="Training"):
        inputs, labels = [x.to(device) for x in batch]

        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()

        # Gradient clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_grad)

        optimizer.step()
        total_loss += loss.item()

    return total_loss / len(dataloader)

Training Loop with Scheduler

def train(model, train_loader, val_loader, config):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)

    optimizer = torch.optim.AdamW(model.parameters(), lr=config["lr"], weight_decay=0.01)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=config["epochs"])
    criterion = nn.CrossEntropyLoss()

    best_val_loss = float('inf')

    for epoch in range(config["epochs"]):
        train_loss = train_epoch(model, train_loader, optimizer, criterion, device)
        val_loss = evaluate(model, val_loader, criterion, device)

        scheduler.step()
        current_lr = scheduler.get_last_lr()[0]

        print(f"Epoch {epoch+1}: train_loss={train_loss:.4f}, val_loss={val_loss:.4f}, lr={current_lr:.2e}")

        # Save best model
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            torch.save(model.state_dict(), "best_model.pt")

    return model

Mixed Precision Training

from torch.cuda.amp import autocast, GradScaler

def train_amp(model, dataloader, optimizer, criterion, device):
    scaler = GradScaler()
    model.train()

    for batch in dataloader:
        inputs, labels = [x.to(device) for x in batch]

        optimizer.zero_grad()

        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()

Gradient Accumulation

def train_with_accumulation(model, dataloader, optimizer, criterion, device,
                            accumulation_steps=4):
    model.train()
    optimizer.zero_grad()

    for i, batch in enumerate(dataloader):
        inputs, labels = [x.to(device) for x in batch]
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss = loss / accumulation_steps  # normalize
        loss.backward()

        if (i + 1) % accumulation_steps == 0:
            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
            optimizer.step()
            optimizer.zero_grad()

Distributed Training (DDP)

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_ddp(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def train_ddp(rank, world_size, model, dataset):
    setup_ddp(rank, world_size)
    model = model.to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    loader = DataLoader(dataset, batch_size=32, sampler=sampler)

    optimizer = torch.optim.AdamW(ddp_model.parameters(), lr=0.001)

    for epoch in range(10):
        sampler.set_epoch(epoch)
        for batch in loader:
            inputs, labels = [x.to(rank) for x in batch]
            optimizer.zero_grad()
            loss = nn.CrossEntropyLoss()(ddp_model(inputs), labels)
            loss.backward()
            optimizer.step()

    dist.destroy_process_group()

Pitfalls

• Always call model.train() and model.eval() for different modes
• Zero gradients BEFORE each backward pass (or accumulate intentionally)
• Gradient clipping prevents explosion — critical for transformers
• Mixed precision: loss scaling avoids underflow for small gradients
• DDP: each process should have different batch (sampler shuffles differently)