pytorch_memlab: per-line CUDA memory accounting for PyTorch debugging

pip installs two names for a single distribution

Two install lines sit at the top of the README, and they spell the package name two different ways. The shorter one fetches the published build:

pip install pytorch_memlab

The longer one skips the package index and pulls the default branch, which the repository names master, straight from GitHub:

pip install git+https://github.com/stonesjtu/pytorch_memlab

The disagreement sits in the middle of the name. setup.py registers the distribution as `name='pytorch-memlab'`, with a hyphen, and repeats that spelling in the extras line `pip install pytorch-memlab[ipython]`, while the plain install drops that hyphen from the same word. PyPI answers to both, because pip folds those two characters together when matching a distribution name. Installing from git sidesteps the question but drops the version pin entirely.

Beyond the name, setup.py pins `version='0.3.2'` and hands the same licence to setuptools twice, once as `licence='MIT'` and once as `license='MIT'`. The floor it declares is explicit: `torch>=2.0`, `calmsize`, `pandas` and `setuptools`, with `python_requires='>=3.8'`, an `IPython>=0.13` extra under the `ipython` key and `pytest` under `test`. The version list jumps from 0.3.0 dated 2023-07-29 to 0.3.1 at 08:49:54 and 0.3.2 at 08:57:13, both on 2026-06-08, and the last push to the repository came on 2026-08-12, after both were published.

Active bytes and reserved bytes answer different questions

Out-Of-Memory kills in PyTorch hit new programmers and experienced ones alike, and this project lays the blame on code that never learned how CUDA memory behaves rather than on a defect in the framework. Its answer is per-line accounting: the profiler is a modification of Python's `line_profiler` that records memory for every executed line of the chosen function instead of wall-clock time.

Two numbers come back per line. `active_bytes` all peak tracks bytes held by live allocations, while `reserved_bytes` all peak tracks what the caching allocator has taken from the device, which is how a card can sit at 2.00M reserved with a fraction of that actually in use.

The decorator to reach for first is `profile`, and it departs from `line_profiler` in one way that matters. Memory statistics are gathered over repeated runs of the function and only the maximum is printed when it finishes. When the whole sequence is what you need instead of one peak, `profile_every` prints on every N-th execution, and swapping `@profile` for `@profile_every(1)` is the entire change. The two decorators can be mixed across different functions for finer granularity, including on `torch.nn.Module` methods such as `forward`, which matters when the growth sits in one layer rather than in a free function.

Any field name from `memory_stats()` can be handed to `display()` to print a statistic other than those two defaults, so the profiler is not welded to the allocator's headline counters. It remains a decorator-based tracer all the same, and that design is what turns the next detail into a global concern.

set_target_gpu is process-wide state you must restore

CUDA device 0 is the profiling target unless something says otherwise, and `set_target_gpu` is how you say it. The sample below allocates on device 0, switches to device 1 for the second layer, then switches back, which keeps every line's numbers attached to the device that produced them:

import torch
from pytorch_memlab import profile, set_target_gpu
@profile
def func():
    net1 = torch.nn.Linear(1024, 1024).cuda(0)
    set_target_gpu(1)
    net2 = torch.nn.Linear(1024, 1024).cuda(1)
    set_target_gpu(0)
    net3 = torch.nn.Linear(1024, 1024).cuda(0)

func()

The trade-off is scope. GPU selection here is global rather than per decorator, per thread or per context, so a function that never calls `set_target_gpu` still reports against whatever device was selected last anywhere in the process. A library that changes the target on import, or a second profile running inside the same interpreter, can move your numbers without touching your function at all. The project states the hazard in a single sentence: you have to remember which gpu you are profiling for the whole process.

There is no documented reset to a default device, and no device argument appears on the decorators in the API shown here. In a notebook session or a test suite running many profiles, that omission is the first thing to check when a report looks impossible.

The printed table and the sample code disagree

The first sample defines `inner()` and `outer()` and then calls `work()`, a name that appears nowhere else in the file. Run exactly as printed, it raises a NameError before a single byte is measured. The terminal table underneath was captured from a different version of the same function:

## outer

active_bytes reserved_bytes line  code
         all            all
        peak           peak
       0.00B          0.00B    7  def outer():
      40.00K          2.00M    8      linear = torch.nn.Linear(100, 100).cuda()
      80.00K          2.00M    9      linear2 = torch.nn.Linear(100, 100).cuda()
     120.00K          2.00M   10      inner()

In the listing printed above, line 10 is the third `torch.nn.Linear(100, 100).cuda()` call, and `outer()` as written never invokes `inner()` at all. Reading the numbers still pays off. Active bytes climb in exact 40.00K steps, 0.00B then 40.00K then 80.00K then 120.00K, one step per layer, while `reserved_bytes` jumps to 2.00M on the first allocation and never moves again. That fixed overhead is what decides whether a shared GPU runs dry, and a profiler printing only one of the two columns would hide it completely.

