如何降低显存占用

降低显存占用是深度学习训练和推理中的常见需求，尤其在资源有限的情况下。以下是一些通用且实用的策略，涵盖模型设计、训练过程、推理优化等多个层面：

一、模型结构层面

1. 使用更小的模型

• 减少网络层数、通道数
• 使用轻量模型（如 MobileNet、EfficientNet、TinyBERT 等）
• 用 Depthwise Separable Convolution 代替普通卷积

2. 降低精度（Quantization）

• FP16 / BF16（混合精度训练）

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
with autocast():
output = model(input)

• INT8 量化（推理阶段）
• PyTorch: torch.quantization
• TensorRT / ONNX Runtime

3. 参数共享与稀疏化

• 共享 embedding / attention 参数
• 使用 LoRA / QLoRA（大模型微调时非常有效）

二、训练过程优化

4. 减小 Batch Size

• 最直接、最有效
• 配合 Gradient Accumulation

loss = loss / accumulation_steps
loss.backward()
if step % accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()

5. 使用梯度检查点（Gradient Checkpointing）

• 用计算换显存

model.gradient_checkpointing_enable()

• 常见于 Transformer（BERT、GPT、LLaMA）

6. 冻结部分参数

• 微调时只训练部分层

for param in model.base_model.parameters():
param.requires_grad = False

三、显存管理技巧

7. 及时释放无用变量

del tensor
torch.cuda.empty_cache()

8. 避免不必要的 .cuda() 拷贝

• 不要频繁在 CPU / GPU 之间搬运数据
• 尽量使用 pin_memory=True + DataLoader

9. 限制 PyTorch 显存增长（CUDA）

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

四、推理阶段优化

10. 推理时不保存梯度

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

11. 使用推理专用优化

• TensorRT
• ONNX Runtime
• TorchScript

model.eval()
traced = torch.jit.trace(model, example_input)

五、分布式与并行

12. 模型并行 / 流水线并行

• 超大模型拆到多张卡
• DeepSpeed / Megatron-LM

13. ZeRO 优化（DeepSpeed）

• ZeRO-1 / ZeRO-2 / ZeRO-3
• 显著降低显存占用

六、针对大语言模型（LLM）

快速诊断显存问题

torch.cuda.memory_summary()
torch.cuda.max_memory_allocated()

推荐组合（实用）

✅ 训练：
FP16 + Gradient Checkpointing + Gradient Accumulation
✅ 微调大模型：
QLoRA + DeepSpeed ZeRO-2
✅ 推理：

FP16 + TensorRT / ONNX + no_grad

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

• 使用的框架（PyTorch / TensorFlow）
• 模型类型（CNN / Transformer / LLM）
• 显卡型号 & 显存大小
• 训练还是推理

我可以给你更具体、可落地的方案。