M1芯片Mac运行YOLOv7遇CUDA相关AssertionError的解决方法

问题背景 #

在M1芯片的MacBook上使用自定义数据集训练YOLOv7，已知CUDA不支持macOS，已重装适配ARM架构的anaconda3和miniforge3，但训练30个epoch后抛出AssertionError: Torch not compiled with CUDA enabled错误。

运行命令 #

python train.py --workers 1 --batch-size 1 --epochs 200 --img 640 640 --data data/custom_data.yaml --hyp data/hyp.scratch.custom.yaml --cfg cfg/training/custom_yolov7.yaml --name yolov7-output --weights yolov7.pt

核心报错信息 #

matching_bs[i] = torch.tensor([], device='cuda:0', dtype=torch.int64)
  File "/Users/yeti/opt/anaconda3/envs/yolo/lib/python3.10/site-packages/torch/cuda/__init__.py", line 221, in _lazy_init
    raise AssertionError("Torch not compiled with CUDA enabled")
AssertionError: Torch not compiled with CUDA enabled

解决建议 #

• 显式指定训练设备：在训练命令中添加--device cpu强制使用CPU，避免代码自动尝试调用CUDA。修改后的命令：

  python train.py --workers 1 --batch-size 1 --epochs 200 --img 640 640 --data data/custom_data.yaml --hyp data/hyp.scratch.custom.yaml --cfg cfg/training/custom_yolov7.yaml --name yolov7-output --weights yolov7.pt --device cpu
  

• 修复代码中的CUDA硬编码：报错原因是代码中直接写死了device='cuda:0'，找到OTA损失计算相关的代码位置，将设备指定改为动态判断或强制CPU。例如：
  将device='cuda:0'替换为device=torch.device('cpu')，或者复用命令行传入的设备参数args.device。
• 安装适配M1的PyTorch版本：卸载当前PyTorch，重新安装支持Apple Silicon的版本，命令如下：

  pip3 install torch torchvision torchaudio
  

  安装后验证MPS可用性：执行python -c "import torch; print(torch.backends.mps.is_available())"，若返回True，可尝试用MPS加速，训练命令改为指定--device mps，同时修改代码中硬编码CUDA的部分为：

  device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu')
  

• 禁用OTA损失计算：在超参数文件hyp.scratch.custom.yaml中设置loss_ota: 0，关闭OTA损失逻辑，避免触发CUDA相关代码。

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