NotImplementedError：尽管设置和设备发现成功，但仍无法运行“torchvision::nms”

Question

我正在尝试使用 pyCharm 在 Windows 10 上使用 GPU 在我的笔记本电脑上训练模型，因此我遵循了一些在线指南，了解如何正确设置我的 NVidia 驱动程序、cuda、cuDNN 以及要安装的库版本。显然，设置过程很成功，我能够发现 GPU 并获取有关它的信息。但是，如果我尝试训练模型，我会收到错误：

NotImplementedError: Could not run 'torchvision::nms with arguments from the 'CUDA' backend.'

我已按照其网站

上的指示设置了 torch、torchvision 和 torchaudio

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

我已经安装了最新的 NVidia 驱动程序、CUDA 版本 12.1 和 cuDNN，并确保环境变量中的路径正确。我重新启动电脑，还是没有。

如果我尝试检查我的 GPU 信息，它会按预期工作。

import torch

print(f'CUDA version: {torch.version.cuda}')
cuda_id = torch.cuda.current_device()
print(f"ID of current CUDA device:{torch.cuda.current_device()}")
print(f"Name of current CUDA device:{torch.cuda.get_device_name(cuda_id)}")
print(f'Torch version: {torch.__version__}')

输出：

CUDA version: 12.1
ID of current CUDA device:0
Name of current CUDA device:NVIDIA GeForce GTX 1060 with Max-Q Design
Torch version: 2.1.1+cu121

所以 Torch 版本是正确的并且与我的 CUDA 版本相对应。

我尝试训练的代码：

from ultralytics import YOLO
from datetime import datetime, timedelta

model = YOLO('yolov8n.pt')

startTime = datetime.now()
results = model.train(data='data.yaml', epochs=15,
                      imgsz=[480, 352])
endTime = datetime.now()
delta = endTime - startTime
print('\n\n\n')
print(f'Model training took {delta}')
print(f'Model training took {delta.seconds}')

这是否仍然是安装不当的问题？可能是pyCharm引起的？我很茫然。

Answer 1

看来这个问题一定是由

pip

引起的。我放弃了尝试让它工作并屈服并安装了 Anaconda。

我创建了一个新环境，按照指示设置 pyTorch，安装了所有其他依赖项，它按预期工作。