[YOLOv8] 问题分析和解决 - NotImplementedError: Could not run ‘torchvision::nms‘ ...from the ‘CUDA‘ backend

问题描述

当使用YOLOv8训练自定义的数据集时，

from ultralytics import YOLO

# 加载一个模型
model = YOLO('yolov8n.yaml')  # 从YAML建立一个新模型
# 训练模型
results = model.train(
    data='D:/YOLOv8Train/v8_train_datasets/mktk_dataset/data.yaml',
    device='0',
    epochs=5,
    batch=4,
    verbose=True,
    imgsz=640)

抛出如下错误：

 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).

运行环境说明

windows10 + python3.11 + cuda11.8

pip install ultralytics
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

原因分析和解决

查看当前python运行环境的torch版本： 

咱这么奇怪，torch + torchaudio都是coda版本，torchvision却不是，估计问题就出在这里。

为什么会出现这种情况呢？ 当你“pip install ultralytics” 时，torch和torchvision会作为它的依赖一起安装了，当执行“pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118”时，会卸载torch，然后再安装torch+cuda版本，但是不会自动卸载torchvision已有版本，然后再安装torchvision+cuda版本。

解决方法一：

手动卸载“torchvision” ，然后重新安装对应的cuda版本，

重新跑python代码，问题解决。

解决方法二：
或许，先安装pytorch cuda版本，然后再安装ultralytics：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
pip install ultralytics

有空的朋友可以试试看。

有网友，验证了，方法2不可行：