YOLOv5训练时GPU显存不足怎么办？解决方案全汇总

在使用YOLOv5进行训练时，如果遇到GPU显存不足的问题，这通常是因为模型、输入数据的大小或者批量（batch size）设置不合理。以下是一些常见的解决方案汇总，帮助你优化显存使用并顺利训练模型。

1. 降低批量大小（Batch Size）

批量大小是影响显存占用的主要因素之一。较大的批量会显著增加显存需求。可以通过修改配置文件中的batch_size参数来减少显存占用。

• 修改方法：打开data/hyp.scratch.yaml或类似配置文件，找到batch_size参数，将其值调小（例如从64降到16或8）。
• 注意事项：较小的批量可能会导致训练不稳定，可以结合梯度累积技术（见下文）来缓解这一问题。

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

梯度累积是一种通过多次前向和后向传播计算梯度后再更新权重的技术，能够在不改变批量大小的情况下模拟大批次的效果。

• 修改方法：在训练命令中添加--accumulate参数，例如：

  python train.py --img 640 --batch 16 --epochs 50 --data dataset.yaml --weights yolov5s.pt --accumulate 4
  

  这里的--accumulate 4表示每4次迭代累积一次梯度再更新权重。

• 原理：相当于将实际批量大小从16提升到64（16 * 4），但显存消耗仍然保持较低水平。

3. 缩小输入图像尺寸

输入图像的分辨率越高，显存需求越大。可以通过降低输入图像的尺寸来减少显存占用。

• 修改方法：在训练命令中调整--img参数，例如从默认的640改为320或更低：

  python train.py --img 320 --batch 16 --epochs 50 --data dataset.yaml --weights yolov5s.pt
  

• 注意事项：较低的分辨率可能会影响模型精度，需根据具体任务权衡。

4. 选择更小的预训练模型

YOLOv5提供了多个版本的预训练模型（如yolov5s、yolov5m、yolov5l、yolov5x），其中yolov5s是最小的模型，显存占用也最低。

• 修改方法：在训练命令中指定较小的模型，例如：

  python train.py --img 640 --batch 16 --epochs 50 --data dataset.yaml --weights yolov5s.pt
  

5. 禁用Mosaic数据增强

Mosaic是一种数据增强技术，它将多张图片拼接成一张图片作为输入。虽然能提高训练效果，但会显著增加显存占用。

• 修改方法：打开data/hyp.scratch.yaml文件，将mosaic参数设置为0.0，例如：

  mosaic: 0.0
  

• 注意事项：禁用Mosaic可能会影响模型的泛化能力，需要根据任务需求权衡。

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

混合精度训练通过使用半精度浮点数（FP16）代替部分全精度浮点数（FP32）计算，可以显著减少显存占用。

• 修改方法：确保PyTorch版本支持torch.cuda.amp模块，并在训练命令中启用自动混合精度（AMP）。YOLOv5默认已经支持AMP，只需确保GPU硬件支持FP16计算即可。
• 注意事项：某些老旧GPU可能不支持FP16计算，需确认硬件兼容性。

7. 升级或更换GPU设备

如果以上方法仍无法解决问题，可能需要考虑升级或更换显卡。NVIDIA的Ampere架构（如RTX 30系列）对深度学习任务有较好的支持，显存容量更大且性能更强。

8. 分布式训练（Distributed Training）

如果有多块GPU可用，可以通过分布式训练分担显存压力。

• 修改方法：使用torch.distributed模块或多机多卡训练脚本。YOLOv5支持分布式训练，参考官方文档配置环境。
• 注意事项：分布式训练需要额外的网络通信开销，适合大规模数据集和复杂模型。

9. 其他优化技巧

• 冻结部分层：对于微调任务，可以冻结主干网络的部分层，仅训练最后几层。这可以显著减少显存占用。
• 减少数据增强操作：除了禁用Mosaic外，还可以减少其他耗显存的数据增强操作，如随机裁剪、翻转等。
• 调整Anchor数量：如果自定义数据集类别较少，可以减少Anchor的数量以降低显存需求。

流程图：解决显存不足的步骤

flowchart TD
    A[显存不足] --> B{是否可降低
Batch Size?}
    B -- 是 --> C[降低Batch Size]
    B -- 否 --> D{是否可启用
梯度累积?}
    D -- 是 --> E[启用梯度累积]
    D -- 否 --> F{是否可缩小
输入尺寸?}
    F -- 是 --> G[调整输入尺寸]
    F -- 否 --> H{是否可选择
更小模型?}
    H -- 是 --> I[切换到小模型]
    H -- 否 --> J[尝试其他优化方法]