5.3. Inside PyTorch: PyTorch Profiler🔗

PyTorch Profiler instruments the framework and supported device runtime. A typical training capture selects activities, excludes startup with a schedule, labels one application range, advances the schedule after each step, prints operator aggregates, and exports a trace:

import torch
from torch.profiler import ProfilerActivity, profile, record_function

activities = [ProfilerActivity.CPU]
if torch.cuda.is_available():
    activities.append(ProfilerActivity.CUDA)

def save_trace(prof):
    print(prof.key_averages().table(
        sort_by="self_cpu_time_total",
        row_limit=20,
    ))
    prof.export_chrome_trace(f"trace-{prof.step_num}.json")

with profile(
    activities=activities,
    schedule=torch.profiler.schedule(
        wait=1, warmup=1, active=3, repeat=1
    ),
    on_trace_ready=save_trace,
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for batch in loader:
        with record_function("training.step"):
            loss = model(batch)
            loss.backward()
            optimizer.step()
            optimizer.zero_grad()
        prof.step()

The user range is only one event in this capture. PyTorch can also record framework operators and supported runtime activity below it. With CUDA activity enabled, its Kineto integration uses CUPTI where available to add runtime calls and on-device kernels to the trace. Other supported builds may offer XPU or additional accelerator activity; torch.profiler.supported_activities() reports what the current installation exposes.

record_shapes, profile_memory, with_stack, and with_flops answer different questions and all have costs. Shape recording can retain tensor references and affect optimizations. Stack collection adds source information but increases overhead. FLOP estimates cover selected operators rather than arbitrary code. Start with the smallest feature set that answers the question.