linux pytorch为何内存不足

在 Linux + PyTorch 场景下，“内存不足”可能指两种不同层面：

1. 物理内存（RAM）不足 → OutOfMemoryError: unable to allocate ...
2. GPU 显存（CUDA memory）不足 → CUDA out of memory

下面分别说明常见原因和解决办法。

一、GPU 显存不足（最常见）

常见原因

1. batch size 太大
2. 模型太大（如 LLM、ViT、大模型）
3. 中间变量未释放
   • 没有 loss.backward() 后 optimizer.zero_grad()
   • 没有 del x / torch.cuda.empty_cache()
4. 保留了计算图
   • 用了 loss.item() 以外的 Python 操作
   • 没有 with torch.no_grad():
5. DataLoader 占用显存
   • pin_memory=True + 多进程
6. 多个进程共用一块 GPU

解决办法

# 减小 batch size
batch_size = 8  # ↓

# 使用混合精度
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
    output = model(x)
    loss = criterion(output, y)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

# 及时清理
del output, loss
torch.cuda.empty_cache()

# 推理时不记录梯度
with torch.no_grad():
    out = model(x)

其他：

• 使用 torch.utils.checkpoint
• 用 model.half()（FP16）
• 多卡训练（DataParallel / DDP）

二、CPU 内存（RAM）不足

常见原因

1. DataLoader num_workers 太多
2. 数据集一次性读入内存
3. 张量在 CPU 上太大
4. 内存泄漏
   • 全局 list 不断 append
5. 交换分区（swap）太小

解决办法

# 减少 worker
DataLoader(..., num_workers=4)

# 使用内存映射
np.load(..., mmap_mode='r')

# 及时释放
del data
gc.collect()

三、Linux 系统层面原因

1. ulimit 限制

ulimit -a

若 max memory size 很小，可修改。

2. 容器内存限制（Docker）

docker run --memory=16g ...

3. OOM Killer

dmesg | grep -i oom

四、快速排查命令

# GPU 显存
nvidia-smi

# CPU 内存
free -h
top
htop

# 进程内存
ps aux --sort=-rss

五、总结一句话

PyTorch 内存不足，90% 是 GPU 显存问题，核心原因是 batch 大 + 计算图未释放。

如果你愿意，可以贴出：

• 报错信息
• 模型规模
• batch size
• 是否用 GPU

我可以帮你精确定位并给出修改方案。