CUDA环境训练YOLOv8时全损失为NaN的问题求助

问题重现 #

• CUDA环境下训练YOLOv8时所有损失出现NaN，CPU训练完全正常；减小图像尺寸后训练恢复，但相同小尺寸下增大batch size会再次触发NaN问题。
• 训练代码：

from ultralytics import YOLO
import torch
import os
os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
model = YOLO("yolov8n.pt")  # 加载预训练模型（推荐用于训练）
results = model.train(data="data.yaml", epochs=15, workers=0, batch=12)  
results = model.val()  
model.export(format="onnx")

• 训练输出示例：

Epoch    GPU_mem   box_loss   cls_loss   dfl_loss  Instances       Size
1/15      1.74G        nan        nan        nan         51        640:   4%

解决方案 #

1. 显存适配优化 #

• 自动匹配最大batch size：将batch参数设为-1，让YOLOv8自动检测当前GPU显存支持的最大可行batch size，避免显存溢出引发数值异常：

  results = model.train(data="data.yaml", epochs=15, workers=0, batch=-1)
  

• 梯度累积模拟大batch：若需要固定较大的batch逻辑，使用accumulate参数，通过多次小batch的梯度累积实现等效大batch效果，降低显存压力。例如用batch=6、accumulate=2等效batch=12：

  results = model.train(data="data.yaml", epochs=15, workers=0, batch=6, accumulate=2)
  

• 适度降低图像分辨率：继续调小imgsz参数（如从640改为512），减少单张图像的显存占用，缓解显存紧张问题。

2. 数值稳定性调整 #

• 降低初始学习率：GPU训练时可能因计算精度问题触发梯度爆炸，尝试将初始学习率从默认的0.01调低，比如设为0.001：

  results = model.train(data="data.yaml", epochs=15, workers=0, batch=12, lr0=0.001)
  

• 显式启用混合精度训练：添加amp=True参数（默认已开启，显式指定可确保生效），利用半精度计算减少显存占用同时提升数值稳定性：

  results = model.train(data="data.yaml", epochs=15, workers=0, batch=12, amp=True)
  

3. CUDA环境验证 #

• 检查版本兼容性：确保PyTorch与CUDA版本匹配（如CUDA 11.8对应PyTorch 2.0+），可通过以下代码验证：

  print(torch.cuda.is_available())
  print(torch.version.cuda)
  

• 清理GPU缓存：训练前执行缓存清理，避免显存中残留的无效数据干扰：

  torch.cuda.empty_cache()
  

4. 训练配置与数据检查 #

• 调整数据加载workers：workers=0会让数据加载在主线程执行，可能与GPU计算冲突，尝试设置workers=2或根据CPU核心数调整，同时确保数据集路径无中文、特殊字符。
• 检查标签合理性：确认数据集标签无极端值（如坐标超出0-1范围），GPU训练对数值异常的敏感度更高，哪怕数据读取无报错也可能引发损失NaN。

内容的提问来源于stack exchange，提问作者Mohamed Abu ElNasr