TopolearnOptimising without measuring wastes effort; most slow training loops are bottlenecked somewhere unexpected.
Before optimising, find the bottleneck. Common culprits are the data pipeline, CPU-GPU synchronisations and many small kernels, rather than the big matmuls.
`torch.profiler.profile` records operator-level CPU and CUDA activity:
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]chooses what to trace.record_shapes=Trueandprofile_memory=Trueadd shape and allocation information.- Wrap regions with
record_function("data")to label them in the trace. prof.key_averages().table(sort_by=...)prints the top operators;prof.export_chrome_trace("trace.json")writes a timeline you can open in Perfetto.
For training loops, pass schedule=torch.profiler.schedule(wait=1, warmup=1, active=3) and call prof.step() each iteration, so you skip start-up noise and keep traces small.
What to look for in a timeline: gaps where the GPU is idle (data loading or Python overhead), many tiny kernels (candidates for torch.compile), and cudaMemcpy or synchronize calls in the middle of steps.
For micro-benchmarks, use `torch.utils.benchmark.Timer`. It handles warm-up and CUDA synchronisation and reports statistics, which naive time.time() measurements get wrong.
Example
import torch
from torch import nn
from torch.profiler import profile, record_function, ProfilerActivity
import torch.utils.benchmark as benchmark
model = nn.Sequential(nn.Linear(1024, 1024), nn.ReLU(), nn.Linear(1024, 10))
x = torch.randn(256, 1024)
activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
activities.append(ProfilerActivity.CUDA)
with profile(activities=activities, record_shapes=True) as prof:
with record_function("forward"):
model(x)
print(prof.key_averages().table(sort_by="cpu_time_total", row_limit=5))
prof.export_chrome_trace("trace.json") # open in Perfetto
t = benchmark.Timer(stmt="model(x)", globals={"model": model, "x": x})
print(t.timeit(100)) # mean time per run, warm-up handledRemember
- Profile first: data loading and syncs are frequent bottlenecks
- Label regions with
record_function; view traces in Perfetto - Use
torch.utils.benchmarkfor micro-benchmarks
Watch out for
- Profiling the first iterations, which include warm-up, cuDNN autotuning and compilation.