如何通过gc.get_objects()定位PyTorch训练中引发内存泄漏的新建对象

第一步：通过gc快照定位高占用的泄漏对象类型 #

直接对比两个epoch节点的对象差异，过滤临时干扰项后锁定嫌疑对象类型，操作逻辑如下：

• 统一采样时机：每次采样前必须先执行gc.collect()等待垃圾回收完成，再抓取对象快照，避免未回收的临时对象造成差值误判。采样间隔至少覆盖1-2个完整epoch，不要在epoch中间采样，否则会把每步迭代的临时计算对象算入差异。
• 快照预处理逻辑：不对原始gc.get_objects()返回的列表直接做差，而是按对象类型聚合计数、累计内存占用，过滤掉差值小于1MB的低影响项。
  可直接复用的采样对比代码：

  import gc
  import sys
  from collections import Counter
  
  def get_snapshot():
      gc.collect()
      obj_list = gc.get_objects()
      type_count = Counter()
      type_mem = Counter()
      for obj in obj_list:
          t = type(obj)
          type_count[t] += 1
          type_mem[t] += sys.getsizeof(obj)
      return type_count, type_mem
  
  # 第N个epoch结束（完成固定清理逻辑后）采样
  count1, mem1 = get_snapshot()
  # 连续跑3个完整epoch
  # 第N+3个epoch结束（同样完成固定清理逻辑后）采样
  count2, mem2 = get_snapshot()
  
  # 计算类型维度的差值，按内存占用降序排列得到泄漏候选
  leak_candidates = sorted(
      [
          (t, count2[t] - count1[t], mem2[t] - mem1[t])
          for t in set(count1) | set(count2)
          if (count2[t] - count1[t]) > 0 and (mem2[t] - mem1[t]) > 1024 * 1024
      ],
      key=lambda x: -x[2]
  )
  

• 结果解读优先级：如果候选列表里出现torch.Tensor、自定义模型类、数据迭代器、计算图节点，优先排查，这类是PyTorch训练中最常见的泄漏来源。
  注意sys.getsizeof()仅统计对象本身内存，无法计算嵌套引用的关联内存，对自定义类、容器类对象的实际占用会有低估。

第二步：定位产生泄漏对象的具体代码行 #

通过内存追踪工具直接定位对象的创建位置，无需逐行排查：

• 内置tracemalloc方案（无需额外安装，粒度到代码行）：
  程序启动第一行开启追踪，在和上述gc快照完全相同的两个epoch节点采样，直接输出两个节点间内存增长对应的代码行：

  # 程序最开头开启追踪
  import tracemalloc
  tracemalloc.start()
  
  # 第一个快照（第N个epoch结束时）
  snap1 = tracemalloc.take_snapshot()
  # 跑3个完整epoch
  # 第二个快照（第N+3个epoch结束时）
  snap2 = tracemalloc.take_snapshot()
  
  # 按代码行维度对比内存差异，打印前20个增长最多的代码位置
  top_diff = snap2.compare_to(snap1, 'lineno')
  for stat in top_diff[:20]:
      print(stat)
  

  该方法可精准定位CPU侧Python对象、CPU张量的创建代码行；GPU侧的内存增长需要配合PyTorch内置的torch.cuda.memory接口做专项追踪。
• PyTorch场景专属高频漏点排查：
  对照以下常见场景逐行核对训练循环代码：
  • 是否将每步的loss、输出张量通过append写入全局列表，且未对张量做.detach()处理，导致计算图被全局引用持有无法释放
  • 反向传播后是否正确执行optimizer.zero_grad(set_to_none=True)，是否存在未清空的梯度累积引用
  • 自定义PyTorch模块是否在forward中动态新增Parameter对象，未做固定缓存复用
• 引用链定位方案（针对已知泄漏对象类型）：
  锁定第一步识别出的泄漏类型后，顺着引用链找到持有该对象的全局代码位置：

  import objgraph
  # 取5个存活的该类泄漏对象，打印引用关系
  leak_objs = [obj for obj in gc.get_objects() if isinstance(obj, 目标泄漏类型)]
  objgraph.show_backrefs(leak_objs[:5], max_depth=10)
  

  顺着输出的引用层级向上追溯，即可定位到持有泄漏对象、未做释放的变量所在代码行。

需注意：需区分正常内存增长和真泄漏：PyTorch内置的内存缓存分配器会复用已分配内存，不会及时把内存还给操作系统，这类增长的特征是内存上升到固定阈值后进入平稳期，不会随epoch数持续线性上涨。真泄漏的特征是每epoch固定涨数百MB，无明显平台期。

内容的提问来源于stack exchange，提问作者AlphaBetaGamma96