yolov11m进行自己数据集的训练时报错CUDA超出存储

Question

我在使用yolov11的官方文档中的模型yolo11m进行自己数据集的训练时，总是报出下面的错误：
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 200.00 MiB. GPU 0 has a total capacity of 11.76 GiB of which 64.44 MiB is free. Process 3506601 has 11.69 GiB memory in use. Of the allocated memory 11.15 GiB is allocated by PyTorch, and 346.56 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
我不知道是不是因为我的数据集过大，有一万多张，而且在运行时我的GPU上只有这一个模型在跑，12G的显卡不知道怎么就被占用了11.69 G的存储，而且在监测GPU存储使用情况时只有几百M被使用。另外，使用yolo11n模型不会报错，使用yolo11s也会报错。如果使用yolo11m训练只有4张图片的coco数据集时也可以正常使用。
我在网络上找了相关问题的解决方案，但是都没有用。希望能得到大 佬的指导。

Answer 1

该回答引用自GPT-deepseek, 由博主 GISer Liu 编写：

问题分析

用户在训练YOLOv11m模型时遇到了CUDA内存不足的错误。具体表现为：

1. 错误信息：torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 200.00 MiB. GPU 0 has a total capacity of 11.76 GiB of which 64.44 MiB is free.
2. GPU使用情况：GPU总容量为11.76 GiB，当前已使用11.69 GiB，剩余64.44 MiB。
3. 模型表现：
   • 使用YOLOv11n模型时不会报错。
   • 使用YOLOv11s模型时会报错。
   • 使用YOLOv11m模型训练只有4张图片的COCO数据集时正常。

问题定位

1. 模型复杂度：YOLOv11m模型比YOLOv11n模型更复杂，需要更多的显存。YOLOv11s模型介于两者之间，但依然可能导致显存不足。
2. 数据集大小：用户的数据集有一万多张图片，训练时显存需求较大。
3. 显存管理：PyTorch在显存管理上可能存在碎片化问题，导致显存无法充分利用。

解决方案

1. 减少批量大小（Batch Size）

批量大小是影响显存使用的主要因素之一。减少批量大小可以显著降低显存需求。

# 修改训练脚本中的批量大小
batch_size = 8  # 原批量大小
new_batch_size = 4  # 新批量大小

2. 使用梯度累积（Gradient Accumulation）

如果减少批量大小影响训练效果，可以使用梯度累积来模拟更大的批量大小。

# 修改训练脚本中的梯度累积步数
accumulation_steps = 2  # 梯度累积步数
optimizer.zero_grad()
for i, (images, targets) in enumerate(train_loader):
    outputs = model(images)
    loss = criterion(outputs, targets)
    loss = loss / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

3. 使用混合精度训练（Mixed Precision Training）

混合精度训练可以减少显存使用并加速训练。

from torch.cuda.amp import GradScaler, autocast

scaler = GradScaler()
for images, targets in train_loader:
    optimizer.zero_grad()
    with autocast():
        outputs = model(images)
        loss = criterion(outputs, targets)
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

4. 设置PyTorch显存管理环境变量

根据错误信息提示，可以设置PYTORCH_CUDA_ALLOC_CONF环境变量来优化显存管理。

export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

5. 使用更小的模型或剪枝

如果以上方法仍无法解决问题，可以考虑使用更小的模型（如YOLOv11n）或对模型进行剪枝。

总结

通过减少批量大小、使用梯度累积、混合精度训练、优化显存管理等方法，可以有效解决CUDA内存不足的问题。如果问题依然存在，建议进一步检查数据集和模型结构，或考虑使用更大显存的GPU。

参考文档

• PyTorch显存管理文档
• 混合精度训练文档

希望以上解决方案能帮助您顺利训练YOLOv11m模型！

如果该回答解决了您的问题，请采纳！如果没有，请私信联系或评论您的疑惑