ubuntu pytorch如何节省显存

在 Ubuntu + PyTorch 场景下，显存（GPU memory）不足是非常常见的问题。下面按由易到难、由训练到推理的方式，系统总结节省显存的方法，并给出可直接用的示例。

一、基础但非常有效的做法

1. 使用更小的 batch size

最直接：

train_loader = DataLoader(dataset, batch_size=8)  # 从 32 → 8

2. 使用 torch.cuda.empty_cache()

释放未被引用的缓存显存（不是魔法，但有用）：

import torch
torch.cuda.empty_cache()

3. 避免在循环中累积计算图

确保不需要反向传播的地方加：

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

二、模型与计算层面的优化（重点）

4. 使用混合精度训练（AMP）✅ 强烈推荐

可节省 30%~50% 显存，还能加速：

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler()
for x, y in loader:
    optimizer.zero_grad()
    with autocast():
        loss = model(x, y)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

5. 使用 torch.utils.checkpoint（梯度检查点）

用时间换显存，适合 Transformer / 大模型：

from torch.utils.checkpoint import checkpoint

def forward(x):
    return checkpoint(model.layer, x)

或：

model.layers = torch.nn.Sequential(*[
    checkpoint(layer) for layer in model.layers
])

6. 使用更小的模型 / 模型并行

• 换 bert-base → bert-small
• 使用 distilbert, tinyllama, mobilenet
• 多卡时：

model = torch.nn.DataParallel(model)
# 或
model = model.to('cuda:0')

三、推理阶段专用（节省显存效果明显）

7. 推理时关闭梯度

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

8. 使用 torch.inference_mode()（比 no_grad 更省）

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

9. 模型量化（INT8 / FP16）

model.half()  # FP16
# 或
model = torch.quantization.quantize_dynamic(
    model, {nn.Linear}, dtype=torch.qint8
)

四、Ubuntu 系统级技巧

10. 查看显存占用

nvidia-smi

或实时监控：

watch -n 1 nvidia-smi

11. 清理僵尸进程

nvidia-smi | grep python
kill -9 <PID>

12. 限制可见 GPU

export CUDA_VISIBLE_DEVICES=0

五、进阶方案（大模型必备）

示例（FSDP）：

from torch.distributed.fsdp import FSDP
model = FSDP(model)

六、推荐组合（实用方案）

训练：

• 小 batch
• AMP
• gradient checkpoint
• empty_cache

推理：

• inference_mode
• FP16
• 量化

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

• GPU 型号（如 3090 / 4090 / A100）
• 模型类型（CNN / Transformer / LLM）
• 训练还是推理

我可以给你针对性的显存优化方案。