在CPU上训练YOLOv7出现CUDA错误的技术求助

问题情况 #

在无GPU的Colab环境中使用WongKinYiu的YOLOv7仓库训练模型，执行命令如下：

python train_aux.py --workers 1 --device cpu --batch-size 1 --data data/coco.yaml --img 128 128 --cfg /content/yolov7/cfg/training/yolov7-e6e.yaml --weights '' --name yolov7-e6e --hypdata/hyp.scratch.p6.yaml

执行后先收到警告：

warnings.warn('User provided device_type of \'cuda\', but CUDA is not available. Disabling')

进入第一个epoch时触发错误：

RuntimeError: No CUDA GPUs are available

已尝试安装cuda-python但无效果，需解决指定CPU后仍调用CUDA的问题。

解决步骤 #

• 修正命令参数格式
  原命令中--hypdata/hyp.scratch.p6.yaml缺少空格，正确写法应为--hyp data/hyp.scratch.p6.yaml。参数格式错误会导致--device cpu未被正确解析，程序默认尝试使用CUDA。

• 强制代码使用CPU
  若命令修正后仍有问题，需修改train_aux.py：

  1. 找到设备初始化代码，确保device被设置为torch.device('cpu')，强制忽略CUDA检测结果。
  2. 将数据加载器的pin_memory参数设为False（该参数仅适配CUDA，CPU环境下开启会引发异常）。
  3. 检查训练循环中的张量操作，移除所有cuda()或.to('cuda')的硬编码调用，统一使用.to(device)确保在CPU上执行。

• 排查配置文件
  打开yolov7-e6e.yaml配置文件，确认其中没有强制指定CUDA的配置项，如有则修改为自适应设备逻辑。

• 卸载无用依赖
  CPU环境下无需安装cuda-python，该包仅适用于CUDA环境，安装后可能引发环境冲突，建议执行!pip uninstall -y cuda-python移除。

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