问题描述
在使用pytorch加载模型时遇到以下错误:
runtimeerror: attempting to deserialize object on a cuda device but torch.cuda.is_available() is false. if you are running on a cpu-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the cpu.
已中止 (核心已转储)
解决方案
方案一:强制加载到cpu(快速解决)
import torch
# 方法1:直接指定cpu设备
model = torch.load('model.pth', map_location=torch.device('cpu'))
# 方法2:使用字符串指定
model = torch.load('model.pth', map_location='cpu')
# 方法3:使用lambda函数
model = torch.load('model.pth', map_location=lambda storage, loc: storage.cpu())
方案二:检查并修复gpu环境(根本解决)
步骤1:检查显卡驱动
nvidia-smi
步骤2:验证pytorch安装
import torch
print(f"pytorch版本: {torch.__version__}")
print(f"cuda可用: {torch.cuda.is_available()}")
print(f"cuda版本: {torch.version.cuda}")
步骤3:重新安装正确的pytorch版本
# 安装支持cuda的pytorch pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 或使用conda conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
方案三:使用state_dict方式(最佳实践)
# 保存模型时使用state_dict
torch.save(model.state_dict(), 'model_weights.pth')
# 加载模型时指定设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
特殊情况处理
amd显卡用户
# amd显卡不支持cuda,使用cpu
device = torch.device('cpu')
apple m系列芯片
# apple m系列芯片使用mps加速
if torch.backends.mps.is_available():
device = torch.device('mps')
else:
device = torch.device('cpu')
完整解决方案代码
import torch
import os
def load_model_safely(model_path, model_class, **kwargs):
"""
安全加载模型,自动处理设备兼容性问题
"""
# 检测可用设备
if torch.cuda.is_available():
device = torch.device('cuda')
print("✅ 检测到cuda设备,使用gpu加速")
elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
device = torch.device('mps')
print("✅ 检测到apple m系列芯片,使用mps加速")
else:
device = torch.device('cpu')
print("💻 未检测到gpu设备,使用cpu")
# 创建模型实例
model = model_class(**kwargs)
try:
# 尝试加载state_dict
state_dict = torch.load(model_path, map_location=device)
# 处理checkpoint格式
if isinstance(state_dict, dict) and 'model_state_dict' in state_dict:
state_dict = state_dict['model_state_dict']
# 处理多gpu训练的模型
if list(state_dict.keys())[0].startswith('module.'):
state_dict = {k.replace('module.', ''): v for k, v in state_dict.items()}
model.load_state_dict(state_dict)
model = model.to(device)
model.eval()
print(f"✅ 模型成功加载到 {device}")
return model, device
except exception as e:
print(f"❌ 模型加载失败: {str(e)}")
raise
# 使用示例
if __name__ == "__main__":
# 假设你有一个模型类mymodel
# model, device = load_model_safely('model.pth', mymodel, num_classes=10)
pass
预防措施
- 保存模型时使用state_dict:
torch.save(model.state_dict(), 'model_weights.pth')
- 加载模型时指定设备:
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.load_state_dict(torch.load('model_weights.pth', map_location=device))
- 定期检查环境配置:
print(f"pytorch版本: {torch.__version__}")
print(f"cuda可用: {torch.cuda.is_available()}")
print(f"gpu数量: {torch.cuda.device_count()}")
if torch.cuda.is_available():
print(f"gpu名称: {torch.cuda.get_device_name(0)}")
常见问题排查
问题1:安装了gpu版pytorch但cuda不可用
- 检查显卡驱动是否最新
- 确认pytorch版本与cuda版本匹配
- 重启python环境
问题2:版本不兼容
- 卸载当前pytorch:
pip uninstall torch torchvision torchaudio - 重新安装匹配版本:参考pytorch官网安装命令
问题3:多环境冲突
- 确认当前使用的python环境
- 检查虚拟环境中的pytorch安装
- 使用
which python和which pip确认路径
注意:在生产环境中,建议始终使用state_dict方式保存和加载模型,以获得最佳的兼容性和灵活性。
到此这篇关于pytorch cuda设备不可用错误的解决方案的文章就介绍到这了,更多相关pytorch cuda设备不可用错误内容请搜索代码网以前的文章或继续浏览下面的相关文章希望大家以后多多支持代码网!
发表评论