PyTorch Profiler: read operator costs and export a CPU trace

PyTorch Profiler records operations inside a chosen part of your program. Use it to find expensive operators, connect them to named stages of your code, and inspect the order of execution. A useful capture starts with a specific question: which operations account for time in this forward pass?

This guide runs a small CPU inference workload, captures two batches, and exports a trace alongside timing and memory tables. The example was tested on Windows with Python 3.12.14 and PyTorch 2.15.0+cpu. It uses synthetic inputs and needs no GPU, model download, or dataset. The measurements illustrate how to read a capture; they are not a model-performance benchmark.

Choose the work you want to measure

The lab uses a two-layer network with a ReLU between the layers. Each batch has 128 rows and 256 input features; the output has 64 features. Model construction and input generation happen before profiling. The capture therefore answers a question about repeated forward passes, not startup time, data loading, or an entire application request.

The script selects evaluation behavior with model.eval() and runs inside torch.inference_mode(). Keep that choice aligned with the task you intend to measure. A training capture needs the actual training path, including backward and optimizer work. Our guide to eval, no_grad, and inference_mode explains the differences.

One CPU thread makes this lab's configuration explicit. It is not a recommendation to reduce your application's thread count. Record your own thread settings, batch shape, model mode, device, and PyTorch version when collecting evidence.

Run a scheduled CPU capture

Save the following code as profile_cpu.py in an environment with PyTorch installed. It writes into a new directory and refuses to overwrite an existing run. The callback saves both tables and the trace from the completed capture.

import json
from pathlib import Path

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


def run_profile(output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=False)
    torch.manual_seed(7)
    torch.set_num_threads(1)  # A fixed lab setting, not a tuning recommendation.
    model = nn.Sequential(nn.Linear(256, 512), nn.ReLU(), nn.Linear(512, 64))
    model.eval()
    inputs = torch.randn(128, 256)
    captures = []

    def save_trace(prof):
        stem = output_dir / f"cpu-step-{prof.step_num}"
        averages = prof.key_averages()
        time_table = averages.table(sort_by="self_cpu_time_total", row_limit=12)
        memory_table = averages.table(sort_by="self_cpu_memory_usage", row_limit=12)
        stem.with_suffix(".time.txt").write_text(time_table, encoding="utf-8")
        stem.with_suffix(".memory.txt").write_text(memory_table, encoding="utf-8")
        prof.export_chrome_trace(str(stem.with_suffix(".json")))
        captures.append({
            "step": prof.step_num,
            "operators": [{"name": e.key, "calls": e.count,
                           "self_cpu_us": e.self_cpu_time_total,
                           "total_cpu_us": e.cpu_time_total,
                           "self_cpu_bytes": e.self_cpu_memory_usage}
                          for e in averages],
        })
        print(time_table)
        print(memory_table)

    with torch.inference_mode():
        reference = model(inputs)
        for _ in range(3):
            model(inputs)  # Warm up the workload before starting the profiler.
        with profile(
            activities=[ProfilerActivity.CPU],
            schedule=schedule(wait=1, warmup=1, active=2, repeat=1),
            on_trace_ready=save_trace,
            profile_memory=True,
            record_shapes=False,
            with_stack=False,
        ) as prof:
            for step in range(6):
                with record_function(f"batch_{step}"):
                    with record_function("forward"):
                        output = model(inputs)
                    del output
                prof.step()  # Advance once after each logical batch.
        torch.testing.assert_close(model(inputs), reference, rtol=0, atol=0)

    report = {"torch": torch.__version__, "device": "cpu",
              "threads": torch.get_num_threads(), "input_shape": list(inputs.shape),
              "output_shape": list(reference.shape), "captures": captures}
    (output_dir / "summary.json").write_text(json.dumps(report, indent=2), encoding="utf-8")
    return report


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("output_dir", help="New directory for this capture")
    args = parser.parse_args()
    run_profile(args.output_dir)

Run it with a directory name you have not used before:

python profile_cpu.py profiler-run-01

The directory contains cpu-step-4.json, cpu-step-4.time.txt, cpu-step-4.memory.txt, and summary.json. The script also compares a prediction before and after profiling, checking that the result is unchanged. It does not save model weights or modify a dataset.

Check which iterations were recorded

