🎓 A1Learner
🔍

PyTorch Profiling with torch.profiler

The PyTorch profiler is a tool that records which operations a program runs and how long each one takes. Its summary names the operations worth optimising first.

Recording a short run

The profile block records five forward passes. Nothing prints until the block ends.

import torch
from torch.profiler import ProfilerActivity, profile

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 8, 3, padding=1),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(8 * 16 * 16, 4),
)
model.eval()
x = torch.randn(8, 3, 16, 16)

with profile(activities=[ProfilerActivity.CPU]) as prof:
    with torch.no_grad():
        for _ in range(5):
            model(x)
print("Profiling finished:", prof is not None)

The summary names the operations

key_averages lists one row per operation that ran. The convolution appears under its built-in name.

names = [event.key for event in prof.key_averages()]
print("aten::conv2d was profiled:", "aten::conv2d" in names)

Counts are exact

Five passes through one convolution give a count of five. A rerun prints the same count.

for event in prof.key_averages():
    if event.key == "aten::conv2d":
        print("aten::conv2d calls:", event.count)

Example

import torch
from torch.profiler import ProfilerActivity, profile

torch.manual_seed(0)
model = torch.nn.Sequential(
    torch.nn.Conv2d(3, 8, 3, padding=1),
    torch.nn.ReLU(),
    torch.nn.Conv2d(8, 8, 3, padding=1),
    torch.nn.ReLU(),
    torch.nn.Flatten(),
    torch.nn.Linear(8 * 16 * 16, 4),
)
model.eval()
x = torch.randn(8, 3, 16, 16)

with torch.no_grad():
    model(x)

with profile(activities=[ProfilerActivity.CPU]) as prof:
    with torch.no_grad():
        for _ in range(5):
            model(x)

counts = {}
for event in prof.key_averages():
    if event.key in ("aten::conv2d", "aten::addmm", "aten::relu"):
        counts[event.key] = event.count

for name in sorted(counts):
    print(f"{name} calls: {counts[name]}")

Key points

  • Use torch.profiler.profile around the code to record a run
  • ProfilerActivity.CPU selects CPU recording for the profile block
  • key_averages() returns one summary row per operation
  • The count column says how many times an operation ran
  • Call counts are exact, while recorded times vary from run to run
  • The aten prefix marks PyTorch's built-in operations in the output
  • In short, the profiler shows where the time goes before any optimising starts
  • The closing example prints counts only

In this example, a small conv model runs five times inside a profile block. The model has two convolutions, two ReLU layers, and one Linear layer. The summary counts ten convolution calls, which is two per pass. ReLU also counts ten, and the Linear layer's addmm counts five. The counts are exact and repeat on every run. The program prints no times, because recorded times vary and would not repeat. Read the counts first, then chase the operation that runs most or costs most.

Note

The raw profiler table also prints memory figures and addresses that change between runs.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

Why does the closing example print call counts instead of times?