讲师中心 微信公众号
AI工具推荐 视频效率加速

Python怎么定位PyTorch代码中的显存泄漏点_使用torch.cuda.memory_summary诊断

风枫酱_9496

风枫酱_9496

发布时间:2026-04-27 16:36:19

|

871人浏览过

|

来源于php中文网

原创

显存泄漏的典型表现是PyTorch中torch.cuda.memory_allocated()持续上升,常见诱因包括循环中累积GPU tensor、未释放的hook或buffer、多进程dataloader引用等;可用torch.cuda.memory_summary()对比分配快照定位泄漏点,配合memory_history和引用计数辅助排查。

python怎么定位pytorch代码中的显存泄漏点_使用torch.cuda.memory_summary诊断

显存泄漏的典型表现和初步确认

PyTorch中显存持续增长、torch.cuda.memory_allocated() 不断上升,但模型训练/推理逻辑没变,GPU显存占用却越跑越高——这往往是显存泄漏的信号。注意:不是所有“显存没释放”都算泄漏,比如 torch.no_grad() 块里创建的 tensor 若被意外保留引用(如塞进全局 list),就可能阻断回收。

常见诱因包括:

  • 在循环中不断 append GPU tensor 到 Python list 或 dict
  • 使用 tensor.detach().cpu() 后仍保留对原始 GPU tensor 的引用
  • 自定义 nn.Module 中缓存了未清空的中间结果(如 hooks、buffer)
  • 多进程 dataloader 里 worker 持有 tensor 引用(尤其用了 pin_memory=True 但没及时 .item() 或 .cpu())

用 torch.cuda.memory_summary() 定位泄漏源头

这个函数不返回结构化数据,而是打印人类可读的显存分配快照,关键在对比「两次调用之间新增了哪些 persistent allocations」。

执行建议:

立即学习“Python免费学习笔记(深入)”;

python学习
python学习

Python 编程学习助手 skill,面向零基础用户。 当用户表达以下意图时触发: - 学习 Python 基础概念(变量、数据类型、循环、函数、类等) - 请求 Python 练习题或测验 - 请求审查或调试 Python 代码 - 请求 Python 项目实战指导(爬虫、数据分析、Web 开发等入门项目) - 泛化表述如"教我 Python""我想学编程""Python 怎么写""帮我学 Python" 此 skill

下载
  • 在疑似泄漏点前后各调一次:torch.cuda.memory_summary(device=None, abbreviated=False)
  • 加 abbreviated=False 确保显示每一块 allocation 的 size 和 stack trace(需编译时启用了 debug info,conda/pip 官方包默认支持)
  • 如果看到大量重复出现的 stack trace(尤其含 <strong>call</strong>、forward、collate_fn),基本就是泄漏发生处

示例片段:

# 在 epoch 开始前
print("Before loop:")
torch.cuda.memory_summary()
<p>for i, (x, y) in enumerate(dataloader):
out = model(x.cuda())</p><h1>❌ 错误:把 GPU tensor 直接存下来</h1><pre class='brush:python;toolbar:false;'>history.append(out)  # ← 这里会累积显存

# ✅ 正确:只存标量或 CPU 数据
loss_history.append(out.mean().item())

print("After loop:") torch.cuda.memory_summary()

为什么 memory_summary() 有时看不到 stack trace?

不是所有 allocation 都带 traceback —— 只有通过 PyTorch C++ 分配器(如 at::empty, at::randn)创建的 tensor 才能关联 Python 调用栈;而底层 CUDA malloc(如自定义 kernel、第三方库调用)或某些 inplace 操作(tensor.add_(...))可能不记录。

排查时注意:

  • 确保用的是 debug 版本 PyTorch(torch.<strong>version</strong> 含 +cu 但不含 dev 通常也够,pip install torch --force-reinstall --no-deps 可重装带符号版本)
  • 关闭 torch.backends.cudnn.benchmark = True,避免 cudnn 缓存干扰分配模式
  • 运行前加 torch.cuda.memory._record_memory_history(enabled=True, max_entries=100000),之后可用 torch.cuda.memory._dump_snapshot("snapshot.pickle") 导出供离线分析

比 memory_summary() 更细粒度的辅助手段

单靠打印 summary 容易漏掉“缓慢泄漏”,建议组合使用:

  • 每轮迭代后检查:torch.cuda.memory_allocated() 和 torch.cuda.memory_reserved() 是否单调增长
  • 用 gc.collect() + torch.cuda.empty_cache() 强制触发清理,再看是否回落——不回落说明仍有活跃引用
  • 对可疑对象做引用计数检查:sys.getrefcount(obj)(注意该函数本身会+1,需对比基准值)
  • 在关键位置插入 torch.cuda.memory._get_memory_history() 获取最近分配记录(需提前开启 history)

