Stonesjtu/pytorch_memlab

pytorch_memlab – A lightweight CUDA‑memory debugging toolbox for PyTorch

What it is – A pure‑Python package that lets you see how much GPU memory each line of your PyTorch code consumes, list the tensors that occupy the device, and temporarily move everything to CPU. It works both from the command line and inside Jupyter/IPython via magic commands.

Key features

• Line‑by‑line memory profiler (LineProfiler / @profile decorator) – similar to line_profiler but reports CUDA active_bytes and reserved_bytes per source line.
• Memory reporter (MemReporter) – enumerates all live torch.Tensor objects, shows the underlying UntypedStorage size, and can handle shared weights, gradients, and verbose storage‑reuse diagrams.
• Courtesy mode – a helper class that can move all GPU tensors back to CPU on demand, useful for freeing the GPU without checkpointing.
• IPython integration – %mlrun and %%mlrun magics let you profile a function or an entire notebook cell with a single command.
• GPU selection – set_target_gpu lets you switch which device the profiler watches during a run.

How to install

# stable release from PyPI
pip install pytorch_memlab

# or the latest code directly from GitHub
pip install git+https://github.com/stonesjtu/pytorch_memlab

(For the IPython magics add the optional extra: pip install pytorch_memlab[ipython].)

Typical workflow

1. Profile a function

   from pytorch_memlab import profile
   
   @profile
   def train_step(x):
       net = torch.nn.Linear(1024, 1024).cuda()
       return net(x).mean()
   

   After the script finishes the profiler prints a table showing how much GPU memory each line allocated.
2. Inspect live tensors

   from pytorch_memlab import MemReporter
   reporter = MemReporter(model)   # model is optional
   reporter.report()               # prints a concise table
   reporter.report(verbose=True)  # shows storage sharing arrows
   

3. Use in a notebook

   %load_ext pytorch_memlab
   %%mlrun -f train_step
   train_step(torch.randn(512, 1024).cuda())
   

4. Free the GPU temporarily

   from pytorch_memlab import Courtesy
   c = Courtesy()
   c.yield_memory()   # moves tensors to CPU
   # …wait for a signal…
   c.restore()        # moves them back
   

Why it matters – Out‑of‑Memory (OOM) crashes are a common pain point when developing deep‑learning models. pytorch_memlab makes the hidden CUDA allocation behaviour visible, helping you spot inefficient tensor lifetimes, unintended storage sharing, or forgotten intermediate buffers.

Limitations / known issues

• The profiler only sees tensors that Python knows about; low‑level C‑side buffers created by autograd are reported as “invisible” and can still consume memory.
• Even when all tensors are on CPU, the PyTorch CUDA context itself reserves ~1 GB of GPU memory, which cannot currently be released by the library.
• Works only on CUDA‑enabled GPUs; no support for AMD or CPU‑only runs.

Project status – Actively maintained (latest release 0.3.2, June 2026). CI runs tests and publishes wheels to PyPI. The README includes usage examples, a demo notebook, and a changelog.

---

All information above is taken directly from the repository’s README; no external claims have been added.