Profiling — cProfile, line_profiler, py-spy

TL;DR

When something is slow, the wrong move is to start rewriting code. The right move is to measure first. Almost always, 90% of the time is spent in 10% of the code. Profile, find the hotspot, optimise the hotspot. Don’t guess.

Three profilers cover the cases:

The workflow: cProfile to find the slow function. line_profiler to find the slow line. py-spy when you can’t (or don’t want to) modify the code, or when the program is running in production. memory_profiler or scalene when you suspect memory.

The picture in your head

A profiler watches your program and records where it spends time. cProfile instruments every function call (deterministic): exact counts, modest overhead. py-spy samples the call stack at a high frequency from outside the process (sampling): lower overhead, slightly fuzzy, no instrumentation needed.

Both produce a flat or hierarchical breakdown of “this function took N% of total time.” That’s all you need: the top of the list is where to look.

cProfile — start here

Built into Python. Profile a script:

python -m cProfile -o profile.out -s cumulative my_script.py

-s cumulative sorts by cumulative time (function + everything it called); -s tottime sorts by exclusive time spent in the function itself.

Programmatic:

import cProfile, pstats

with cProfile.Profile() as pr:
    expensive_function()

stats = pstats.Stats(pr).sort_stats("cumulative")
stats.print_stats(20)

The output looks like:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.001    0.001   12.453   12.453 train.py:42(train_epoch)
      100    8.234    0.082   10.101    0.101 model.py:120(forward)
     6400    1.823    0.000    1.823    0.000 model.py:88(attention)
        1    0.954    0.954    0.954    0.954 dataloader.py:33(load)
        ...

cumtime = total time including subcalls. tottime = time spent in this function exclusive of subcalls. Sort by cumtime to see “where did the program spend its time?”; sort by tottime to see “which function is the hotspot?”

For a visual tree, view the profile with snakeviz:

pip install snakeviz
snakeviz profile.out

Browser opens a sunburst chart. Way easier to reason about than the text dump.

line_profiler — drill into a function

pip install line_profiler

Decorate the functions you want profiled:

@profile     # injected by kernprof; no import needed
def train_step(model, batch):
    x, y = batch
    pred = model(x)
    loss = F.cross_entropy(pred, y)
    loss.backward()
    optimizer.step()
    return loss.item()

Run with kernprof:

kernprof -l -v my_script.py

Output (per line):

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
   42                                           @profile
   43                                           def train_step(model, batch):
   44       100      1240.0     12.4      0.4       x, y = batch
   45       100    105800.0   1058.0     33.4       pred = model(x)
   46       100      8200.0     82.0      2.6       loss = F.cross_entropy(pred, y)
   47       100    198400.0   1984.0     62.7       loss.backward()
   48       100       890.0      8.9      0.3       optimizer.step()
   49       100       170.0      1.7      0.1       return loss.item()

Now you know loss.backward() is 62.7% of the train step — and that optimising the forward pass would help less than improving the backward.

Use line_profiler after cProfile has pointed you at the function. Don’t profile your whole codebase line-by-line; the overhead is large.

py-spy — sample a running process

pip install py-spy

py-spy reads memory of a running Python process from outside. No need to modify code, no need to restart the program.

# Show the top consumers in a running process, like `top`
sudo py-spy top --pid 12345

# Record N seconds and produce a flame graph
sudo py-spy record -o profile.svg --pid 12345 --duration 30

py-spy record produces a flame graph SVG you can open in a browser. The width of each block = total time spent there. The classic flame graph reading: look for “plateaus” — wide flat regions are where the time goes.

This is the right tool for:

• Production. A live training job. You can attach without restarting.
• Long-running services where adding instrumentation is costly.
• Debugging “it’s just slow” without a clear hypothesis.

memory_profiler — line-by-line memory

pip install memory_profiler

@profile
def big_op():
    x = np.zeros((10_000, 10_000))   # 800 MB
    y = x * 2                         # another 800 MB
    return y.sum()

python -m memory_profiler my_script.py

Output shows current and incremental memory per line. Useful when RAM keeps climbing and you don’t know which step is at fault.

For “which lines retain memory between calls” — that’s leak detection, which memory_profiler won’t catch. Use tracemalloc (stdlib) or pympler for that.