真正难缠的泄漏往往藏在跨函数边界的 tensor 生命周期管理里,比如一个 hook 函数里保存了 input tensor,而 hook 没被 remove;或者分布式训练中 torch.distributed.all_gather 返回的 list 里混入了未 .cpu() 的 tensor。这些地方 memory_summary() 能指方向,但得靠人工顺着 stack trace 查引用链。

热门AI工具

更多
UP简历
UP简历 Hot

一款AI办公效率工具,主要用于基于AI技术的免费在线简历制作工具,适合需要提升相关任务效率的用户。

墨刀AI
墨刀AI Hot

一款AI图像与设计工具,主要用于产品经理的专属智能体,适合需要提升相关任务效率的用户。

蛙蛙写作

一款AI论文写作工具,主要用于超级AI智能写作助手,适合需要提升相关任务效率的用户。

Atoms
Atoms Hot

Atoms是一款AI智能体工具,第一支自动构建真实业务的 AI 团队。

DeepSeek

DeepSeek是一款面向对话、写作、编程和推理场景的AI大模型工具。

讯飞绘文

讯飞绘文是一款由科大讯飞推出的一站式 AIGC 内容运营平台。

WorkBuddy

一款AI办公效率工具,主要用于腾讯云推出的AI原生桌面智能体工作台,适合需要提升相关任务效率的用户。

二狗PPT
二狗PPT Hot

一款AI演示文稿工具,主要用于专为中式职场打造的AI PPT生成工具,适合需要提升相关任务效率的用户。

豆包大模型

豆包大模型是一款由字节跳动推出的企业级大语言模型服务平台。

相关专题

更多
python打包成可执行文件
python打包成可执行文件

本专题为大家带来python打包成可执行文件相关的文章,大家可以免费的下载体验。

1791

2023.07.20

python能做什么
python能做什么

python能做的有:可用于开发基于控制台的应用程序、多媒体部分开发、用于开发基于Web的应用程序、使用python处理数据、系统编程等等。本专题为大家提供python相关的各种文章、以及下载和课程。

4684

2023.07.25

format在python中的用法
format在python中的用法

Python中的format是一种字符串格式化方法,用于将变量或值插入到字符串中的占位符位置。通过format方法,我们可以动态地构建字符串,使其包含不同值。php中文网给大家带来了相关的教程以及文章,欢迎大家前来阅读学习。

1749

2023.07.31

python教程
python教程

Python已成为一门网红语言,即使是在非编程开发者当中,也掀起了一股学习的热潮。本专题为大家带来python教程的相关文章,大家可以免费体验学习。

27717

2023.08.03

python环境变量的配置
python环境变量的配置

Python是一种流行的编程语言,被广泛用于软件开发、数据分析和科学计算等领域。在安装Python之后,我们需要配置环境变量,以便在任何位置都能够访问Python的可执行文件。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

3347

2023.08.04

python eval
python eval

eval函数是Python中一个非常强大的函数,它可以将字符串作为Python代码进行执行,实现动态编程的效果。然而,由于其潜在的安全风险和性能问题,需要谨慎使用。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

3347

2023.08.04

scratch和python区别
scratch和python区别

scratch和python的区别:1、scratch是一种专为初学者设计的图形化编程语言,python是一种文本编程语言;2、scratch使用的是基于积木的编程语法,python采用更加传统的文本编程语法等等。本专题为大家提供scratch和python相关的文章、下载、课程内容,供大家免费下载体验。

1203

2023.08.11

python合并两个列表
python合并两个列表

Python是一种强大的编程语言,具有许多方便的功能和工具。在Python中,有多种方法可以合并两个列表。php中文网给大家带来了相关的教程以及文章,欢迎大家前来学习阅读。

596

2023.08.10

Gradle插件开发入门教程
Gradle插件开发入门教程

本专题介绍Gradle插件开发的基本流程,涵盖插件类编写、自定义任务注册、扩展参数设计和插件应用方法,帮助你理解插件如何组织构建逻辑,减少多个项目之间的重复配置。

0

2026.10.22

热门下载

更多
网站特效
/
网站源码
/
网站素材
/
前端模板

精品课程

更多
相关推荐
/
热门推荐
/
最新课程
Uvicorn官方安装与入门手册
Uvicorn官方安装与入门手册

共0课时 | 0人学习

Python uv官方项目管理指南
Python uv官方项目管理指南

共0课时 | 0人学习

Python uv官方安装与入门文档
Python uv官方安装与入门文档

共0课时 | 0人学习

关于我们 免责申明 举报中心 意见反馈 讲师合作 广告合作 最新更新
php中文网:公益在线php培训,帮助PHP学习者快速成长!
关注服务号
PHP中文网订阅号
每天精选资源文章推送

Copyright 2014-2026 https://www.php.cn/ All Rights Reserved | php.cn