MacBook M2用MPS训练YOLOv8出现CUDA报错求助

问题背景 #

使用MacBook Air M2训练YOLOv8自定义数据集，代码如下：

model = YOLO('yolov8s.pt')

# train
results = model.train(
    data=data,
    epochs=epochs,
    batch=batch_size,
    imgsz=img_size,
    project=project_name,
    name=model_name,
    device=torch.device('mps')
)

未指定device时依赖CPU训练，每个epoch耗时约1小时；指定device='mps'尝试GPU加速时，触发报错：

AssertionError: Torch not compiled with CUDA enabled

报错根源是YOLOv8的AutoBatch模块在计算最优batch size时，错误调用了CUDA相关API，而M2芯片不支持CUDA，导致逻辑崩溃。

解决办法 #

• 手动指定batch size，跳过AutoBatch自动计算
  直接设置具体的batch数值（根据M2内存调整，建议8-16），避免触发CUDA检测逻辑。修改后代码示例：

  model = YOLO('yolov8s.pt')
  
  # train
  results = model.train(
      data=data,
      epochs=epochs,
      batch=16,  # 手动设置适配M2的batch size
      imgsz=img_size,
      project=project_name,
      name=model_name,
      device='mps'  # 用字符串格式更简洁
  )
  

• 确认PyTorch支持MPS
  验证当前PyTorch是否开启MPS支持，执行命令：

  python -c "import torch; print(torch.backends.mps.is_available())"
  

  输出True说明MPS可用；若为False，需重新安装适配macOS的PyTorch 2.0+版本。

• 优化训练参数进一步提速

  • 开启数据集缓存：添加cache='ram'或cache='disk'参数，减少IO等待；
  • 缩小输入尺寸：降低imgsz（如从640改为480），减少计算量；
  • 调整工作线程数：将workers设为2或4，适配M2内存限制。

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