Ubuntu上PyTorch内存不足怎么解决

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

一快速定位问题

区分是GPU显存不足还是系统内存不足：
- GPU显存：运行命令查看nvidia-smi，关注显存占用与进程列表；在代码中用**torch.cuda.memory_summary()**查看分配与缓存。
- 系统内存：运行命令查看free -h，观察可用内存与Swap使用情况。
常见现象与对应线索：
- 训练/推理报CUDA out of memory多为GPU显存不足。
- **torch.save()**或大对象序列化时报“Killed”多为系统内存或Swap不足，系统OOM Killer终止了进程。
建议先做一次“单步占用评估”：在模型创建后、前向前、反向后分别打印torch.cuda.memory_allocated() / memory_reserved()，定位占用峰值位置。

二通用优化清单（GPU显存）

降低Batch Size，必要时配合梯度累积（保持“虚拟批量”同时减少显存）。
启用自动混合精度 AMP：用torch.cuda.amp.autocast() + GradScaler()，在保持精度的同时显著降低显存占用并提速。
使用梯度检查点（Gradient Checkpointing）：以计算换显存，特别适合Transformer/深网络。
优化数据加载：设置合适的num_workers、开启pin_memory=True，避免数据预处理成为瓶颈。
减少优化器状态占用：在允许的情况下用SGD替代Adam（Adam为每个参数维护动量与方差，状态翻倍）。
降低激活/参数占用：结合FSDP（FullyShardedDataParallel）进行张量分片；或使用CPU/RAM卸载（如DeepSpeed ZeRO-Inference/Offload）。
缓解显存碎片：设置环境变量PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:<值>（如512），在“reserved >> allocated”时尤为有效。
推理场景的分块/切片：如超分、生成类模型可用**–tile**等参数降低单次处理分辨率，显著降低峰值显存。

三系统内存不足的处理

释放无用对象并回收Python内存：在关键阶段使用del 变量与import gc; gc.collect()，随后调用**torch.cuda.empty_cache()**清理未使用的GPU缓存（注意这只释放缓存，不会释放被张量占用的显存）。
适度清理Page Cache（仅在必要时、对性能影响可接受时执行）：
- 查看内存：运行free -h；
- 手动清理：执行sync && echo 1 > /proc/sys/vm/drop_caches（或2/3，谨慎使用，避免影响系统性能）。
增加Swap（当物理内存不足且训练/保存阶段需要更大内存缓冲时）：
- 临时创建并启用交换文件（示例为16GB）：
  - sudo fallocate -l 16G /swapfile
  - sudo chmod 600 /swapfile
  - sudo mkswap /swapfile
  - sudo swapon /swapfile
- 验证：运行free -h查看Swap是否增加；如需持久化，将其加入**/etc/fstab**。
若torch.save阶段仍被系统杀死，优先增加Swap或改为分批保存/流式保存，降低单次内存峰值。

四推理与超大模型场景的实用建议

优先采用AMP与分块推理（如**–tile**），必要时结合梯度检查点与CPU/RAM卸载。
使用更精简的优化器（如SGD）或低精度/混合精度以降低状态占用。
多卡/大模型采用FSDP进行张量分片；超大模型可结合ZeRO或DeepSpeed进行参数/梯度/优化器状态分片与卸载。
持续监控显存：在训练/推理循环中打印torch.cuda.memory_allocated() / memory_reserved()，配合nvidia-smi观察峰值与泄漏。

五最小可用代码示例 AMP + 梯度累积 + 清理缓存

示例展示了在不改变逻辑的前提下，组合使用AMP、梯度累积与缓存清理来降低显存占用：

import torch, gc
from torch.cuda.amp import autocast, GradScaler
from torch.utils.data import DataLoader

device = torch.device('cuda')
model = MyModel().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scaler = GradScaler()

accum_steps = 4  # 累积4个小批量的梯度
loader = DataLoader(dataset, batch_size=16, num_workers=4, pin_memory=True)

for i, (x, y) in enumerate(loader, 1):
    x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
    optimizer.zero_grad(set_to_none=True)  # 更省显存的梯度清零方式

    with autocast():
        out = model(x)
        loss = criterion(out, y) / accum_steps  # 注意缩放

    scaler.scale(loss).backward()

    if i % accum_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

    # 可选：释放本轮不再使用的临时张量与缓存
    del x, y, out, loss
    if i % 50 == 0:
        torch.cuda.empty_cache()
        gc.collect()

若仍报OOM，可进一步降低batch_size、开启梯度检查点，或在多卡环境下改用FSDP分片。

最新问答

相关标签