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 进行分片训练。