使用PyTorch时，如何避免遇CUDA内存不足重启Jupyter Notebook？

• 强制释放未引用张量与CUDA缓存
  直接删除变量后，Jupyter的历史记录（如_、In/Out列表）可能保留引用，导致内存无法回收。执行以下代码彻底清理：

  import gc
  # 删除指定的大变量/模型
  del large_tensor, train_model
  # 触发Python垃圾回收
  gc.collect()
  # 清空CUDA未使用的缓存
  torch.cuda.empty_cache()
  

  同时可通过Jupyter菜单栏Cell > All Output > Clear清空输出历史，避免隐藏引用。

• 用上下文管理器隔离内存使用
  对于推理或临时大张量运算，使用torch.no_grad()关闭梯度计算，减少内存占用；临时变量会在上下文结束后更快被回收：

  with torch.no_grad():
      # 执行推理或大张量运算
      inference_output = pre_trained_model(input_data)
  # 离开上下文后，inference_output的内存可被及时释放
  

• 定位隐藏的内存占用
  运行以下代码找出所有在CUDA上的张量，排查意外引用的变量：

  import gc
  import torch
  
  for obj in gc.get_objects():
      try:
          if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
              if obj.is_cuda:
                  print(f"CUDA张量类型: {type(obj)}, 尺寸: {obj.size()}")
      except Exception:
          pass
  

  根据输出找到未被删除的大张量，针对性删除并触发回收。

• 使用Jupyter魔术命令清理命名空间
  若全局变量过多难以逐个删除，执行%reset -f强制清空当前命名空间，再清理CUDA缓存：

  %reset -f
  import gc
  import torch
  gc.collect()
  torch.cuda.empty_cache()
  

  注意：此操作会删除所有变量，需提前保存必要结果。

• 优化模型与运算的内存效率

  • 训练时开启混合精度：使用torch.cuda.amp.GradScaler()与torch.autocast(device_type='cuda')，大幅降低内存占用；
  • 拆分大批次：将大批次数据拆分为小批次处理，或用梯度累积替代大批次训练；
  • 简化模型：移除不必要的偏置层、减少模型参数量，避免创建冗余中间张量。

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