PyTorch GPU Training Optimization: 10 Tips to Speed Up Your Models

import torch
from torch.cuda.amp import GradScaler, autocast

# Initialize the scaler
scaler = GradScaler()

# Training loop with mixed precision
for batch_idx, (data, target) in enumerate(train_loader):
    optimizer.zero_grad()
    
    # Forward pass with autocast
    with autocast():
        output = model(data)
        loss = criterion(output, target)
    
    # Backward pass with scaling
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

def find_optimal_batch_size(model, input_shape, max_batch_size=512):
    """Find the largest batch size that fits in GPU memory"""
    batch_size = 2
    
    while batch_size <= max_batch_size:
        try:
            # Create dummy batch
            dummy_input = torch.randn(batch_size, *input_shape).cuda()
            
            # Forward pass
            with torch.no_grad():
                output = model(dummy_input)
            
            print(f"Batch size {batch_size}: OK")
            batch_size *= 2
            
        except RuntimeError as e:
            if "out of memory" in str(e):
                optimal_batch_size = batch_size // 2
                print(f"Optimal batch size: {optimal_batch_size}")
                return optimal_batch_size
            else:
                raise e
    
    return max_batch_size

# Simulate batch size of 128 with accumulation_steps=4 and batch_size=32
accumulation_steps = 4
effective_batch_size = 128

for batch_idx, (data, target) in enumerate(train_loader):
    # Forward pass
    with autocast():
        output = model(data)
        loss = criterion(output, target) / accumulation_steps
    
    # Backward pass
    scaler.scale(loss).backward()
    
    # Update weights every accumulation_steps
    if (batch_idx + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad()

from torch.utils.data import DataLoader

# Optimized DataLoader settings
train_loader = DataLoader(
    dataset=train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=8,  # Use 4-8 workers for optimal performance
    pin_memory=True,  # Faster GPU transfer
    persistent_workers=True,  # Keep workers alive between epochs
    prefetch_factor=2,  # Prefetch 2 batches per worker
    drop_last=True  # Avoid small last batch
)

# For even better performance, use a custom collate function
def fast_collate(batch):
    """Optimized collate function for image data"""
    imgs = [img[0] for img in batch]
    targets = torch.tensor([img[1] for img in batch], dtype=torch.int64)
    
    # Stack images efficiently
    w, h = imgs[0].size
    tensor = torch.zeros((len(imgs), 3, h, w), dtype=torch.uint8)
    
    for i, img in enumerate(imgs):
        numpy_array = np.asarray(img, dtype=np.uint8)
        tensor[i] = torch.from_numpy(numpy_array.transpose((2, 0, 1)))
    
    return tensor, targets

import torch.utils.checkpoint as checkpoint

class CheckpointedResNet(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model
    
    def forward(self, x):
        # Apply checkpointing to each residual block
        for i, layer in enumerate(self.model.children()):
            if i > 0:  # Skip first layer (usually conv + bn + relu)
                x = checkpoint.checkpoint(layer, x)
            else:
                x = layer(x)
        return x

# Enable checkpointing for transformers
from transformers import AutoModel
model = AutoModel.from_pretrained("bert-base-uncased")
model.gradient_checkpointing_enable()

# 1. Built-in PyTorch profiler
from torch.profiler import profile, record_function, ProfilerActivity

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for batch_idx, (data, target) in enumerate(train_loader):
        if batch_idx >= 10:  # Profile first 10 batches
            break
        
        with record_function("forward"):
            output = model(data)
            loss = criterion(output, target)
        
        with record_function("backward"):
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()

# Export profiling results
prof.export_chrome_trace("trace.json")

# 2. Simple GPU utilization monitoring
import GPUtil

def monitor_gpu():
    gpus = GPUtil.getGPUs()
    for gpu in gpus:
        print(f"GPU {gpu.id}: {gpu.load*100:.1f}% | Memory: {gpu.memoryUtil*100:.1f}%")

# 3. Real-time monitoring during training
import threading
import time

def gpu_monitor_thread():
    while training:
        monitor_gpu()
        time.sleep(5)

monitor_thread = threading.Thread(target=gpu_monitor_thread)
monitor_thread.start()