Part 4: Training and Optimization

Part of the PyTorch 101 Series

My First Training Loop

I trained my first model on a Tuesday afternoon. Left it running, checked the next morning: 89% accuracy!

Felt like magic.

Then on real data: 23% accuracy. Worse than random guessing (25% for 4 classes).

What went wrong? Everything about training:

• Wrong learning rate (too high)

• Wrong optimizer (basic SGD)

• No learning rate scheduling

• Data loading bottleneck

• No validation split

Spent a week debugging. Now I know training is 80% of the work. The model architecture is easy - training it well is hard.

Let me share what I learned.

The Training Loop

Every PyTorch training loop follows this pattern:

import torch
import torch.nn as nn
import torch.optim as optim

# Model, loss, optimizer
model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Training loop
for epoch in range(num_epochs):
    for batch_idx, (data, targets) in enumerate(train_loader):
        # Forward pass
        predictions = model(data)
        loss = criterion(predictions, targets)
        
        # Backward pass
        optimizer.zero_grad()  # Clear gradients
        loss.backward()        # Compute gradients
        optimizer.step()       # Update weights

This is the foundation. Everything else builds on it.

Complete Training Example

Here's my actual training code (simplified):

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

def train_one_epoch(model, train_loader, criterion, optimizer, device):
    """Train for one epoch."""
    model.train()
    
    running_loss = 0.0
    correct = 0
    total = 0
    
    for batch_idx, (data, targets) in enumerate(train_loader):
        # Move to device
        data = data.to(device)
        targets = targets.to(device)
        
        # Forward pass
        predictions = model(data)
        loss = criterion(predictions, targets)
        
        # Backward pass
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        
        # Statistics
        running_loss += loss.item()
        _, predicted = predictions.max(1)
        total += targets.size(0)
        correct += predicted.eq(targets).sum().item()
        
        # Print progress
        if (batch_idx + 1) % 100 == 0:
            print(f'Batch {batch_idx + 1}/{len(train_loader)} | '
                  f'Loss: {loss.item():.4f} | '
                  f'Acc: {100. * correct / total:.2f}%')
    
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100. * correct / total
    
    return epoch_loss, epoch_acc

def validate(model, val_loader, criterion, device):
    """Validate model."""
    model.eval()
    
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():  # No gradients needed
        for data, targets in val_loader:
            data = data.to(device)
            targets = targets.to(device)
            
            predictions = model(data)
            loss = criterion(predictions, targets)
            
            running_loss += loss.item()
            _, predicted = predictions.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
    
    val_loss = running_loss / len(val_loader)
    val_acc = 100. * correct / total
    
    return val_loss, val_acc

# Main training loop
def train_model(model, train_loader, val_loader, num_epochs=10):
    """Complete training pipeline."""
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    model = model.to(device)
    
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    
    best_val_acc = 0.0
    
    for epoch in range(num_epochs):
        print(f'\nEpoch {epoch + 1}/{num_epochs}')
        print('-' * 60)
        
        # Train
        train_loss, train_acc = train_one_epoch(
            model, train_loader, criterion, optimizer, device
        )
        
        # Validate
        val_loss, val_acc = validate(model, val_loader, criterion, device)
        
        print(f'\nTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%')
        print(f'Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%')
        
        # Save best model
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            torch.save(model.state_dict(), 'best_model.pth')
            print(f'Saved best model (Val Acc: {val_acc:.2f}%)')
    
    return model

This is production-ready code I use in all my projects.

Data Loading

Critical for training speed. Bad data loading = GPU sitting idle.

Creating Dataset

from torch.utils.data import Dataset, DataLoader
import torch

class CustomDataset(Dataset):
    """Custom dataset class."""
    
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        label = self.labels[idx]
        
        if self.transform:
            sample = self.transform(sample)
        
        return sample, label

# Create dataset
data = torch.randn(1000, 3, 224, 224)  # 1000 images
labels = torch.randint(0, 10, (1000,))  # 10 classes

dataset = CustomDataset(data, labels)

DataLoader

from torch.utils.data import DataLoader

