linux pytorch如何省内存

在 Linux 下使用 PyTorch 训练或推理时，**省内存（显存 / RAM）**通常从以下几个方面入手。下面按「实用程度 + 常见场景」给你一个清单。

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

1. 使用 torch.cuda.empty_cache()

import torch
torch.cuda.empty_cache()

⚠️ 只能释放未被引用的缓存，不能凭空省钱，但能避免碎片。

2. 减小 Batch Size

最直接有效：

batch_size = 8  # 从 32 降到 8

如果 batch 太小影响效果，可用 梯度累积。

3. 梯度累积（Gradient Accumulation）

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

✅ 小 batch 模拟大 batch

4. 使用 torch.no_grad()（推理 / 验证）

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

避免保存计算图，省大量显存。

5. 混合精度训练（AMP）

scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
    loss = model(x).loss
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

✅ 显存 ↓ 30%~50%

6. 使用更小的模型 / 参数 dtype

model.half()        # FP16
# 或
model = model.to(torch.float32)  # 避免默认 float64

7. 及时释放变量

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

8. 使用 model.eval()（验证时）

model.eval()

关闭 dropout / batch norm 更新。

二、CPU 内存（RAM）优化（Linux 服务器常见）

1. 使用 DataLoader 的 num_workers

DataLoader(..., num_workers=4, pin_memory=True)

⚠️ 太多反而吃内存

2. 避免把数据全加载进内存

❌ 错误：

data = [load(i) for i in range(100000)]

✅ 正确：

• 用 Dataset
• 用 yield
• 用内存映射（np.memmap, lmdb, webdataset）

3. 使用 del + gc.collect()

import gc
del large_var
gc.collect()

三、进阶（大模型 / 训练不起）

1. 梯度检查点（Gradient Checkpointing）

model.gradient_checkpointing_enable()

✅ 用时间换显存

2. ZeRO / DeepSpeed / FSDP

• DeepSpeed ZeRO-2 / ZeRO-3
• PyTorch FSDP

适合：

• 大模型
• 多卡
• 显存严重不足

3. 模型分片 / CPU offload

# HuggingFace
model = AutoModel.from_pretrained(..., device_map="auto")

或：

model.cpu()

四、Linux 系统层辅助

查看显存

nvidia-smi

查看内存

free -h
htop

限制进程内存（不推荐训练）

ulimit -v 20000000

五、推荐组合（最实用）

✅ 训练时：

• 小 batch + 梯度累积
• AMP
• empty_cache
• gradient checkpointing

✅ 推理时：

• torch.no_grad()
• model.eval()
• FP16

如果你愿意，可以告诉我：

• GPU 型号 / 显存大小
• 模型类型（CNN / Transformer / LLM）
• 报错信息（OOM 日志）

我可以直接给你针对性的省内存方案。