PyTorch猫狗分类训练加验证时CUDA内存不足问题求助

核心排查与解决方法 #

从报错信息看，PyTorch已预留7.29GiB显存，已分配达7.21GiB，剩余显存为0，说明显存碎片化或验证流程存在显存未释放/泄漏问题。以下是针对性方案：

1. 强制关闭验证阶段的梯度计算 #

验证不需要计算梯度，必须用torch.no_grad()包裹整个验证逻辑，避免不必要的显存占用：

model.eval()  # 切换到验证模式，禁用Dropout/BatchNorm训练行为
with torch.no_grad():
    for data in val_loader:
        images, labels = data
        images = images.to('cuda')
        outputs = model(images)
        # 后续验证指标计算逻辑...
model.train()  # 验证完成切回训练模式

PyTorch 1.9+可改用torch.inference_mode()，能更彻底地禁用梯度相关显存分配。

2. 清理验证循环内的临时变量 #

避免在循环中累积大张量，用完即删除并手动清理显存：

val_acc = 0.0
model.eval()
with torch.no_grad():
    for data in val_loader:
        images, labels = data
        images = images.to('cuda')
        labels = labels.to('cuda')
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
        val_acc += torch.sum(preds == labels.data)
        
        # 手动删除临时变量，释放显存
        del outputs, preds, images, labels
        torch.cuda.empty_cache()
val_acc = val_acc / len(val_dataset)

3. 优化验证数据加载器配置 #

• 将验证集的pin_memory设为False，避免额外占用GPU显存；
• 关闭验证集的shuffle，减少数据加载时的临时显存波动。

4. 缓解显存碎片化 #

按报错提示设置环境变量，调整显存块拆分阈值：

• Linux/macOS终端：

  export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
  

• Windows命令行：

  set PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128
  

该参数能减少因碎片化导致的无法分配连续显存问题。

5. 定位显存飙升节点 #

在训练和验证的关键步骤打印显存使用详情，精准定位问题：

print(torch.cuda.memory_summary())

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