iwantcoding.com
🔥 Daily 👥 Rooms 🏆 Top Log in Sign up

Profiler

By · Reviewed & last updated · Part of PyTorch Tutorial

torch.profiler captures op-level CPU and CUDA timings, memory usage, and stack traces, then visualises them in TensorBoard or as Chrome trace JSON. Reach for it when training is slow and you cannot tell whether the bottleneck is data loading, kernel launches, copies, or actual compute. Profiling beats guessing every time.

Profile a training step and find the bottleneck

EXAMPLE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from torch.profiler import profile, record_function, ProfilerActivity, schedule, tensorboard_trace_handler

# 1) A small model + fake data
class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 512)
        self.fc2 = nn.Linear(512, 256)
        self.fc3 = nn.Linear(256, 10)
    def forward(self, x):
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return self.fc3(x)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model  = MLP().to(device)
opt    = torch.optim.Adam(model.parameters(), lr=1e-3)
loss   = nn.CrossEntropyLoss()

X = torch.randn(50_000, 784); y = torch.randint(0, 10, (50_000,))
loader = DataLoader(TensorDataset(X, y), batch_size=512, shuffle=True, num_workers=2)

# 2) Profile with a schedule: warm up, then capture, then stop
sched = schedule(wait=1, warmup=1, active=3, repeat=1)

with profile(
    activities=[ProfilerActivity.CPU,
                ProfilerActivity.CUDA] if device == 'cuda' else [ProfilerActivity.CPU],
    schedule=sched,
    on_trace_ready=tensorboard_trace_handler('./logs/torch-profile'),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step, (xb, yb) in enumerate(loader):
        xb, yb = xb.to(device, non_blocking=True), yb.to(device, non_blocking=True)

        # record_function annotates a span you will see in the UI
        with record_function('forward'):
            logits = model(xb)
        with record_function('loss'):
            l = loss(logits, yb)
        with record_function('backward'):
            opt.zero_grad(set_to_none=True)
            l.backward()
        with record_function('step'):
            opt.step()

        prof.step()
        if step >= 6: break

# 3) Inspect summary in the console — fastest first look
print(prof.key_averages().table(sort_by='cpu_time_total', row_limit=10))
print(prof.key_averages().table(sort_by='self_cuda_time_total', row_limit=10))

# 4) Open in TensorBoard
# pip install torch_tb_profiler
# tensorboard --logdir ./logs/torch-profile

# 5) Chrome trace export (drag-and-drop into chrome://tracing or speedscope)
# prof.export_chrome_trace('trace.json')

# 6) Common fixes profiling reveals:
# - DataLoader-bound: increase num_workers, persistent_workers=True, pin_memory=True
# - CPU<->GPU copies dominate: keep tensors on device, avoid .item() in hot loops
# - Small batch / kernel-launch-bound: use larger batches or torch.compile()
# - Memory-bound: switch to bf16/fp16 (autocast), check for activation reservation
# - .backward() dominates: inspect for retain_graph=True or oversized model

# 7) Lightweight alternative for quick checks (no scheduler)
with torch.profiler.profile() as p:
    for _ in range(3):
        out = model(X[:1024].to(device))
print(p.key_averages().table(sort_by='cpu_time_total', row_limit=5))

Why it matters

Profile a *steady-state* iteration, not the first batch — first-batch timings include CUDA init, kernel autotuning, and DataLoader spin-up that you will never see again. The wait+warmup+active schedule exists exactly so the captured slice represents what production training spends its time on.

Tip: Tweak the snippet with Try it Yourself », then sit the quiz at the bottom of the page.

Example

Example
from torch.profiler import profile, ProfilerActivity, record_function
with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    with record_function('forward'):
        model(xb)
print(prof.key_averages().table(sort_by='cuda_time_total'))
Try it Yourself »

Discussion

Loading…