PyTorch Optimization and Debugging: Loss Curves, Overfitting and Speed

Common Training Patterns

Training a model is an experiment. When it fails, you need a debugging process. Is the data wrong? Are labels misencoded? Is the learning rate too high? Is the model too small? Is validation leaking? Is the loss function mismatched?

Strong PyTorch developers debug from simple to complex. They overfit one batch, inspect shapes and gradients, check label ranges, compare train and validation curves, and only then add advanced tricks.

If training loss does not decrease, suspect learning rate, model output shape, loss function, frozen parameters, bad labels, or missing optimizer step. If training loss decreases but validation worsens, suspect overfitting, data split issues, or distribution shift.

Overfit one batch to prove the model and loop can learn.
Plot train and validation loss.
Inspect gradient norms for vanishing or exploding gradients.
Use weight decay, dropout, augmentation, or early stopping for overfitting.

Performance Improvements

Once correctness is proven, improve speed with larger batches, pinned memory, multiple dataloader workers, mixed precision, and avoiding unnecessary CPU-GPU transfers.

Use mixed precision on compatible GPUs.
Avoid calling .item() too often inside hot loops.
Profile before optimizing complex code.

Read Loss Curves Before Changing Hyperparameters

A falling training loss with worsening validation loss suggests overfitting, while both losses remaining high suggests underfitting, weak features, optimization trouble, or incorrect targets. A sudden NaN points to numerical instability, invalid data, excessive learning rate, or exploding gradients. Plot the same metric and aggregation for both splits before comparing them.

Profile the Pipeline in Measured Stages

Separate data loading, host-to-device transfer, forward pass, loss, backward pass, optimizer step, synchronization, and validation. GPU operations are asynchronous, so use the profiler or explicit synchronization around a diagnostic measurement. Record batch size, precision, device, model mode, and warm-up before comparing changes.

Overfit One Batch Test

If a model cannot overfit one small batch, fix the data, loss, model, or loop before training on the full dataset.

Overfit One Batch Test

def overfit_one_batch(model, loader, loss_fn, optimizer, device, steps=200):
    model.train()
    features, labels = next(iter(loader))
    features = features.to(device)
    labels = labels.to(device)

    for step in range(steps):
        logits = model(features)
        loss = loss_fn(logits, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if step % 25 == 0:
            acc = (logits.argmax(dim=1) == labels).float().mean().item()
            print(f"step={step} loss={loss.item():.4f} acc={acc:.3f}")

A healthy model should drive training loss very low on one batch.
If it cannot, do not waste time on full training runs yet.

Mixed Precision Skeleton

Mixed precision can speed up training on modern GPUs while keeping model quality stable.

Mixed Precision Skeleton

scaler = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())

for features, labels in train_loader:
    features = features.to(device)
    labels = labels.to(device)

    optimizer.zero_grad()
    with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
        logits = model(features)
        loss = loss_fn(logits, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Use mixed precision after the normal training loop is correct.
Some operations may still require full precision; test metrics carefully.

Before you move on