PyTorch训练CUDA内存不足求助，已尝试缓存清理仍未解决

你已经尝试了初始化时的内存清理，但仍在loss.backward()阶段触发CUDA OOM，结合报错里的内存占用细节（PyTorch已分配8.96 GiB、存在未分配的保留内存），可以从以下几个方向解决：

1. 修复CUDA内存碎片问题 #

报错提示内存碎片可能是诱因，通过设置PyTorch的CUDA分配参数优化内存利用：

import os
# 设置max_split_size_mb，值可根据情况调整（如64、256）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

建议在导入torch前设置这个环境变量，或者重启内核后添加该代码再执行训练。

2. 训练过程中动态清理内存 #

仅在代码开头执行清理是不够的，需在训练循环中针对性释放内存：

for epoch in range(epochs):
    for input_tensor, mask in dataloader:
        input_tensor, mask = input_tensor.cuda(), mask.cuda()
        optim.zero_grad()
        
        # 前向传播、计算损失
        predicted_img = model(input_tensor)
        loss_mask = torch.mean((predicted_img - input_tensor) ** 2 * mask / mask_ratio)
        
        # 反向传播与优化
        loss_mask.backward()
        optim.step()
        
        # 清理当前batch的资源
        loss_mask.detach()  # 解除损失张量与计算图的关联
        # 显式删除不再使用的张量
        del loss_mask, predicted_img
        gc.collect()
        torch.cuda.empty_cache()

同时注意：Jupyter单元格执行后变量会残留，训练前重启内核，避免之前的模型、张量占用内存。

3. 核心内存优化手段 #

• 减小Batch Size：直接降低单次迭代的内存占用，这是最快速有效的方法。
• 启用混合精度训练：通过torch.cuda.amp将部分张量转为半精度，大幅减少内存消耗：

  from torch.cuda.amp import GradScaler, autocast
  
  scaler = GradScaler()
  for input_tensor, mask in dataloader:
      input_tensor, mask = input_tensor.cuda(), mask.cuda()
      optim.zero_grad()
      
      with autocast():
          predicted_img = model(input_tensor)
          loss_mask = torch.mean((predicted_img - input_tensor) ** 2 * mask / mask_ratio)
      
      scaler.scale(loss_mask).backward()
      scaler.step(optim)
      scaler.update()
  

• 冻结冗余模型参数：如果是微调模型，冻结不需要训练的底层参数，减少梯度计算的内存占用：

  # 示例：冻结模型backbone层
  for param in model.backbone.parameters():
      param.requires_grad = False
  

• 检查GPU张量冗余：确保所有临时张量用完即删除，避免不必要的GPU内存占用。

4. Jupyter专属优化 #

• 将完整训练代码放在单个单元格执行，避免多单元格执行后残留大量未使用变量。
• 可使用%reset -f命令清空当前命名空间的变量（执行前确认无需要保留的内容）。

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