train_loader = DataLoader(
    dataset,
    batch_size=32,
    shuffle=True,          # Shuffle training data
    num_workers=4,         # Parallel data loading
    pin_memory=True        # Faster GPU transfer
)

# Iterate
for batch_data, batch_labels in train_loader:
    # Train on batch
    pass

I learned the hard way: num_workers=0 (single process) made my training 5x slower.

Real Image Dataset

from torch.utils.data import Dataset
from PIL import Image
import os

class ImageDataset(Dataset):
    """Load images from directory."""
    
    def __init__(self, image_dir, transform=None):
        self.image_dir = image_dir
        self.transform = transform
        self.images = os.listdir(image_dir)
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.images[idx])
        image = Image.open(img_path).convert('RGB')
        
        # Extract label from filename (e.g., "cat_001.jpg")
        label = 0 if 'cat' in self.images[idx] else 1
        
        if self.transform:
            image = self.transform(image)
        
        return image, label

# Transforms
from torchvision import transforms

train_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                       std=[0.229, 0.224, 0.225])
])

dataset = ImageDataset('data/images', transform=train_transforms)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

Optimizers

The optimizer updates model weights. Choice matters!

SGD (Stochastic Gradient Descent)

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

Basic but effective with momentum.

Adam

optimizer = optim.Adam(
    model.parameters(),
    lr=0.001,
    betas=(0.9, 0.999),
    weight_decay=1e-4
)

My default choice. Adaptive learning rates, works well most of the time.

AdamW

optimizer = optim.AdamW(
    model.parameters(),
    lr=0.001,
    weight_decay=0.01
)

Better weight decay than Adam. I use this for transformers.

Comparing Optimizers

# SGD with momentum - good for CNNs
optimizer_sgd = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)

# Adam - general purpose
optimizer_adam = optim.Adam(model.parameters(), lr=0.001)

# AdamW - better regularization
optimizer_adamw = optim.AdamW(model.parameters(), lr=0.001, weight_decay=0.01)

# RMSprop - good for RNNs
optimizer_rmsprop = optim.RMSprop(model.parameters(), lr=0.001)

Rule of thumb:

• CNNs: SGD with momentum or Adam

• Transformers: AdamW

• RNNs: Adam or RMSprop

• When in doubt: Adam

Learning Rate Scheduling

Fixed learning rate rarely works best. I use schedulers in every project.

Step Decay

from torch.optim.lr_scheduler import StepLR

optimizer = optim.Adam(model.parameters(), lr=0.001)
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)

for epoch in range(num_epochs):
    train_one_epoch(model, train_loader, criterion, optimizer, device)
    scheduler.step()  # Decay learning rate
    
    print(f"Epoch {epoch + 1}, LR: {optimizer.param_groups[0]['lr']}")

Reduces LR every step_size epochs by factor gamma.

Cosine Annealing

from torch.optim.lr_scheduler import CosineAnnealingLR

scheduler = CosineAnnealingLR(optimizer, T_max=100, eta_min=1e-6)

# Smoothly decreases LR following cosine curve

My favorite for long training. Smooth decay, good results.

Reduce on Plateau

from torch.optim.lr_scheduler import ReduceLROnPlateau

scheduler = ReduceLROnPlateau(
    optimizer,
    mode='min',
    factor=0.1,
    patience=5,
    verbose=True
)

for epoch in range(num_epochs):
    train_loss = train_one_epoch(...)
    val_loss = validate(...)
    
    scheduler.step(val_loss)  # Monitor validation loss

Reduces LR when validation loss stops improving. I use this when I don't know optimal schedule.

One Cycle Policy

from torch.optim.lr_scheduler import OneCycleLR

optimizer = optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = OneCycleLR(
    optimizer,
    max_lr=0.1,
    steps_per_epoch=len(train_loader),
    epochs=num_epochs
)

for epoch in range(num_epochs):
    for batch in train_loader:
        train_step(...)
        scheduler.step()  # Step per batch, not epoch!

Fast training with super-convergence. Got me 30% faster training on image classification.

Warm-up Schedule

from torch.optim.lr_scheduler import LambdaLR

def warmup_lambda(epoch):
    if epoch < 5:
        return (epoch + 1) / 5  # Linear warmup
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=warmup_lambda)

