memorylane 0.1.0

Super-lightweight line-by-line memory profiler for numerical Python code. See where those allocations are coming from! Supports NumPy, PyTorch, and more to come.

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:

<style> .r1 {color: #800080; text-decoration-color: #800080; font-weight: bold} .r2 {font-weight: bold} .r3 {color: #008080; text-decoration-color: #008080; font-weight: bold} .r4 {color: #5f8787; text-decoration-color: #5f8787} .r5 {color: #5f8787; text-decoration-color: #5f8787; font-weight: bold} .r6 {color: #008000; text-decoration-color: #008000; font-weight: bold} .r7 {color: #f8f8f2; text-decoration-color: #f8f8f2; background-color: #272822} .r8 {color: #ff4689; text-decoration-color: #ff4689; background-color: #272822} .r9 {color: #ae81ff; text-decoration-color: #ae81ff; background-color: #272822} .r10 {color: #e6db74; text-decoration-color: #e6db74; background-color: #272822} .r11 {color: #7f7f7f; text-decoration-color: #7f7f7f} .r12 {color: #800000; text-decoration-color: #800000; font-weight: bold} .r13 {color: #66d9ef; text-decoration-color: #66d9ef; background-color: #272822} body { color: #000000; background-color: #ffffff; } /* Disable wrapping in

 and  blocks */
pre, code {
    white-space: pre;
    overflow-x: auto;
}
</style>


    
 ━━━━━━ MemoryLane: Line-by-Line Memory Profiler ━━━━━━
 Tracing 'my_function' (file: /home/psharpe/GitHub/memorylane/examples/make_report/make_reports.py:9):
 Mem:    100 MB | ΔMem:    100 MB | Peak:    100 MB | ΔPeak:    100 MB | make_reports.py:11   |     x = torch.randn(5120, 5120, device="cuda")
 Mem:    108 MB | ΔMem:      8 MB | Peak:    208 MB | ΔPeak:    108 MB | make_reports.py:12   |     x = x @ x
 Mem:    108 MB | ΔMem:      0 MB | Peak:    208 MB | ΔPeak:      0 MB | make_reports.py:13   |     x = x.relu()
 Mem:      8 MB | ΔMem:   -100 MB | Peak:    208 MB | ΔPeak:      0 MB | make_reports.py:14   |     x = x.mean()
 Mem:      8 MB | ΔMem:      0 MB | Peak:    208 MB | ΔPeak:      0 MB | make_reports.py:15   |     return x



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, 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!