PyTorch模型重复初始化引发GPU内存泄漏问题求助

1. 内存持续增长的原因 #

• Python引用未及时回收：当你执行my_model = model(params)时，旧的模型对象并不会立即销毁——Python的垃圾回收（GC）是惰性机制，只有当对象的引用计数降为0时才会被回收。第一次移至CUDA的模型张量，因存在未清理的隐式引用（比如计算图残留、梯度缓存，或是你在列表、全局变量中保存了旧模型的引用），导致GPU内存无法释放。
• torch.cuda.empty_cache()的局限性：该函数仅会释放已被Python GC回收的GPU内存块，如果旧模型对象还未被GC回收，对应GPU内存不会被清空。你之前的操作仅处理了最后一次的模型，旧模型的引用未清理，因此内存仅释放了最后一次分配的部分。
• del的作用有限：del my_model只是删除了my_model变量名对模型对象的引用，但如果还有其他引用存在（比如PyTorch内部临时缓存、自定义的其他关联变量），对象依然会留在内存中，GPU内存也不会释放。

2. 释放泄漏内存的解决方案 #

按以下步骤严格执行，确保旧模型的GPU内存被完全释放：

• 清理旧模型的GPU关联资源：在重新赋值新模型前，先将旧模型移回CPU，并清空梯度（若模型经过训练）：

  # 清空模型梯度并释放梯度内存
  if hasattr(my_model, 'parameters'):
      for param in my_model.parameters():
          param.grad = None
  my_model.to('cpu')
  

• 删除所有引用并触发GC：

  del my_model
  # 强制触发Python垃圾回收，确保旧模型对象被销毁
  import gc
  gc.collect()
  

• 清空CUDA缓存：

  torch.cuda.empty_cache()
  

• 排查隐式引用：检查代码中是否将模型存入列表、字典、全局变量等容器，需同时清空这些容器中的旧模型引用：

  # 示例：若旧模型存在全局列表中
  if 'model_list' in globals() and my_model in model_list:
      model_list.remove(my_model)
  

• 使用内存诊断工具定位问题：若仍存在泄漏，用PyTorch工具排查：

  # 打印详细内存使用情况
  print(torch.cuda.memory_summary())
  # 生成内存快照，查看未释放的张量
  snapshot = torch.cuda.memory_snapshot()
  

• 若使用了torch.compile：PyTorch 2.x的编译模型会有额外缓存，需先卸载编译：

  my_model = my_model.uncompile()
  

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