Advanced

PyTorch Best Practices

Optimize training performance with mixed precision and compilation, scale to multiple GPUs with distributed training, debug effectively, and deploy models to production.

Mixed Precision Training

Python

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()

for inputs, targets in train_loader:
    optimizer.zero_grad()

    # Forward pass in mixed precision
    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    # Scaled backward pass
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

torch.compile (PyTorch 2.0+)

Python

# One-line speedup with torch.compile
model = torch.compile(model)  # Up to 2x faster!

# Different modes for different tradeoffs
model = torch.compile(model, mode="reduce-overhead")  # Best for small models
model = torch.compile(model, mode="max-autotune")     # Best absolute performance

DataLoader Optimization

Python

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

Distributed Training

Python

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

# Initialize process group
dist.init_process_group(backend='nccl')
local_rank = int(os.environ['LOCAL_RANK'])

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

# Launch with: torchrun --nproc_per_node=4 train.py

Debugging Tips

Check tensor shapes
Print shapes at each step: print(x.shape). Shape mismatches are the most common bug.
Detect NaN/Inf
Use torch.autograd.set_detect_anomaly(True) to get a traceback when NaN appears in gradients.
Verify data pipeline
Visualize a batch before training: plot images, check label distributions, verify normalization.
Gradient checking
Use torch.autograd.gradcheck for custom operations to verify gradient correctness.

Production Deployment

Python

# Export with TorchScript for C++ inference
scripted = torch.jit.script(model)
scripted.save('model_scripted.pt')

# Export to ONNX for cross-platform deployment
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(model, dummy_input, 'model.onnx',
                  input_names=['image'], output_names=['prediction'])

# Serve with TorchServe
# torch-model-archiver --model-name resnet --version 1.0 \
#   --serialized-file model_scripted.pt --handler image_classifier
# torchserve --start --model-store model_store --models resnet=resnet.mar

Quick Reference Checklist

Practice	Impact
Use torch.compile()	Up to 2x speedup with one line
Enable mixed precision	1.5-3x faster training on modern GPUs
Set num_workers > 0	Prevent data loading bottleneck
Use pin_memory=True	Faster CPU-to-GPU data transfer
Call model.eval()	Correct BatchNorm/Dropout in inference
Use torch.no_grad()	Save memory during inference
Set random seeds	Reproducible results

Course Complete!

You now have a solid foundation in PyTorch. Continue your deep learning journey by exploring JAX for high-performance computing or FastAI for rapid prototyping.

Next Course: JAX →