I use warm-up for large models - prevents early instability.

Handling Common Training Issues

1. Exploding Gradients

import torch.nn as nn

# Gradient clipping
max_grad_norm = 1.0

for batch in train_loader:
    loss = train_step(...)
    loss.backward()
    
    # Clip gradients
    nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
    
    optimizer.step()

Essential for RNNs/LSTMs.

2. Class Imbalance

import torch.nn as nn

# Weighted loss
class_weights = torch.tensor([1.0, 5.0])  # Weight minority class more
criterion = nn.CrossEntropyLoss(weight=class_weights)

3. Overfitting

# Regularization techniques
model = nn.Sequential(
    nn.Linear(100, 50),
    nn.Dropout(0.5),     # Dropout
    nn.Linear(50, 10)
)

# L2 regularization via optimizer
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-4)

# Data augmentation
transforms = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.RandomRotation(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2)
])

# Early stopping
best_val_loss = float('inf')
patience = 10
patience_counter = 0

for epoch in range(num_epochs):
    val_loss = validate(...)
    
    if val_loss < best_val_loss:
        best_val_loss = val_loss
        patience_counter = 0
        save_model()
    else:
        patience_counter += 1
    
    if patience_counter >= patience:
        print("Early stopping")
        break

4. Slow Training

# Mixed precision training (PyTorch 1.6+)
from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for batch in train_loader:
    optimizer.zero_grad()
    
    # Forward in mixed precision
    with autocast():
        predictions = model(data)
        loss = criterion(predictions, targets)
    
    # Backward with scaling
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Got me 2-3x speedup on my image classifier with minimal code changes.

Production Training Pipeline

Here's my complete production setup:

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torch.cuda.amp import autocast, GradScaler
import wandb  # For experiment tracking

