Ubuntu上PyTorch的内存管理技巧

以下是Ubuntu上PyTorch的内存管理技巧：

1. 降低批次大小：减少单次输入样本量，直接降低显存占用，但可能影响训练速度。
2. 混合精度训练：使用torch.cuda.amp进行FP16/FP32混合精度计算，减少显存占用并加速训练。
3. 梯度累积：通过小批次累积梯度模拟大批量效果，避免显存爆炸。
4. 清理缓存与释放内存：
   • 使用torch.cuda.empty_cache()释放未使用的缓存。
   • 手动删除无用变量并调用gc.collect()触发垃圾回收。
5. 优化模型结构：
   • 采用轻量化模型（如MobileNet、DistilBERT）或深度可分离卷积。
   • 使用梯度检查点（gradient_checkpointing）减少中间激活值存储。
6. 数据加载优化：
   • 减少num_workers或设置pin_memory=False避免内存锁竞争。
   • 对大规模数据集启用流式加载（streaming=True）。
7. 分布式训练：通过数据并行或多卡训练分摊显存压力。
8. 监控与分析工具：
   • 使用torch.cuda.memory_summary()查看显存分配情况。
   • 通过nvidia-smi监控GPU内存使用，定位异常进程。

参考来源：