CUDA Out of Memory原因总结和方法
原因总结
-
显存不足：
- 深度学习模型（尤其是大型模型）在训练或推理时需要大量的显存。如果显存容量不足，会导致CUDA Out of Memory错误。
-
批处理大小过大：
- 在训练过程中，批处理（batch）大小设置过大时，会占用过多的显存，导致显存溢出。
-
模型过大：
- 模型参数过多，导致模型在显存中的占用空间过大。
-
显存碎片化：
- 在频繁的显存分配和释放过程中，可能会导致显存碎片化，尽管总的空闲显存足够，但由于碎片化导致无法分配连续的大块显存空间。
-
未及时释放显存：
- 有些情况下，未及时释放显存或显存未被正确回收，也会导致显存溢出。
-
多模型或多任务并行运行：
- 同时运行多个模型或任务时，共享的显存资源可能不足。
解决方法
- 减少批处理大小：
- 将批处理大小减小，以降低每次训练迭代时的显存需求。
batch_size = 32 # 示例：减小批处理大小- 模型精简：
- 使用较小的模型或对模型进行裁剪，减少模型参数数量和显存占用。
- 使用模型蒸馏等技术，得到一个较小的模型。
- 使用混合精度训练（Mixed Precision Training），以减少显存占用。
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for data, target in train_loader:
optimizer.zero_grad()
with autocast():
output = model(data)
loss = criterion(output, target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
- 优化代码：
确保及时释放不再使用的显存。例如，在PyTorch中使用torch.cuda.empty_cache()来释放显存。
import torch
torch.cuda.empty_cache()
- 分布式训练：
使用数据并行或模型并行技术，将模型训练分布到多个GPU上，以减小单个GPU的显存压力。
model = torch.nn.DataParallel(model)