自定义数据集训练YOLOv5、YOLOv6时无法启用GPU加速问题咨询

按优先级从高到低排查以下问题，基本能覆盖绝大多数同类异常：

• 首先核对训练启动参数的设备指定
  不管是YOLOv5还是YOLOv6，启动训练时如果没有显式传入--device参数，部分版本的代码会默认 fallback 到CPU运行，和torch本身能不能识别CUDA没有直接关系。单卡训练直接在启动命令后加--device 0，多卡训练写--device 0,1即可，参考正确启动命令：
  python train.py --data custom_dataset.yaml --weights yolov5s.pt --img 640 --device 0
• 检查训练配置文件的硬编码设置
  打开你使用的训练配置yaml文件，找到device字段，很多人调整数据集、超参数的时候会不小心把这个字段改成cpu，或者留空触发默认CPU逻辑，直接把字段值修改为对应GPU编号0即可。
• 验证torch的CUDA调用是否真的可用
  部分环境会出现torch.cuda.is_available()返回True，但实际张量无法加载到GPU的假识别问题，跑以下测试代码验证：

  import torch
  # 测试张量能否正常加载到GPU做计算
  test_tensor = torch.randn(20000, 20000).cuda()
  res = torch.matmul(test_tensor, test_tensor)
  print(f"计算设备: {res.device}, 计算结果维度: {res.shape}")
  

  运行代码时打开任务管理器/nvidia-smi看GPU显存占用，如果显存明显上涨，说明torch本身CUDA功能正常，问题出在YOLO代码的设备判断逻辑；如果这步就报错，重装对应CUDA版本的torch即可，torch1.12匹配CUDA11.3/11.6版本。
• 修复代码分支的设备判断bug
  如果最近拉取过YOLOv5/YOLOv6的最新仓库代码，部分老版本分支对torch1.12的兼容性有问题：比如YOLOv6 3.0早期版本的设备判断逻辑，会把空值的device参数错误解析为CPU，哪怕CUDA可用也不会切换。直接打开仓库下utils/torch_utils.py找到select_device函数，加一行打印看最终返回的设备类型，如果返回cpu顺着逻辑往上定位覆盖设备参数的代码段修正即可。
• 排查环境变量干扰
  检查运行训练脚本的终端，有没有设置过CUDA_VISIBLE_DEVICES=""，这个环境变量为空时，torch虽然能检测到CUDA组件存在，但枚举不到可用物理卡，会自动切CPU运行。启动训练前先执行对应命令指定可用显卡：
  • Windows cmd：set CUDA_VISIBLE_DEVICES=0
  • Linux/Windows PowerShell：export CUDA_VISIBLE_DEVICES=0

相关参考截图：

内容的提问来源于stack exchange，提问作者Kkweil