scalene — the modern all-in-one

pip install scalene

scalene my_script.py

Profiles CPU, memory, and GPU together with low overhead. Output is a detailed line-by-line breakdown in a browser UI. Newer than the others; worth knowing.

A worked example — debugging a slow training loop

Suppose train.py is unexpectedly slow. The workflow:

# 1) Get a coarse view
python -m cProfile -o profile.out -s cumulative train.py
snakeviz profile.out

You see dataloader.next() taking 60% of the time. The model itself is only 30%.

# 2) Drill into the dataloader
@profile  # add to dataloader.next()
def next(self):
    raw = self._read_next_chunk()
    decoded = self._decode(raw)
    return self._collate(decoded)

kernprof -l -v train.py

You see _decode is 90% of next(). Specifically a Python loop inside _decode is the hotspot.

You vectorize the loop, rerun:

python -m cProfile -o profile2.out -s cumulative train.py

Now the dataloader is 5% of total time, and the model is 80% — exactly where you want it.

The lesson: always profile, never guess. The wrong instinct in that scenario would be “the model must be slow, let me try mixed precision.” That would have helped some, but the actual win was a 50-line vectorization in a place you wouldn’t have looked otherwise.

ML-specific profilers

For ML workloads specifically:

• torch.profiler — see exactly which CUDA kernels and CPU ops are running. Outputs Chrome-trace JSON viewable in chrome://tracing or via torch.profiler.tensorboard_trace_handler for TensorBoard.
• NVIDIA Nsight Systems / Nsight Compute — kernel-level GPU profiling.
• nvtx annotations — mark named ranges in your Python that show up in Nsight traces.

from torch.profiler import profile, ProfilerActivity, record_function

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    with record_function("forward"):
        out = model(x)
    with record_function("backward"):
        out.backward()

print(prof.key_averages().table(sort_by="cuda_time_total"))

For training-loop optimisation, torch.profiler tells you what the GPU is actually doing — including how much time is wasted waiting on the Python side feeding it.

Common gotchas

• Profiling without a baseline. “It feels slow” is not enough. Time something, then compare after the change.
• Optimising the wrong function. cProfile first, always.
• Profiling the wrong workload. Profile a real-sized run, not a toy one. Hotspots change with input size.
• cProfile overhead skews relative timing for very fast functions. Take results with a grain of salt for sub-microsecond functions.
• Production overhead from cProfile / line_profiler. Don’t leave them on. py-spy is the production-safe option.
• GPU work doesn’t show up in cProfile. PyTorch CUDA operations return immediately (kernels run async). Use torch.profiler or cudaDeviceSynchronize calls before timing.

When NOT to profile

• The code is fast enough. Don’t optimise what doesn’t need it.
• You’re early in development. Get it working first; optimise once.
• You haven’t measured. “I think this is slow” is not enough; measure.

The 80/20 rule applies extremely strongly to performance: most code doesn’t matter for speed. Find the 20% that does, then focus there.

Where this shows up in real ML codebases

• torch.profiler in training scripts to debug data-loading vs model-forward vs optimizer-step.
• py-spy attached to a running training job whose tokens/sec dropped.
• cProfile + snakeviz to debug slow eval pipelines.
• scalene for “is this slow because of CPU, GPU, or RAM?” questions.
• Custom timing decorators for production services (Datadog/Prometheus instrumentation), built on the patterns from PY 105.

The defensive habit: measure before, change one thing, measure after, keep the change only if it helped. Engineering, not vibes.

Resources

• Python docs — cProfile — docs.python.org — stdlib reference.
• line_profiler — github.com/pyutils/line_profiler — line-level timing.
• py-spy — github.com/benfred/py-spy — sampling profiler for live processes.
• memory_profiler — pypi.org/project/memory-profiler — line-level memory.
• scalene — github.com/plasma-umass/scalene — modern combined CPU/GPU/memory profiler.
• snakeviz — github.com/jiffyclub/snakeviz — browser-based cProfile viewer.
• High Performance Python (2nd ed.) — Profiling chapter — oreilly.com — the canonical treatment.
• PyTorch — torch.profiler — pytorch.org — for ML-specific profiling.
• Brendan Gregg — flame graphs — brendangregg.com/flamegraphs.html — how to read py-spy’s output.