MemReporter counts storage bytes behind every tensor

Line totals cannot say which tensor is holding the memory, so a second tool walks the object graph. `MemReporter` iterates every `Tensor` and reads the underlying `UntypedStorage`, the type previously called `Storage`, because the size on the tensor surface says nothing about bytes when several views share one buffer.

import torch
from pytorch_memlab import MemReporter
linear = torch.nn.Linear(1024, 1024).cuda()
reporter = MemReporter()
reporter.report()

Five lines is as far as the docs go, and the sample output paired with them stops after a single column heading, `Element type`. Nobody reading the README sees which fields follow it, how rows are ordered, whether the total is grouped by storage or by tensor, or whether `report()` takes arguments at all. The walk is process-wide as well, inspecting whatever tensors are alive in the interpreter when it runs, so the total depends on when you call it inside a loop, and no documented flag scopes it to one module or one device.

For an adoption decision that gap outranks the feature list. The mechanism itself is a short, readable loop over live objects, but you will be reading the source under `pytorch_memlab/` to learn the output format.

IPython magics wrap the same decorators

Inside a notebook the profiler is reachable through two magics instead of decorators, which takes the edit-and-rerun cycle out of a debugging session. IPython has to be installed already or pulled in through the `ipython` extra, and the extension loads first with `%load_ext pytorch_memlab`. Cell magic then names the function with `-f` and profiles everything the cell defines and runs:

%%mlrun -f func
import torch
from pytorch_memlab import profile, set_target_gpu
def func():
    net1 = torch.nn.Linear(1024, 1024).cuda(0)
    set_target_gpu(1)
    net2 = torch.nn.Linear(1024, 1024).cuda(1)
    set_target_gpu(0)
    net3 = torch.nn.Linear(1024, 1024).cuda(0)

The line form profiles a single call instead, as in `%mlrun -f func func(2048)`, and `%mlrun?` opens the supported arguments: which device to profile, where to dump results to a file, and how to get the `LineProfiler` object back for inspection after the run.

None of that is a second profiler. The magics drive the very same decorators, which is why the global device rule from `set_target_gpu` applies inside a notebook exactly as it does in a script, and why a session that leaves the target on device 1 hands you numbers you never asked for. A worked session lives in `demo.ipynb`, and more decorator samples sit in `test/test_line_profiler.py`.

Courtesy has a feature entry and no documented call

The feature list names four parts. Three of them come with worked examples: the line profiler, the reporter and the IPython support. The fourth, Courtesy, is described in the bullet list as a way to temporarily move all CUDA tensors into CPU memory and back again, and the table of contents links a section for it.

No import line, function name or code sample for Courtesy appears anywhere else in the README, so the call you would need is never written down. That gap deserves weight, because moving live tensors to host memory and back is the operation most able to change what a program computes rather than how much memory it holds. Nothing in the docs says whether references are cached or rebuilt on the return trip, what happens to a tensor that was already on the CPU, or whether the round trip is safe inside autograd.

Treat that surface as unverified until you read the implementation. `pytorch_memlab/` holds the package and `test/` holds the suite that pins its behaviour down.

The same counters come from torch.cuda.memory_stats

The serious alternative is to skip the decorators and read PyTorch's allocator directly, because the numbers this tool prints already originate there. The README points at that source itself, sending readers to `memory_stats()` for what each column means and letting any field from that dict be passed to `display()`. The difference is granularity rather than data. Sampling `torch.cuda.memory_stats` before and after a block gives one number per block, with no line attribution, no `profile_every` sequence and no grouping; what pytorch_memlab adds is an automatic mapping from a line of Python to a row of the table.

That gap marks a clear adoption boundary. Take the package when an Out-Of-Memory has to be pinned to a single line and torch 2.0 is already installed with spare device memory for instrumentation. Leave it alone for CPU-only workloads, where every part of the tool is built around CUDA device memory, and for a tensor-level audit you cannot verify from the docs. Check two things before depending on either: what `report()` prints after `Element type`, and which function actually implements Courtesy.

Editorial conclusion

Adopt pytorch_memlab when a CUDA Out-Of-Memory has to be traced back to one line of Python and you already run torch 2.0 with device memory to spare; the split between active and reserved bytes per line is what no built-in counter hands you. Do not adopt it for CPU-only work, for a tensor-level audit the docs never finish describing, or for the Courtesy round trip whose entry point is never written down. Before trusting any of it, run the 0.3.2 sample and check that its table matches your own function.

Frequently asked questions

Official sources

1. Issues
2. License: MIT
3. README
4. Releases
5. Stonesjtu/pytorch_memlab on GitHub