Ubuntu下PyTorch内存不足怎么解决

Ubuntu下PyTorch内存不足的排查与优化

一 快速定位问题

• 先看是GPU显存不足还是系统内存不足：用nvidia-smi查看显存占用与进程列表；若显存已满，优先释放或优化显存；若训练进程被系统“Killed”，多为系统内存或交换空间不足。必要时在保存大对象前用**torch.cuda.empty_cache()清理未使用缓存，验证阶段用with torch.no_grad()**避免保存中间梯度。若日志显示“Killed”，结合系统监控确认是否为内存压力导致。

二 解决GPU显存不足

• 降低Batch Size，必要时用梯度累积保持等效批量；验证/测试阶段使用torch.no_grad()；在合适位置调用**torch.cuda.empty_cache()**释放未使用缓存。
• 启用混合精度训练 AMP：用torch.cuda.amp.autocast与GradScaler降低显存占用并加速。
• 优化数据加载：设置合适的num_workers与预取，避免数据管线成为瓶颈。
• 减少显存碎片：设置环境变量PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:<值>（如512），缓解“reserved >> allocated”的碎片问题。
• 清理占用显存的进程：用nvidia-smi查PID，必要时kill -9 PID释放显存。
• 若仍不足，考虑更轻量的模型架构或升级GPU显存。

三 解决系统内存不足

• 释放Python与CUDA缓存：在代码中及时del不再使用的张量/中间变量，调用gc.collect()，并用**torch.cuda.empty_cache()**清理GPU缓存。
• 谨慎清理系统页面缓存：仅在必要时以root执行sync && echo 1 > /proc/sys/vm/drop_caches（或2/3），避免影响系统稳定性。
• 扩充Swap空间（应对大对象序列化如torch.save时被系统杀死）：创建交换文件并启用，例如创建约16GB的交换文件，然后激活并检查。

四 常用命令与代码片段

• 查看与清理GPU
  • 查看显存与进程：nvidia-smi
  • 结束进程：kill -9
• 释放缓存与无梯度推理
  • 清理缓存：torch.cuda.empty_cache()
  • 验证阶段：with torch.no_grad():
• 混合精度训练模板
  • 代码示例：
    • from torch.cuda.amp import autocast, GradScaler
    • scaler = GradScaler()
    • for data, target in loader:
      • optimizer.zero_grad()
      • with autocast():
        • output = model(data)
        • loss = criterion(output, target)
      • scaler.scale(loss).backward()
      • scaler.step(optimizer)
      • scaler.update()
• 缓解显存碎片
  • 环境变量：export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:512
• 扩充Swap（示例为约16GB）
  • 创建文件：sudo dd if=/dev/zero of=swapfile bs=1073741824 count=16
  • 设置权限并启用：sudo chmod 600 swapfile && sudo mkswap swapfile && sudo swapon swapfile
  • 验证：free -h 或 swapon --show
• 清理系统页面缓存（谨慎）
  • 执行：sync && echo 1 > /proc/sys/vm/drop_caches（需要root，按需使用）