PyTorch加载预训练模型至GPU时出现CUDA unknown error问题求助

针对你遇到的这个CUDA未知错误，结合你的设备（Nvidia 940m）和已经尝试过的方案，我整理了几个可能的原因和对应的解决办法：

1. 检查PyTorch版本与显卡算力的兼容性 #

Nvidia 940m属于Kepler架构，CUDA算力为3.5。而最新版的PyTorch已经不再支持算力3.5及以下的显卡——这很可能是核心问题。你可以尝试安装仍支持该算力的PyTorch版本，比如：

• 选择PyTorch 1.12.x版本（这是最后一批支持Kepler架构的正式版之一）
• 对应CUDA版本的安装命令示例（以CUDA 10.0为例）：

  pip install torch==1.12.1+cu102 torchvision==0.13.1+cu102 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu102
  

  若你使用CUDA 9.0，替换为cu92版本的包即可。

2. 强制笔记本使用独立GPU运行程序 #

很多Windows笔记本默认会给Python程序分配集成显卡，即使torch.cuda.is_available()返回True，实际调用GPU时仍会出问题。你可以这样设置：

• 打开Nvidia控制面板 → 管理3D设置 → 程序设置
• 找到你运行代码的Python解释器（比如Anaconda环境里的python.exe）
• 将其首选图形处理器设置为高性能Nvidia处理器
• 应用设置后重启Python环境再运行代码

3. 分步执行模型初始化与设备迁移 #

不要一次性把double()和to(device)连在一起调用，拆分步骤避免初始化冲突：

# 先在CPU加载模型
vgg_model1 = VGGFeatureExtractor(True)
# 转换数据类型
vgg_model1 = vgg_model1.double()
# 再迁移到GPU
vgg_model1 = vgg_model1.to(device)

也可以提前手动初始化CUDA再加载模型：

import torch
# 提前初始化CUDA
torch.cuda.init()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 之后再加载模型
vgg_model1 = VGGFeatureExtractor(True).double().to(device)

4. 排查显存不足问题 #

Nvidia 940m的显存通常只有2GB，预训练VGG模型本身占用不少显存，加上数据类型转换可能额外消耗资源。你可以尝试：

• 先将预训练权重加载到CPU，再迁移到GPU：

  model = VGGFeatureExtractor(pretrained=False)
  state_dict = torch.load('vgg_pretrained_weights.pth', map_location='cpu')
  model.load_state_dict(state_dict)
  model = model.double().to(device)
  

• 关闭不必要的后台程序，释放显存空间。

5. 确认CUDA驱动与Toolkit版本匹配 #

虽然你试过重装，但要确保CUDA Toolkit版本和显卡驱动版本严格对应：

• 比如CUDA 10.0对应的驱动版本至少是411.31，CUDA 9.0对应的驱动至少是384.81
• 可以在Nvidia控制面板的「系统信息」中查看当前驱动版本，再调整CUDA Toolkit版本适配

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