MemoryLane 💾🛣️

by Peter Sharpe

A super-lightweight line-by-line memory profiler for numerical Python code. See where those pesky allocations are coming from!

• Supports PyTorch CUDA memory measurement, and more to come.
• Minimal dependencies (just Rich + your favorite numerical library)

Installation

pip install memorylane[torch]  # For PyTorch support

Usage

To use MemoryLane, just import it and decorate your function with @profile:

import torch
from memorylane import profile

@profile
def my_function():
    x = torch.randn(5120, 5120, device="cuda")
    x = x @ x
    x = x.relu()
    x = x.mean()
    return x

my_function()

This will print your line-by-line memory usage:

Features

• For complicated functions, filter the report to only show lines with non-negligible changes in memory usage: @profile(only_show_significant=True)
• When used from terminal via most editors (e.g., VSCode/Cursor, PyCharm, etc.), the printouts like make_reports.py:11 become clickable links that will take you directly to the offending line in your code
• Profiling of multiple functions, including nested ones (these will be shown with indentation, to allow you to see where the allocations are coming from)
• Report generation in HTML and text formats
• (Work in progress) Support for measuring memory usage of:
  • PyTorch CPU operations
  • NumPy operations
  • JAX operations
  • Python whole-process memory usage
  • ...and more!

Examples

Under construction - for now, see the examples folder!