如何为PyTorch增加GPU显存分配 解决训练RES101时CUDA显存不足问题

前置排查 #

先确认非配置错误导致的显存限制：

• 检查环境变量CUDA_VISIBLE_DEVICES是否设置正确：
  Linux下执行echo $CUDA_VISIBLE_DEVICES，Windows cmd下执行echo %CUDA_VISIBLE_DEVICES%，确认返回值为目标显卡ID（单卡场景一般为0）。如果设置错误，执行export CUDA_VISIBLE_DEVICES=0（Linux）或set CUDA_VISIBLE_DEVICES=0（Windows）修复，也可写入系统环境变量永久生效。
• 检查代码中是否手动设置了显存分配比例：搜索代码中是否存在torch.cuda.set_per_process_memory_fraction()相关配置，若参数为0.33左右恰好对应24GB显存的1/3即8GB，删除该配置或调整参数到0.9（预留10%显存供系统、驱动调用）即可。
• 验证PyTorch对显卡显存的识别是否正确：执行以下命令

  import torch
  print(torch.cuda.get_device_properties(0).total_memory / 1024**3)
  

  若返回值远小于24GB，说明PyTorch版本与本机CUDA版本不匹配，卸载现有PyTorch后安装对应CUDA版本的PyTorch即可修复。

显存分配调整配置 #

代码内修改 #

在训练代码导入torch后，第一时间添加以下配置：

import torch
# 开启显存动态分配，避免初始化时提前预留固定额度显存
torch.backends.cudnn.benchmark = True
# 设置单进程显存最大分配比例为95%，适配24GB显存场景
torch.cuda.set_per_process_memory_fraction(0.95, device=0)
# 清空已占用的空闲显存缓存
torch.cuda.empty_cache()

运行前操作 #

启动训练脚本前，通过nvidia-smi（Linux）或任务管理器GPU面板（Windows）关闭所有占用GPU显存的无关进程（如浏览器硬件加速、视频编辑软件、其他Python进程等），避免显存被抢占。

可选：显存占用优化（变相提升可用显存） #

如果调整分配比例后仍存在显存不足问题，可同时开启以下优化降低训练显存开销：

• 梯度累积：降低单步batch size，每N步更新一次梯度，在不改变训练效果的前提下降低单步显存占用
• 混合精度训练：使用torch.cuda.amp开启自动混合精度，可降低30%~50%显存占用
• 梯度检查点：开启RES101的梯度检查点配置，用少量计算开销换取40%左右的显存占用降低

效果验证 #

启动训练后，通过nvidia-smi查看对应Python进程的显存占用，确认可突破8GB上限、最高占用可达20GB以上即配置生效。

内容的提问来源于stack exchange，提问作者赵嘉琛