本地运行SAM演示Notebook遇‘torchvision::nms’CUDA后端错误求助

我在本地Jupyter Notebook运行Segment-Anything的automatic_mask_generator_example.ipynb时遇到以下错误：

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

环境对比 #

• Google Colab（可正常运行）：

PyTorch version: 2.0.1+cu118
Torchvision version: 0.15.2+cu118
CUDA is available: True

• 本地环境：

PyTorch version: 2.0.1+cu117
Torchvision version: 0.15.2+cpu
CUDA is available: True

已尝试但无效的安装命令：

• pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117
• pip install torchvision==0.15.0+cu117
• conda install torchvision==0.15.0+cu117 -c pytorch -c nvidia

解决步骤 #

1. 彻底卸载现有CPU版本依赖 #

先清理已安装的包，避免版本冲突：

# pip环境
pip uninstall -y torch torchvision torchaudio

# conda环境（如果用conda管理）
conda uninstall -y torch torchvision torchaudio

2. 确认CUDA 11.7环境变量配置 #

检查本地CUDA 11.7的环境变量是否生效：

# Linux/macOS
echo $CUDA_HOME
echo $PATH | grep cuda
echo $LD_LIBRARY_PATH | grep cuda

# Windows（命令提示符）
echo %CUDA_PATH%

确保输出指向CUDA 11.7的安装路径，未配置则手动添加。

3. 重新安装适配CUDA 11.7的PyTorch套件 #

方法一：pip安装（优先推荐）

执行指定版本的安装命令，确保从官方CUDA 11.7源获取包：

pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu117

方法二：conda安装

若使用conda环境，运行：

conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia

4. 验证安装结果 #

在Jupyter Notebook中运行以下代码确认：

import torch
import torchvision

print(f"PyTorch version: {torch.__version__}")
print(f"Torchvision version: {torchvision.__version__}")
print(f"CUDA is available: {torch.cuda.is_available()}")
print(f"Torchvision CUDA enabled: {torchvision._is_cuda_enabled()}")

需确保torchvision版本带有+cu117后缀，且Torchvision CUDA enabled输出为True。

5. 临时替代方案（安装失败时使用） #

若上述方法仍无法解决，可临时强制SAM在CPU上运行，修改Notebook中的设备设置：

device = torch.device('cpu')

此方案会降低运行速度，仅作为临时应急手段。

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