PyTorch内存泄漏排查实战：使用memray定位问题

最近在优化一个图像分类模型时，遇到了严重的内存泄漏问题。训练过程中，GPU内存逐渐增长直至显存溢出，即使使用了torch.cuda.empty_cache()也无法释放。为了解决这个问题，我决定使用memray工具进行深入排查。

环境准备

首先安装memray：

pip install memray

排查步骤

使用memray追踪内存分配：

import torch
import torch.nn as nn
from memray import AllocatorType, Tracker

# 模拟训练循环
model = nn.Sequential(
    nn.Conv2d(3, 64, 3),
    nn.ReLU(),
    nn.AdaptiveAvgPool2d((1, 1)),
    nn.Flatten(),
    nn.Linear(64, 10)
).cuda()

optimizer = torch.optim.Adam(model.parameters())

with Tracker('trace.bin', allocator=AllocatorType.CUDA):
    for i in range(100):  # 模拟100个batch
        x = torch.randn(32, 3, 224, 224).cuda()
        y = model(x)
        loss = y.sum()
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

分析追踪结果：

memray flamegraph trace.bin --output flamegraph.html

排查结果

通过火焰图发现，内存增长主要来自于模型参数的累积梯度。在loss.backward()后，没有正确清理中间变量。解决方法是添加torch.cuda.empty_cache()并使用del显式删除不需要的张量。

性能对比

方法	内存峰值(MB)	运行时间(s)
无优化	2400	120
添加del + empty_cache	800	115

最终问题得到解决，内存使用稳定在合理范围。这个经验对于大规模训练任务很有价值。

紫色风铃 · 2026-01-08T10:24:58

memray确实是个好工具，能精准定位GPU内存泄漏点。建议在模型训练中加入定期的`torch.cuda.empty_cache()`，避免显存持续增长。

TallTara · 2026-01-08T10:24:58

代码里加了`del`和`empty_cache`后效果明显，但要注意别删得太早导致计算图断开。可以结合`torch.autograd.set_detect_anomaly(True)`排查梯度问题。

WiseNinja · 2026-01-08T10:24:58

火焰图分析很直观，我之前也遇到过类似问题，主要是optimizer.step()后没及时清理梯度。建议把优化逻辑封装成函数，便于统一管理。

深海游鱼姬 · 2026-01-08T10:24:58

这种内存泄漏问题在大模型训练中特别常见，除了工具排查外，还可以用`torch.utils.checkpoint`来节省显存，减少中间变量的累积

PyTorch内存泄漏排查实战：使用memray定位问题