Stonesjtu/pytorch_memlab

pytorch_memlab – PyTorch 的轻量级 CUDA 内存调试工具箱

它是啥 – 一个纯 Python 包，可让你查看 PyTorch 代码中每一行消耗了多少 GPU 内存，列出设备上占用的张量，并可临时将所有数据移至 CPU。它既支持命令行，也支持通过魔法命令在 Jupyter/IPython 中使用。

主要功能

• 逐行内存分析器 (LineProfiler / @profile 装饰器) – 类似于 line_profiler，但会报告每行源码的 CUDA active_bytes 和 reserved_bytes。
• 内存报告器 (MemReporter) – 列出所有活跃的 torch.Tensor 对象，显示底层 UntypedStorage 的大小，并可处理共享权重、梯度，以及详细的存储复用图示。
• 礼让模式 – 一个辅助类，可按需将所有 GPU 张量移回 CPU，有助于在不保存检查点的情况下释放 GPU。
• IPython 集成 – %mlrun 和 %%mlrun 魔法命令可让你仅用一条命令就对函数或整个笔记本单元进行分析。
• GPU 选择 – set_target_gpu 允许你在运行时切换分析器监控的设备。

安装方法

# 从 PyPI 安装稳定版
pip install pytorch_memlab

# 或从 GitHub 直接安装最新代码
pip install git+https://github.com/stonesjtu/pytorch_memlab

(若需使用 IPython 魔法命令，请安装可选依赖：pip install pytorch_memlab[ipython]。)

典型工作流程

1. 分析函数

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

   脚本运行结束后，分析器会输出一个表格，显示每一行分配了多少 GPU 内存。
2. 检查活跃张量

   from pytorch_memlab import MemReporter
   reporter = MemReporter(model)   # model 为可选参数
   reporter.report()               # 输出简洁表格
   reporter.report(verbose=True)  # 显示存储共享箭头
   

3. 在笔记本中使用

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

4. 临时释放 GPU

   from pytorch_memlab import Courtesy
   c = Courtesy()
   c.yield_memory()   # 将张量移至 CPU
   # …等待信号…
   c.restore()        # 将其移回
   

为何重要 – 在开发深度学习模型时，内存不足（OOM）崩溃是常见痛点。pytorch_memlab 使隐藏的 CUDA 分配行为变得可见，帮助你发现低效的张量生命周期、意外的存储共享，或被遗忘的中间缓冲区。

限制 / 已知问题

• 分析器只能看到 Python 知道的张量；autograd 创建的底层 C 层缓冲区会被报告为“不可见”，但仍会占用内存。
• 即使所有张量都在 CPU 上，PyTorch CUDA 上下文本身仍会保留约 1 GB 的 GPU 内存，目前无法由该库释放。
• 仅支持 CUDA 兼容 GPU；不支持 AMD 或纯 CPU 运行。

项目状态 – 活跃维护中（最新版本 0.3.2，2026 年 6 月）。CI 运行测试并发布轮子包到 PyPI。README 包含使用示例、演示笔记本和变更日志。

---

以上所有信息均直接来自仓库的 README，未添加任何外部声明。