retain_graph=True 会阻止PyTorch在backward后释放计算图,导致中间变量持续占用显存并线性累积直至OOM;其泄漏本质是跨step的强引用滞留,需用memory_snapshot定位重复backward栈及_saved_*字段。

retain_graph=True 会阻止计算图自动释放
PyTorch 默认在 loss.backward() 执行后立即释放整个计算图——包括所有中间变量(如前向传播中保存的输入、权重、激活值等)。一旦释放,这些张量占用的显存就可被复用或回收。但设置 retain_graph=True 会强制保留整条计算图及其所有 _saved_* 变量,哪怕你后续只读取梯度、不重复反向传播,这些内存也一直挂着。
泄漏不是“一次多占”,而是“跨 step 累积”
问题常出现在训练循环里误用 retain_graph=True,比如:
- 本不需要多次反向传播,却加了
retain_graph=True - 在日志收集或调试时,把带图的
loss或output直接塞进全局列表(如loss_history.append(loss)) - 验证阶段没关梯度,又手动加了
retain_graph=True做分析
每次迭代都保留一份完整图,而 Python 列表持有对这些图节点的强引用,gc.collect() 也无法清理——因为图节点仍被列表项引用着。显存就随 step 线性增长,直到 OOM。
和 detach() / item() 的行为冲突极易被忽略
很多人以为“只要最后调用了 .item() 就安全”,但错了:如果中间某步用了 retain_graph=True,那之前所有已构建的图节点都不会释放,哪怕你后面只取标量。
-
loss.item()只返回 Python float,不释放图本身 -
loss.detach()断开梯度链,但若图已被retain_graph=True锁住,detached 张量仍指向原存储 - 真正释放图的唯一时机是:最后一次
backward()且retain_graph=False(默认)
排查时别只看 memory_allocated,要盯 memory_snapshot
torch.cuda.memory_allocated() 显示的是当前活跃张量占用,但 retain_graph 泄漏往往藏在“已分配但未释放”的缓存块里。更可靠的做法是:
- 每 50 步执行一次
torch.cuda.memory_snapshot() - 用
snapshot[0].traceback定位哪行代码反复申请了大块显存 - 重点检查是否出现大量重复的
backward调用栈 +_saved_*字段
最隐蔽的泄漏点,往往不是某一行写错了,而是你在某个调试分支里加了一次 retain_graph=True,然后忘了删——它会让之后所有迭代都背负额外图开销,且几乎不报错。
Python免费学习笔记(深入):立即使用
在学习笔记中,你将探索 Python 的核心概念和高级技巧!











