如何在Ubuntu上解决PyTorch的内存不足问题

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

一 快速定位与监控

• 使用 nvidia-smi 实时查看显存占用与进程列表，定位是否有其他进程占用大量显存：
  • 持续监控：运行命令：watch -n 1 nvidia-smi
  • 结束无关进程：先用 nvidia-smi 找到占用显存的 PID，再执行：kill -9
• 在 Python 中打印显存与张量占用，辅助定位内存热点：
  • 查看显存：torch.cuda.memory_allocated() / torch.cuda.memory_reserved()
  • 清理缓存：torch.cuda.empty_cache()（仅释放未使用的缓存，不会释放被张量占用的显存）
• 若报错提示 CUDA out of memory 且伴随“无法为新分配的内存找到足够空间”，常见原因是显存碎片化或他进程占用，可结合清理缓存、减少并行度、设置环境变量等手段缓解。

二 训练阶段的高效优化

• 降低批量大小（batch_size）：最直接有效；若影响收敛，可配合梯度累积。
• 梯度累积（Gradient Accumulation）：用小批量多次累积梯度后再更新，等效增大批量：
  • 伪代码：每步 loss.backward() 后累积；每 N 步执行一次 optimizer.step() 与 optimizer.zero_grad()
• 自动混合精度（AMP）：用 torch.cuda.amp 在 FP16/BF16 与 FP32 间自动切换，显著降低显存占用并常加速训练：
  • 示例：
    • from torch.cuda.amp import autocast, GradScaler
    • scaler = GradScaler()
    • with autocast(): outputs = model(inputs); loss = criterion(outputs, targets)
    • scaler.scale(loss).backward(); scaler.step(optimizer); scaler.update()
• 低精度数据类型：在不使用 AMP 时，可将模型与数据转为 torch.float16 / torch.bfloat16（需设备支持）。
• 梯度检查点（Gradient Checkpointing）：以计算换显存，常见可减少激活占用约 40–50%：
  • 示例：from torch.utils.checkpoint import checkpoint；output = checkpoint(segment_fn, x)
• 激活与参数卸载（Offloading）：将部分中间激活或参数临时移至 CPU，缓解单卡显存瓶颈（可结合 DeepSpeed、Fabric 等框架）。
• 张量分片与分布式训练：使用 FSDP（Fully Sharded Data Parallel） 将参数/梯度/优化器状态分片至多卡，单卡显存需求可显著下降（部分场景可达约 10× 内存降低）。

三 数据与加载优化

• 高效数据加载：
  • 设置 pin_memory=True 与合适的 num_workers，加速主机到设备的传输并减少数据瓶颈。
  • 使用生成器/迭代器与流式读取，避免一次性将整个数据集加载进内存。
• 避免不必要的数据复制：优先使用张量视图与就地操作（如 add_），减少临时张量分配。
• 内存分析：在 Python 中用 sys.getsizeof / psutil 监控对象与进程内存；必要时用 valgrind 排查本地内存问题。

四 系统与CUDA层面的调优

• 缓解显存碎片化：设置环境变量 PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True，允许使用可扩展内存段，降低碎片导致的分配失败概率。
• 释放无关进程占用的显存：通过 nvidia-smi 结束占用显存的无关训练/推理进程，避免资源争用。
• 谨慎清理系统级缓存（仅在必要时）：
  • 查看内存：free -h
  • 写入不同级别以释放页缓存/目录项/索引节点：echo 1/2/3 > /proc/sys/vm/drop_caches（需 root，可能影响系统性能，不建议频繁执行）。

五 仍不足时的兜底方案

• 扩充系统内存与交换空间（Swap）：当 torch.save 或数据预处理阶段因系统内存不足被 OOM Killer 终止时，可临时增加 Swap 以完成保存或预处理：
  • 示例：创建 16GB 交换文件
    • sudo dd if=/dev/zero of=/swapfile bs=1G count=16
    • sudo mkswap /swapfile & & sudo swapon /swapfile
  • 注意：Swap 速度远慢于物理内存，仅作为临时或过渡方案。
• 升级硬件与模型策略：增加 GPU 显存 或改用更轻量的模型架构（如 MobileNet/EfficientNet 相对 ResNet/VGG 更省显存）；在单卡无法容纳时，采用 多卡/FSDP 进行分片训练。

声明：本文内容由网友自发贡献，本站不承担相应法律责任。对本内容有异议或投诉，请联系2913721942#qq.com核实处理，我们将尽快回复您，谢谢合作！