【Bug已解决】Getting NotImplementedError: Could not run torchvision::nms with arguments from CUDA backend despite having all necessary libraries and imports 解决方案问题描述在使用 PyTorch 和 torchvision 进行目标检测或非极大值抑制NMS相关任务时开发者可能会遇到一个令人困惑的错误NotImplementedError: Could not run torchvision::nms with arguments from the CUDA backend.这个错误通常出现在以下场景中使用 torchvision 的nms函数处理 GPU 上的边界框使用 Faster R-CNN、SSD、YOLO 等目标检测模型进行推理使用 torchvision.ops 中的batched_nms、roi_align、roi_pool等操作在将模型和数据移到 GPU 后调用包含 NMS 的后处理操作令人困惑的是开发者可能已经确认PyTorch 已正确安装且 CUDA 可用torch.cuda.is_available()返回Truetorchvision 已正确安装所有输入张量都在 GPU 上所有必要的库和导入都已完成但仍然遇到这个NotImplementedError。错误复现场景一直接使用 torchvision.ops.nmsimport torch import torchvision.ops as ops # 创建 GPU 上的边界框和分数 boxes torch.tensor([ [10, 10, 50, 50], [12, 12, 52, 52], [100, 100, 150, 150] ], devicecuda) scores torch.tensor([0.9, 0.8, 0.7], devicecuda) # 执行 NMS result ops.nms(boxes, scores, iou_threshold0.5)报错NotImplementedError: Could not run torchvision::nms with arguments from the CUDA backend.场景二使用目标检测模型import torch import torchvision from torchvision.models.detection import fasterrcnn_resnet50_fpn # 加载预训练模型并移到 GPU model fasterrcnn_resnet50_fpn(pretrainedTrue) model model.to(cuda) model.eval() # 创建输入并移到 GPU images [torch.randn(3, 300, 400, devicecuda)] # 推理 with torch.no_grad(): predictions model(images)报错NotImplementedError: Could not run torchvision::nms with arguments from the CUDA backend.场景三使用 batched_nmsimport torch import torchvision.ops as ops boxes torch.randn(100, 4, devicecuda) scores torch.randn(100, devicecuda) idxs torch.randint(0, 10, (100,), devicecuda) result ops.batched_nms(boxes, scores, idxs, 0.5)报错NotImplementedError: Could not run torchvision::nms with arguments from the CUDA backend.根因分析1. torchvision 与 PyTorch 版本不匹配这是最常见的原因。torchvision 的 C 扩展包括 NMS、ROI Align 等操作是针对特定版本的 PyTorch 编译的。如果 torchvision 和 PyTorch 的版本不匹配CUDA 后端的 C 扩展可能无法正确加载或运行。版本对应关系示例PyTorch 2.1.0 → torchvision 0.16.0PyTorch 2.0.1 → torchvision 0.15.2PyTorch 1.13.1 → torchvision 0.14.1如果安装了 PyTorch 2.1.0 但 torchvision 是 0.15.2为 PyTorch 2.0 编译的NMS 的 CUDA 实现可能不兼容。2. torchvision 安装了 CPU-only 版本与 PyTorch 类似torchvision 也有 CPU 和 CUDA 版本。如果通过默认的pip install torchvision安装可能安装的是 CPU 版本。CPU 版本的 torchvision 不包含 CUDA 后端的 C 扩展。3. torchvision 的 C 扩展未正确编译从源码编译安装 torchvision 时如果编译环境中缺少 CUDA Toolkit 或编译参数不正确可能只编译了 CPU 后端没有编译 CUDA 后端。4. CUDA 架构不匹配torchvision 的 CUDA 扩展是针对特定的 GPU 架构compute capability编译的。如果你的 GPU 架构不在编译目标列表中CUDA 内核无法运行。例如torchvision 可能针对 compute capability 3.7 编译而你的 GPU 是 3.5。5. 混合安装方式导致冲突通过 conda 安装 PyTorch 但通过 pip 安装 torchvision或反之可能导致 ABI 不兼容。6. torchvision 版本过旧旧版本的 torchvision 可能存在 CUDA NMS 实现的 bug在某些 PyTorch 版本上无法正常工作。解决方案方案一重新安装匹配版本的 PyTorch 和 torchvision首先卸载当前版本pip uninstall torch torchvision torchaudio -y然后安装匹配的 CUDA 版本# CUDA 12.1 版本 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 # CUDA 11.8 版本 pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118使用 conda 安装conda install pytorch torchvision torchaudio pytorch-cuda12.1 -c pytorch -c nvidia方案二验证版本兼容性import torch import torchvision print(fPyTorch version: {torch.__version__}) print(ftorchvision version: {torchvision.__version__}) print(fCUDA available: {torch.cuda.is_available()}) print(fCUDA version: {torch.version.cuda}) # 检查 torchvision 的 CUDA 扩展是否可用 # 尝试在 GPU 上运行 NMS try: boxes torch.tensor([[0.0, 0.0, 10.0, 10.0]], devicecuda) scores torch.tensor([0.9], devicecuda) result torchvision.ops.nms(boxes, scores, 0.5) print(fNMS CUDA test passed: {result}) except NotImplementedError: print(NMS CUDA not available! Version mismatch likely.)方案三从源码编译 torchvision带 CUDA 支持如果预编译包不适用于你的环境可以从源码编译# 安装编译依赖 sudo apt install -y ninja-build # Linux # 或 brew install ninja # macOS # 克隆 torchvision git clone https://github.com/pytorch/vision.git cd vision # 检出与 PyTorch 匹配的版本 git checkout v0.16.0 # 与 PyTorch 2.1.0 匹配 # 设置 CUDA 架构根据你的 GPU 调整 export TORCH_CUDA_ARCH_LIST7.0;7.5;8.0;8.6;8.9PTX # 编译安装 python setup.py install方案四临时回退到 CPU 执行 NMS如果无法立即修复 CUDA NMS可以临时将 NMS 操作移到 CPU 执行import torch import torchvision.ops as ops def nms_cuda_fallback(boxes, scores, iou_threshold): NMS CUDA 回退方案将数据移到 CPU 执行 NMS再移回 GPU original_device boxes.device boxes_cpu boxes.cpu() scores_cpu scores.cpu() result ops.nms(boxes_cpu, scores_cpu, iou_threshold) return result.to(original_device) # 使用回退方案 boxes torch.tensor([[10, 10, 50, 50], [12, 12, 52, 52]], devicecuda) scores torch.tensor([0.9, 0.8], devicecuda) result nms_cuda_fallback(boxes, scores, 0.5) print(result) # 在 GPU 上完整修复代码以下是一个完整的诊断和修复脚本 torchvision NMS CUDA 问题诊断与修复脚本 import sys import torch import torchvision def diagnose(): 诊断 torchvision NMS CUDA 问题 print( * 60) print(torchvision NMS CUDA 诊断) print( * 60) # 1. 检查 PyTorch 和 torchvision 版本 print(f\n[1] 版本信息) print(f PyTorch: {torch.__version__}) print(f torchvision: {torchvision.__version__}) print(f CUDA available: {torch.cuda.is_available()}) print(f CUDA version: {torch.version.cuda}) # 2. 检查版本兼容性 print(f\n[2] 版本兼容性检查) compatible_pairs { 2.1.0: 0.16.0, 2.1.1: 0.16.1, 2.1.2: 0.16.2, 2.0.0: 0.15.0, 2.0.1: 0.15.2, 1.13.0: 0.14.0, 1.13.1: 0.14.1, 1.12.0: 0.13.0, 1.12.1: 0.13.1, } torch_ver torch.__version__.split()[0] # 去掉 cu121 等后缀 tv_ver torchvision.__version__.split()[0].split(a)[0] # 去掉后缀 expected_tv compatible_pairs.get(torch_ver) if expected_tv: if tv_ver expected_tv: print(f [OK] torchvision {tv_ver} 与 PyTorch {torch_ver} 匹配) else: print(f [MISMATCH] PyTorch {torch_ver} 需要 torchvision {expected_tv}) print(f 但当前安装的是 torchvision {tv_ver}) print(f 修复: pip install torchvision{expected_tv}) else: print(f [UNKNOWN] 无法验证 PyTorch {torch_ver} 的兼容性) print(f 请查看 https://github.com/pytorch/vision#installation) # 3. 检查是否为 CPU 版本 print(f\n[3] 安装类型检查) if cpu in torch.__version__: print(f [WARNING] PyTorch 是 CPU 版本 ({torch.__version__})) print(f 需要安装 CUDA 版本) elif cu in torch.__version__: print(f [OK] PyTorch 是 CUDA 版本 ({torch.__version__})) else: print(f [INFO] PyTorch 版本: {torch.__version__} (可能是 conda 安装)) # 4. 测试 NMS CPU print(f\n[4] NMS CPU 测试) try: boxes torch.tensor([[0.0, 0.0, 10.0, 10.0], [1.0, 1.0, 11.0, 11.0]]) scores torch.tensor([0.9, 0.8]) result torchvision.ops.nms(boxes, scores, 0.5) print(f [OK] NMS CPU 测试通过: {result}) except Exception as e: print(f [ERROR] NMS CPU 测试失败: {e}) # 5. 测试 NMS CUDA print(f\n[5] NMS CUDA 测试) if not torch.cuda.is_available(): print(f [SKIP] CUDA 不可用跳过 CUDA 测试) return try: boxes torch.tensor([[0.0, 0.0, 10.0, 10.0], [1.0, 1.0, 11.0, 11.0]], devicecuda) scores torch.tensor([0.9, 0.8], devicecuda) result torchvision.ops.nms(boxes, scores, 0.5) print(f [OK] NMS CUDA 测试通过: {result}) except NotImplementedError as e: print(f [ERROR] NMS CUDA 测试失败: {e}) print(f 这表明 torchvision 的 CUDA 扩展未正确安装或版本不匹配) except Exception as e: print(f [ERROR] NMS CUDA 测试失败: {type(e).__name__}: {e}) # 6. 检查 torchvision ops 的 CUDA 支持 print(f\n[6] torchvision ops CUDA 支持检查) try: # 检查 torchvision._C 是否有 CUDA 相关符号 import torchvision._C as C print(f torchvision._C 模块加载成功) # 检查是否有 _cuda 后端 if hasattr(C, _cuda): print(f [OK] _cuda 后端存在) else: print(f [WARNING] _cuda 后端不存在) except ImportError: print(f [ERROR] torchvision._C 模块加载失败) print(f torchvision 可能未正确编译) def test_all_ops(): 测试所有常用的 torchvision CUDA ops print(\n * 60) print(torchvision CUDA Ops 全面测试) print( * 60) if not torch.cuda.is_available(): print(CUDA 不可用跳过测试) return device cuda tests [ (nms, lambda: torchvision.ops.nms( torch.tensor([[0., 0., 10., 10.]], devicedevice), torch.tensor([0.9], devicedevice), 0.5 )), (batched_nms, lambda: torchvision.ops.batched_nms( torch.tensor([[0., 0., 10., 10.]], devicedevice), torch.tensor([0.9], devicedevice), torch.tensor([0], devicedevice), 0.5 )), (roi_align, lambda: torchvision.ops.roi_align( torch.randn(1, 256, 50, 50, devicedevice), torch.tensor([[0, 0., 0., 10., 10.]], devicedevice), (7, 7) )), (roi_pool, lambda: torchvision.ops.roi_pool( torch.randn(1, 256, 50, 50, devicedevice), torch.tensor([[0, 0., 0., 10., 10.]], devicedevice), (7, 7) )), (deform_conv2d, lambda: torchvision.ops.deform_conv2d( torch.randn(1, 3, 10, 10, devicedevice), torch.randn(1, 27, 8, 8, devicedevice), torch.randn(3, 3, 3, 3, devicedevice) )), (focal_loss, lambda: torchvision.ops.sigmoid_focal_loss( torch.randn(10, 5, devicedevice), torch.randint(0, 5, (10,), devicedevice), alpha0.25, gamma2.0, reductionsum )), ] passed 0 failed 0 for name, test_fn in tests: try: result test_fn() print(f [OK] {name}) passed 1 except NotImplementedError: print(f [FAIL] {name} - CUDA backend not implemented) failed 1 except Exception as e: print(f [FAIL] {name} - {type(e).__name__}: {e}) failed 1 print(f\n结果: {passed} 通过, {failed} 失败) if failed 0: print(\n修复建议:) print( 1. 确保 PyTorch 和 torchvision 版本匹配) print( 2. 确保安装了 CUDA 版本不是 CPU 版本) print( 3. 重新安装:) print( pip uninstall torch torchvision -y) print( pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121) def safe_nms(boxes, scores, iou_threshold): 安全的 NMS 实现自动回退到 CPU try: return torchvision.ops.nms(boxes, scores, iou_threshold) except NotImplementedError: # 回退到 CPU device boxes.device result torchvision.ops.nms(boxes.cpu(), scores.cpu(), iou_threshold) return result.to(device) def safe_batched_nms(boxes, scores, idxs, iou_threshold): 安全的 batched_nms 实现自动回退到 CPU try: return torchvision.ops.batched_nms(boxes, scores, idxs, iou_threshold) except NotImplementedError: # 回退到 CPU device boxes.device result torchvision.ops.batched_nms( boxes.cpu(), scores.cpu(), idxs.cpu(), iou_threshold ) return result.to(device) def demo_safe_ops(): 演示安全 ops 的使用 print(\n * 60) print(安全 NMS Ops 演示) print( * 60) if not torch.cuda.is_available(): print(CUDA 不可用使用 CPU 演示) device cpu else: device cuda # 生成测试数据 num_boxes 100 boxes torch.randn(num_boxes, 4, devicedevice) # 确保 boxes 格式正确 (x1, y1, x2, y2) 且 x2 x1, y2 y1 boxes[:, 2] boxes[:, 0] torch.abs(boxes[:, 2]) 1 boxes[:, 3] boxes[:, 1] torch.abs(boxes[:, 3]) 1 scores torch.rand(num_boxes, devicedevice) idxs torch.randint(0, 5, (num_boxes,), devicedevice) # 使用安全 NMS print(f\n使用 safe_nms (device{device})...) result safe_nms(boxes, scores, 0.5) print(f NMS 保留 {len(result)} / {num_boxes} 个框) # 使用安全 batched_nms print(f\n使用 safe_batched_nms (device{device})...) result safe_batched_nms(boxes, scores, idxs, 0.5) print(f Batched NMS 保留 {len(result)} / {num_boxes} 个框) def main(): 主函数 print(\n * 60) print( torchvision NMS CUDA 问题诊断与修复工具) print( * 60) diagnose() test_all_ops() demo_safe_ops() print(\n * 60) print( 诊断完成) print( * 60) if __name__ __main__: main()常见陷阱与注意事项1. torchvision 版本必须与 PyTorch 版本严格匹配torchvision 的 C 扩展是针对特定 PyTorch 版本编译的。即使只差一个小版本号也可能导致 ABI 不兼容。务必查看官方版本对应表https://github.com/pytorch/vision#installation2. 不要混用 pip 和 conda 安装通过 conda 安装 PyTorch 但通过 pip 安装 torchvision或反之会导致 ABI 不兼容。始终使用同一种包管理器安装两者。3. CPU 回退的性能影响将 NMS 操作回退到 CPU 会带来数据传输开销GPU → CPU → GPU。对于大量边界框如 YOLO 的数千个候选框这个开销可能很明显。建议尽快修复 CUDA 版本问题。4. 从源码编译时的 CUDA 架构从源码编译 torchvision 时必须设置正确的TORCH_CUDA_ARCH_LIST。如果你的 GPU 架构不在列表中CUDA 内核无法运行。可以通过以下命令查看 GPU 架构import torch print(torch.cuda.get_device_capability()) # 例如 (7, 5) 表示 compute capability 7.55. Docker 容器中的 torchvision在 Docker 中使用 torchvision 时确保基础镜像包含正确版本的 torchvision。推荐使用官方 PyTorch Docker 镜像docker run --gpus all -it pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel6. torchvision.ops 的其他操作除了 NMStorchvision.ops 中的其他操作如roi_align、deform_conv2d、focal_loss等也可能遇到同样的 CUDA 后端问题。修复方法相同确保版本匹配和 CUDA 安装。7. 预训练模型的 NMS使用 torchvision 的预训练检测模型如fasterrcnn_resnet50_fpn时NMS 是在模型内部调用的。如果 NMS CUDA 不可用需要修改模型的后处理代码或使用安全回退方案。8. 检查 torchvision 的编译信息import torchvision print(torchvision.__file__) # 检查是否在 site-packages 中pip 安装还是其他位置9. conda 环境中的 cudatoolkit使用 conda 安装时确保cudatoolkit或cuda-toolkit版本与 PyTorch 的 CUDA 版本匹配conda list | grep cuda总结NotImplementedError: Could not run torchvision::nms with arguments from the CUDA backend错误的根本原因是 torchvision 的 CUDA C 扩展未正确安装或与 PyTorch 版本不匹配。解决步骤如下检查版本匹配确保 torchvision 版本与 PyTorch 版本严格对应检查安装类型确保安装的是 CUDA 版本而非 CPU 版本统一安装方式不要混用 pip 和 conda重新安装卸载后使用正确的--index-url重新安装临时回退在修复前使用 CPU 回退方案保持代码可运行通过本文提供的诊断脚本可以快速定位问题所在并根据具体情况进行修复。对于生产环境建议始终使用官方推荐的版本组合并在部署前进行全面的 CUDA ops 测试。