Tensor是如何让你的内存/显存泄漏的_torch tensorch 内存-CSDN博客

本文链接：https://blog.csdn.net/qq_29788741/article/details/125762734

本文详细分析了代码库中模型间内存泄漏差异，通过最小复现代码揭示了Tensor在动态图中保存AutogradMeta导致的内存泄漏。重点讲解了如何通过源码追踪，解释为何`requires_grad=True`引发内存泄漏，并讨论了Python代码中看似无害却可能导致内存问题的实例。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

某codebase出现了奇怪的泄漏现象，奇怪的点有以下几个方面：

（1）不同的模型，内存/显存泄漏的现象不一样。比如A模型和B模型泄露的速度是不一样的

（2）训练同一个模型的时候，如果在dataset中增加了数据量，相比不加数据，会在更早的epoch就把内存泄漏完。

是不是听起来现象非常离谱，本着”code never lies“的世界观，我开始探求这个现象的真正原因。

要想解决一个大的问题，首先就要降低问题的复杂度。最小复现代码是我们找问题的基础，而这个写最小复现代码的过程其实也是遵循了一定套路的，此处一并分享给大家：

如果突然出现了历史上没有出现过的问题（比如在某个版本之后突然内存开始泄漏了），用git bisect找到 first bad commit（前提项目管理的比较科学，不会出现很多feature杂糅在一个commit里面；还有就是git checkout之后复现问题的成本不高）。如果bisect大法失效，考虑下面的复现流程。
首先排除data的问题，也就是只创建一个dataloader，让这个loader不停地供数据，看看内存会不会涨（通常data是一系列对不上点、内存泄漏的重灾区）。
其次排除训练的问题，找一个固定数据，不停地让网络训练固定数据进行，看看是否发生泄漏。这一步主要是检查模型、优化器等组件的问题（通常模型本身不会发生泄漏，这一步经常能查出来一些自定义op的case）
最后就是检查一些外围组件了。比如各种自己写的utils/misc的内容。这块通常不是啥重灾区。

最后给出来我的最小复现（psutil需要通过pip安装一下）：

import torchimport osimport psutil

def log_device_usage(count, use_cuda):    mem_Mb = psutil.Process(os.getpid()).memory_info().rss / 1024 ** 2    cuda_mem_Mb = torch.cuda.memory_allocated(0) / 1024 ** 2 if use_cuda else 0    print(f"iter {count}, mem: {int(mem_Mb)}Mb, gpu mem:{int(cuda_mem_Mb)}Mb")

def leak():    use_cuda = torch.cuda.is_available()    val = torch.rand(1).cuda() if use_cuda else torch.rand(1)    count = 0    log_iter = 20000    while True:        value = torch.rand(1).cuda() if use_cuda else torch.rand(1)        val += value.requires_grad_()        if count % log_iter == 0:            log_device_usage(count, use_cuda)        count += 1

if __name__ == "__main__":    leak()

试着运