class Trainer:
    """Production training class."""
    
    def __init__(self, model, train_loader, val_loader, config):
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config
        
        # Device
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.model = self.model.to(self.device)
        
        # Loss and optimizer
        self.criterion = nn.CrossEntropyLoss()
        self.optimizer = optim.AdamW(
            model.parameters(),
            lr=config['lr'],
            weight_decay=config['weight_decay']
        )
        
        # Scheduler
        self.scheduler = optim.lr_scheduler.CosineAnnealingLR(
            self.optimizer,
            T_max=config['epochs']
        )
        
        # Mixed precision
        self.scaler = GradScaler()
        
        # Tracking
        self.best_val_acc = 0.0
        
    def train_epoch(self):
        """Train one epoch."""
        self.model.train()
        
        total_loss = 0
        correct = 0
        total = 0
        
        for data, targets in self.train_loader:
            data = data.to(self.device)
            targets = targets.to(self.device)
            
            self.optimizer.zero_grad()
            
            # Mixed precision forward
            with autocast():
                predictions = self.model(data)
                loss = self.criterion(predictions, targets)
            
            # Backward
            self.scaler.scale(loss).backward()
            
            # Gradient clipping
            self.scaler.unscale_(self.optimizer)
            nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
            
            self.scaler.step(self.optimizer)
            self.scaler.update()
            
            # Statistics
            total_loss += loss.item()
            _, predicted = predictions.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        avg_loss = total_loss / len(self.train_loader)
        accuracy = 100. * correct / total
        
        return avg_loss, accuracy
    
    @torch.no_grad()
    def validate(self):
        """Validate model."""
        self.model.eval()
        
        total_loss = 0
        correct = 0
        total = 0
        
        for data, targets in self.val_loader:
            data = data.to(self.device)
            targets = targets.to(self.device)
            
            predictions = self.model(data)
            loss = self.criterion(predictions, targets)
            
            total_loss += loss.item()
            _, predicted = predictions.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
        
        avg_loss = total_loss / len(self.val_loader)
        accuracy = 100. * correct / total
        
        return avg_loss, accuracy
    
    def train(self):
        """Complete training loop."""
        for epoch in range(self.config['epochs']):
            print(f"\nEpoch {epoch + 1}/{self.config['epochs']}")
            print("-" * 60)
            
            # Train
            train_loss, train_acc = self.train_epoch()
            
            # Validate
            val_loss, val_acc = self.validate()
            
            # Scheduler step
            self.scheduler.step()
            
            # Log
            print(f"Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
            print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
            print(f"LR: {self.optimizer.param_groups[0]['lr']:.6f}")
            
            # Save best model
            if val_acc > self.best_val_acc:
                self.best_val_acc = val_acc
                self.save_checkpoint('best_model.pth')
                print(f"✓ Saved best model (Val Acc: {val_acc:.2f}%)")
    
    def save_checkpoint(self, path):
        """Save model checkpoint."""
        torch.save({
            'model_state_dict': self.model.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'scheduler_state_dict': self.scheduler.state_dict(),
            'best_val_acc': self.best_val_acc,
        }, path)
    
    def load_checkpoint(self, path):
        """Load model checkpoint."""
        checkpoint = torch.load(path)
        self.model.load_state_dict(checkpoint['model_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        self.best_val_acc = checkpoint['best_val_acc']

# Usage
config = {
    'lr': 0.001,
    'weight_decay': 0.01,
    'epochs': 50
}

trainer = Trainer(model, train_loader, val_loader, config)
trainer.train()

This is what I use. Handles everything:

• Mixed precision training

• Gradient clipping

• Learning rate scheduling

• Checkpointing

• Proper train/val split

Monitoring Training

TensorBoard

from torch.utils.tensorboard import SummaryWriter

writer = SummaryWriter('runs/experiment_1')

for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(...)
    val_loss, val_acc = validate(...)
    
    # Log metrics
    writer.add_scalar('Loss/train', train_loss, epoch)
    writer.add_scalar('Loss/val', val_loss, epoch)
    writer.add_scalar('Accuracy/train', train_acc, epoch)
    writer.add_scalar('Accuracy/val', val_acc, epoch)
    writer.add_scalar('Learning_rate', optimizer.param_groups[0]['lr'], epoch)

writer.close()

Run TensorBoard:

tensorboard --logdir=runs

Weights & Biases

import wandb

# Initialize
wandb.init(project='my-project', config={
    'lr': 0.001,
    'epochs': 50,
    'batch_size': 32
})

# Log metrics
for epoch in range(num_epochs):
    train_loss, train_acc = train_epoch(...)
    val_loss, val_acc = validate(...)
    
    wandb.log({
        'train_loss': train_loss,
        'train_acc': train_acc,
        'val_loss': val_loss,
        'val_acc': val_acc,
        'lr': optimizer.param_groups[0]['lr']
    })

I use Weights & Biases for all experiments. Great for tracking hyperparameters and comparing runs.

Best Practices

From training hundreds of models:

1. Always use validation set:

# Split data
from sklearn.model_selection import train_test_split
train_data, val_data = train_test_split(data, test_size=0.2)

2. Monitor multiple metrics:

metrics = {
    'loss': loss.item(),
    'accuracy': accuracy,
    'precision': precision,
    'recall': recall,
    'f1': f1_score
}

3. Save checkpoints:

if val_acc > best_val_acc:
    torch.save(model.state_dict(), f'model_epoch_{epoch}.pth')

4. Use reproducible seeds:

import random
import numpy as np

seed = 42
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True

5. Profile training:

import time

start_time = time.time()
for epoch in range(num_epochs):
    epoch_start = time.time()
    train_epoch(...)
    epoch_time = time.time() - epoch_start
    print(f"Epoch time: {epoch_time:.2f}s")

total_time = time.time() - start_time
print(f"Total training time: {total_time / 60:.2f} minutes")

What I Learned

Training is an iterative process:

1. Start with reasonable defaults (Adam, lr=0.001)

2. Train a few epochs

3. Check for issues (overfitting, underfitting, slow convergence)

4. Adjust (learning rate, regularization, data augmentation)

5. Repeat

No single recipe works for everything. But the patterns above solve 90% of problems.

What's Next?

You now know how to train models effectively. In Part 5, we'll learn how to deploy these models to production.

Next: Part 5 - Production Deployment and Best Practices

---

Previous: Part 3 - Building Neural Networks with torch.nn

This article is part of the PyTorch 101 series. All examples use Python 3 and are based on real projects.