torch.compile and a faster input pipeline

Where training time actually goes

torch.compile(model) is PyTorch 2's one-line speed-up. It captures your model's Python code into graphs (via TorchDynamo) and generates optimised, fused kernels (via TorchInductor by default), so several small operations run as one, reducing memory traffic and Python overhead. The first calls are slow because compilation happens then, and again if input shapes change; later calls reuse the result. Code it cannot capture causes a graph break and simply runs in eager mode. Gains vary by model and hardware, so measure.

Often the GPU is not the bottleneck at all. If utilisation shown by nvidia-smi is low, the GPU is waiting for data. Raise num_workers, set pin_memory=True and use .to(device, non_blocking=True) so copies overlap with compute, and use a batch size large enough to keep the GPU busy.

Avoid Python loops over tensor elements; express the computation as whole-tensor operations. And profile before guessing: torch.profiler shows which operations and data steps consume the time.

Compile the model, feed it well

SmallCNN, device and train_ds come from earlier topics. non_blocking=True only helps when the source batch is in pinned memory.

import torch
from torch.utils.data import DataLoader

model = SmallCNN().to(device)
compiled = torch.compile(model)        # first call compiles, later calls reuse

loader = DataLoader(train_ds, batch_size=128, shuffle=True,
                    num_workers=4, pin_memory=True)
for x, y in loader:
    x = x.to(device, non_blocking=True)   # copy overlaps with compute
    y = y.to(device, non_blocking=True)
    logits = compiled(x)
    ...

Vectorise instead of looping

Both functions compute row-wise squared distances from a point. The loop launches many tiny operations from Python; the vectorised version does one batched operation and is far faster on large inputs, especially on a GPU.

import torch

points = torch.randn(10_000, 3)
center = torch.zeros(3)

def slow(points, center):
    out = torch.empty(len(points))
    for i in range(len(points)):           # one Python iteration per row
        out[i] = ((points[i] - center) ** 2).sum()
    return out

def fast(points, center):
    return ((points - center) ** 2).sum(dim=1)   # broadcasting, one call

print(torch.allclose(slow(points, center), fast(points, center)))   # True

Measure first

Change one thing at a time and time a fixed number of steps after a warm-up (compilation and cuDNN autotuning make the first steps unrepresentative). Remember to synchronise the GPU before stopping the clock, or you will time only how fast Python queued the work.

Quick check: Training is slow and `nvidia-smi` shows the GPU mostly idle. What is the most likely first fix?

• Switch from Adam to SGD
• Add more layers so the GPU has more work
• Call `torch.cuda.synchronize()` after every batch
• Speed up data loading, for example with more `num_workers` and `pin_memory=True`