Back to blog
• 4 min read • Deep Learning

PyTorch Performance Optimization Tips

Training deep learning models can be computationally expensive. These optimization techniques can significantly reduce training time and memory usage without sacrificing model quality.

Mixed Precision Training

Modern GPUs support half-precision (FP16) operations that are significantly faster than single-precision (FP32). PyTorch’s Automatic Mixed Precision (AMP) feature automatically chooses which operations to run in FP16:

from torch.cuda.amp import autocast, GradScaler

# Initialize gradient scaler for loss scaling
scaler = GradScaler()

for data, target in train_loader:
    data, target = data.cuda(), target.cuda()
    optimizer.zero_grad()

    # Automatic mixed precision context
    with autocast():
        output = model(data)
        loss = criterion(output, target)

    # Scale loss and backward
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Mixed precision can provide 2-3x speedup on modern GPUs with minimal accuracy impact. The GradScaler handles loss scaling to prevent gradient underflow in FP16.

Optimized Data Loading

Data loading is often the bottleneck in training pipelines. The DataLoader provides several optimization options:

from torch.utils.data import DataLoader

train_loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,        # Parallel data loading
    pin_memory=True,      # Faster CPU->GPU transfer
    persistent_workers=True,  # Keep workers alive between epochs
    prefetch_factor=2     # Prefetch batches per worker
)

Key parameters:

  • num_workers: Set to number of CPU cores for parallel loading
  • pin_memory: Speeds up host-to-device transfers
  • persistent_workers: Avoids worker spawn overhead each epoch

Gradient Accumulation

When GPU memory is limited, gradient accumulation allows effective larger batch sizes:

accumulation_steps = 4  # Effective batch_size * 4

for i, (data, target) in enumerate(train_loader):
    output = model(data)
    loss = criterion(output, target) / accumulation_steps
    loss.backward()

    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

This technique is essential when training large models on consumer GPUs.

Memory Optimization Techniques

Gradient Checkpointing

Trade computation for memory by recomputing activations during backward pass:

from torch.utils.checkpoint import checkpoint

class MemoryEfficientModel(nn.Module):
    def forward(self, x):
        # Only store checkpoints, recompute rest
        x = checkpoint(self.layer1, x)
        x = checkpoint(self.layer2, x)
        return x

This can reduce memory usage by 50-70% with ~20% computation overhead.

In-Place Operations

Use in-place operations carefully to save memory:

# Good: Saves memory
x.relu_()  # In-place ReLU

# Bad: Breaks autograd
# x = x + y  # Creates new tensor
x.add_(y)   # In-place addition

Warning: In-place operations can break the computation graph. Only use them when you’re certain they won’t interfere with gradient computation.

Model Compilation with torch.compile

PyTorch 2.0 introduced torch.compile, which can significantly speed up model execution:

model = torch.compile(model, mode='reduce-overhead')

Modes:

  • default: Good balance of performance and compilation time
  • reduce-overhead: For smaller models, reduces framework overhead
  • max-autotune: Maximum performance, longer compile time

Typical speedups range from 1.2-2x depending on model architecture.

Profiling Your Code

Use PyTorch Profiler to identify bottlenecks:

from torch.profiler import profile, ProfilerActivity

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True
) as prof:
    model(input_data)

print(prof.key_averages().table(sort_by="cuda_time_total"))

This reveals which operations consume the most time, guiding optimization efforts.

Distributed Training

For multi-GPU training, use DistributedDataParallel (DDP):

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

# Initialize process group
dist.init_process_group(backend='nccl')

# Wrap model
model = DDP(model, device_ids=[local_rank])

DDP provides near-linear scaling—2 GPUs give ~1.8x speedup, 4 GPUs ~3.5x.

Inference Optimization

For production inference, consider:

TorchScript

# Trace model for deployment
example_input = torch.rand(1, 3, 224, 224)
traced_model = torch.jit.trace(model, example_input)
traced_model.save('model.pt')

# Load and run
loaded_model = torch.jit.load('model.pt')
output = loaded_model(input_data)

ONNX Export

Export to ONNX for cross-platform deployment:

torch.onnx.export(
    model,
    example_input,
    'model.onnx',
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}}
)

TensorRT

For NVIDIA GPUs, convert to TensorRT for maximum inference performance:

import torch_tensorrt

trt_model = torch_tensorrt.compile(
    model,
    inputs=[torch_tensorrt.Input(shape=[1, 3, 224, 224])],
    enabled_precisions={torch.float16}
)

Practical Checklist

Before training:

  • Enable mixed precision if using modern GPUs
  • Optimize DataLoader with appropriate num_workers
  • Use gradient accumulation for large effective batch sizes
  • Profile to identify bottlenecks
  • Consider torch.compile for repeated model execution

During training:

  • Monitor GPU utilization (should be >90%)
  • Check for data loading stalls
  • Verify memory usage isn’t near OOM

These optimizations compound—applying several can yield 5-10x speedups in training time.

Back to top