There are two separate warm-ups here. Three unrecorded forward passes prepare the workload before the profiler starts. Inside the profiler, the schedule waits for one iteration, warms up for one, then records two. The loop still executes all six iterations.

• Batch 0: wait, with no captured batch range.
• Batch 1: profiler warm-up, excluded from the exported capture.
• Batches 2 and 3: active capture.
• Batches 4 and 5: no further capture because repeat=1.

Our trace contained exactly batch_2 and batch_3, with two forward ranges. The callback ran once, after the fourth call to prof.step(), which explains the filename. Keep that call at the end of each logical step. If a step means an optimizer update rather than a batch in your application, label and advance it accordingly. The profiler API reference documents schedules and callbacks.

Read self time before adding up totals

The table aggregates repeated operator names. Self CPU excludes time in nested calls; CPU total includes them. CPU time avg is the total divided by the call count. The official profiler recipe explains these columns and memory reporting.

These selected rows come from our captured run, in microseconds. They are reformatted from summary.json; timings will differ on another run.

Name                Calls   Self CPU (us)   CPU total (us)
forward                 2           358.1           3085.1
aten::linear            4           265.3           2455.3
aten::addmm             4          1522.5           1579.8
aten::clamp_min         2            40.6             40.6

The two forward passes generated four linear calls because the model has two linear layers. Here, aten::addmm accounted for the largest self time. That makes the matrix operations a reasonable place to investigate in this particular capture. It does not establish that every model is limited by matrix multiplication.

The forward, aten::linear, and aten::addmm totals overlap because those operations are nested. Adding those totals would count some work several times. Likewise, a high total with a small self time can identify a wrapper around expensive children. Use the trace to see the relationship before choosing a code change.

The named forward range also includes Python work between its child operations. Its self time is not a separate mathematical layer. For a small example like this, profiler and labeling overhead can be a substantial part of the observed time.

Interpret memory as allocation activity

With profile_memory=True, the second table sorts by self_cpu_memory_usage. Positive and negative entries describe tensor allocation and release attributed to the recorded work. Self memory excludes child calls. These columns do not report the process's peak resident memory.

In this run, aten::addmm reported 589,824 bytes of self allocation: two captured passes × 128 rows × (512 + 64) outputs × four bytes per float32 value. The ReLU allocation appeared separately. Intermediate tensors and the final outputs were released during each batch, so the batch ranges had zero net memory change despite allocating memory while they ran.

Zero net change does not mean zero memory was needed. Model parameters, inputs, and the reference prediction already existed before the active window. This capture cannot reconstruct their earlier allocations or give a complete memory budget for the process.

Open the trace and follow one batch

Open cpu-step-4.json in the Perfetto trace viewer using its local-file option. The Perfetto UI guide covers loading and navigating traces. Search for batch_2, inspect its forward range, and follow the nested linear and ReLU operations. Compare it with batch_3.

The table helps rank aggregated work; the timeline helps explain where that work occurred. Our checks parsed the exported JSON and verified its batch names, operator counts, durations, and allocation/release events. They did not include a visual inspection in the trace viewer.

Make one change and measure it separately

Use the capture to form a testable explanation, then compare the original and changed code without profiling enabled. Keep the input, model mode, thread settings, and device consistent, warm up both versions, and collect repeated measurements. torch.utils.benchmark.Timer provides warm-up and repeated timing support. Check output correctness alongside any timing comparison.

For a second diagnostic pass, enable record_shapes=True when input sizes are the question, then group averages with group_by_input_shape=True. Shape and stack recording add overhead; shape recording also retains tensor references temporarily. Keep them off when they do not help answer the question.

If the delay occurs before the forward pass, expand the captured region to include batch retrieval and preparation. Our DataLoader guide covers those parts of the input path. This script uses an already-created tensor, so it measures no file reading, preprocessing workers, or host-to-device transfer.

What this CPU example verifies

Ten local checks covered the capture window, callback count, nested timing, operator counts, tensor allocation and release, exported files, unchanged output, and protection against overwriting a capture. The JSON contained no accelerator kernel events. The test environment lacked NumPy and emitted an initialization warning, but this tensor-only script completed without using NumPy.

GPU kernel profiling requires an appropriate device, build, and profiling support. Adding a CUDA activity to a CPU-only example does not supply that evidence. This lab validates CPU inference capture and export; it does not establish GPU bottlenecks, training performance, distributed behavior, cloud VM throughput, or a speedup.