# torch.compile and a faster input pipeline — PyTorch

Source: https://www.skillbyai.com/en/pytorch/pt-compile-speed

> Speed up a model with torch.compile, keep the GPU fed with an efficient DataLoader, and profile before optimising.

## 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.

```python
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.

```python
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.

**Quiz:** 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
- [x] Speed up data loading, for example with more `num_workers` and `pin_memory=True`

*Answer:* Speed up data loading, for example with more `num_workers` and `pin_memory=True`. An idle GPU is usually waiting for the input pipeline. Parallel loading and pinned memory keep batches ready; extra synchronisation would slow things further.