CUDA可用但运行torchvision::nms报错，YOLOv8训练失败

问题背景 #

使用Ultralytics YOLOv8构建模型并训练时，尽管torch.cuda.is_available()返回True，仍触发NotImplementedError，提示torchvision::nms不支持CUDA后端。相关代码及错误信息如下：

用户代码：

from ultralytics import YOLO

# Load a model
model = YOLO("yolov8m.yaml")  # build a new model from scratch

# Use the model
results = model.predict(source="0", show=True, stream=True, classes=0, device='0')
model.train(data="conf.yaml", epochs=20, imgsz=600, device=' 0')  # train the model

错误信息：

Eror:NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). If you are a Facebook employee using PyTorch on mobile, please visit https://fburl.com/ptmfixes for possible resolutions. 'torchvision::nms' is only available for these backends: [CPU, QuantizedCPU, BackendSelect, Python, FuncTorchDynamicLayerBackMode, Functionalize, Named, Conjugate, Negative, ZeroTensor, ADInplaceOrView, AutogradOther, AutogradCPU, AutogradCUDA, AutogradXLA, AutogradMPS, AutogradXPU, AutogradHPU, AutogradLazy, AutogradMeta, Tracer, AutocastCPU, AutocastCUDA, FuncTorchBatched, FuncTorchVmapMode, Batched, VmapMode, FuncTorchGradWrapper, PythonTLSSnapshot, FuncTorchDynamicLayerFrontMode, PythonDispatcher].
File "C:\Users\pogti\Downloads\dataset_rust_ai\main.py", line 9, in
model.train(data="conf.yaml", epochs=20, imgsz=600, device=' 0') # train the model

解决步骤 #

• 修正设备参数格式
  训练代码中device=' 0'包含多余空格，导致YOLO无法正确识别CUDA设备，将其改为：

  model.train(data="conf.yaml", epochs=20, imgsz=600, device='0')
  

• 对齐torch与torchvision版本
  torchvision::nms算子对CUDA的支持依赖于torch与torchvision版本严格匹配：

  1. 查看当前torch版本：

     python -c "import torch; print(torch.__version__)"
     

  2. 安装与torch版本匹配的torchvision（示例为torch 2.0.1+cu118对应的版本）：

     pip install torchvision==0.15.2+cu118 --index-url https://download.pytorch.org/whl/cu118
     

     需替换版本号为实际torch对应的版本。

• 临时绕过算子兼容性问题
  如果版本对齐无效，可强制YOLO使用CPU版NMS后转回CUDA，在训练前添加代码：

  from ultralytics.utils.torch_utils import select_device
  import torchvision
  
  def cuda_safe_nms(boxes, scores, iou_threshold):
      return torchvision.ops.nms(boxes.cpu(), scores.cpu(), iou_threshold).to(select_device('0'))
  
  torchvision.ops.nms = cuda_safe_nms
  

• 验证CUDA环境完整性
  执行以下命令确认CUDA硬件识别正常：

  python -c "import torch; print(torch.cuda.device_count()); print(torch.cuda.get_device_name(0))"
  

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