linux pytorch怎么节省内存

在 Linux 上用 PyTorch 训练或推理时，节省内存（显存 + 内存）通常从模型、数据、计算方式、系统配置几个层面入手。下面按实用程度整理：

一、显存（GPU 内存）优化（最常见）

1. 使用混合精度（AMP）

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
with autocast():
    out = model(x)
    loss = crit(out, y)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()

✅ 显存可节省 30%~50%

2. 梯度累积（模拟大 batch）

accum_steps = 4
for i, (x, y) in enumerate(dataloader):
    loss = model(x, y) / accum_steps
    loss.backward()
    if (i+1) % accum_steps == 0:
        opt.step()
        opt.zero_grad()

✅ 小显存跑大 batch

3. 及时释放无用变量

del x, y, out
torch.cuda.empty_cache()

⚠️ empty_cache() 只释放未用显存，不解决根本问题

4. 使用 torch.no_grad()

推理或不需要梯度时：

with torch.no_grad():
    out = model(x)

5. 模型层面优化

• 使用更小模型（MobileNet / DistilBERT 等）
• 使用：

model.half()      # FP16
# 或
model.to(torch.float16)

6. 使用梯度检查点（Trade compute for memory）

from torch.utils.checkpoint import checkpoint

out = checkpoint(model, x)

✅ 显存显著降低，训练稍慢

7. 冻结部分参数

for p in model.backbone.parameters():
    p.requires_grad = False

二、CPU 内存优化（Linux 内存）

1. 使用 DataLoader 的 pin_memory=False

DataLoader(..., pin_memory=False)

2. 减少 num_workers

DataLoader(..., num_workers=2)

⚠️ worker 过多会复制多份数据

3. 使用内存映射 / 流式读取

• 大文件用 np.memmap
• 不要用 list 存全部 tensor

4. 及时 del 和 gc

import gc
del big_var
gc.collect()

三、Linux 系统层面

1. 限制进程内存（防 OOM）

ulimit -v 8000000   # 8GB

2. 使用 watch -n 1 nvidia-smi

观察显存峰值

3. 关闭多余进程

nvidia-smi
kill -9 PID

四、进阶方案（显存极小）

• ZeRO / DeepSpeed
• FSDP（PyTorch 原生）
• 模型offload到CPU（accelerate）

五、快速检查清单 ✅

• [ ] 用 AMP
• [ ] 用梯度累积
• [ ] 不需要梯度就 no_grad
• [ ] del + empty_cache
• [ ] num_workers 合理
• [ ] 模型是否可减小

如果你能说明：

• GPU 型号 / 显存大小
• 模型类型（CNN / Transformer）
• 训练还是推理

我可以给你针对性的最小改动方案。