运行YOLOv7的detect.py时遇torchvision::nms CUDA后端报错求助

问题描述 #

运行YOLOv7检测命令：

python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg

出现核心报错：

NotImplementedError: Could not run 'torchvision::nms' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend...

解决方案 #

1. 确保torch与torchvision版本严格匹配 #

你的环境使用PyTorch 2.1.1+cu118，需安装对应CUDA版本的torchvision，执行命令：

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

2. 强制CPU运行验证问题根源 #

切换到CPU模式运行，确认是否为CUDA相关兼容问题：

python detect.py --weights yolov7.pt --conf 0.25 --img-size 640 --source inference/images/horses.jpg --device cpu

若CPU运行正常，说明问题出在CUDA与torchvision的版本兼容上。

3. 修改YOLOv7的NMS实现代码 #

若版本匹配后仍报错，直接修改utils/general.py中的NMS调用逻辑：
找到文件中调用torchvision.ops.nms的代码行（约第713行），替换为：

# 改用PyTorch原生NMS，适配CUDA后端
i = torch.nms(boxes, scores, iou_thres)

torch.nms是PyTorch原生实现，对CUDA后端支持更稳定，参数顺序与原调用一致。

4. 重装PyTorch与CUDA对应版本 #

如果以上方法无效，完全卸载现有PyTorch组件后重装：

pip uninstall torch torchvision torchaudio -y
pip install torch==2.1.1+cu118 torchvision==0.16.1+cu118 torchaudio==2.1.1+cu118 --index-url https://download.pytorch.org/whl/cu118

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