【问题解决】RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False.

问题描述

在加载 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.

问题原因

这个错误通常由以下原因引起：

1. CUDA 不可用：当前环境没有可用的 CUDA 设备（GPU）
2. 模型保存于 CUDA 设备：模型是在有 GPU 的环境中保存的，包含 CUDA 设备信息
3. 加载方式不正确：没有使用 map_location 参数指定加载到 CPU
4. PyTorch 版本问题：不同版本的 PyTorch 对 CUDA 设备处理不同
5. 环境配置问题：CUDA 驱动或运行时未正确安装

解决方案

方案 1：使用 map_location 参数加载模型

import torch

# 正确加载模型到 CPU
model = torch.load('model.pth', map_location=torch.device('cpu'))

# 或使用更简洁的方式
model = torch.load('model.pth', map_location='cpu')

方案 2：检查 CUDA 可用性并动态选择设备

import torch

def load_model(model_path):
    """加载模型，自动根据设备选择"""
    if torch.cuda.is_available():
        print("Loading model to CUDA device")
        return torch.load(model_path)
    else:
        print("Loading model to CPU")
        return torch.load(model_path, map_location='cpu')

# 使用
model = load_model('model.pth')

方案 3：保存模型时使用正确的方式

在保存模型时，使用 torch.save(model.state_dict(), path) 而不是保存整个模型：

# 保存模型状态字典（推荐）
torch.save(model.state_dict(), 'model_state_dict.pth')

# 加载模型状态字典
model = ModelClass()  # 创建模型实例
model.load_state_dict(torch.load('model_state_dict.pth', map_location='cpu'))

方案 4：检查环境配置

# 检查 PyTorch 版本和 CUDA 可用性
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"

# 检查 CUDA 版本
nvcc --version

# 检查 GPU 驱动
nvidia-smi

方案 5：安装 CPU 版本的 PyTorch

如果确实不需要 GPU，可以安装 CPU 版本的 PyTorch：

# 安装 CPU 版本的 PyTorch
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu

方案 6：处理模型中的特定层

对于复杂模型，可能需要单独处理特定层：

import torch

# 加载模型到 CPU
state_dict = torch.load('model.pth', map_location='cpu')

# 处理可能包含 CUDA 设备信息的层
for key in list(state_dict.keys()):
    # 移除可能的设备信息
    if 'module.' in key:
        state_dict[key.replace('module.', '')] = state_dict.pop(key)

# 加载到模型
model.load_state_dict(state_dict)

示例代码

完整的模型加载和处理示例

import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    """简单的模型类"""
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.fc2 = nn.Linear(50, 1)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

def save_model_example():
    """保存模型示例"""
    model = SimpleModel()
    
    # 移到 GPU（如果可用）
    if torch.cuda.is_available():
        model.to('cuda')
        print("Model saved on CUDA device")
    else:
        print("Model saved on CPU")
    
    # 保存整个模型
    torch.save(model, 'model.pth')
    
    # 保存模型状态字典（推荐）
    torch.save(model.state_dict(), 'model_state_dict.pth')

def load_model_example():
    """加载模型示例"""
    print("=== Loading entire model ===")
    try:
        # 尝试直接加载（可能会失败）
        model = torch.load('model.pth')
        print("Loaded model successfully")
    except RuntimeError as e:
        print(f"Error loading model: {e}")
        # 使用 map_location 加载到 CPU
        model = torch.load('model.pth', map_location='cpu')
        print("Loaded model to CPU successfully")
    
    print("\n=== Loading model state dict ===")
    # 创建模型实例
    new_model = SimpleModel()
    # 加载状态字典到 CPU
    new_model.load_state_dict(torch.load('model_state_dict.pth', map_location='cpu'))
    print("Loaded model state dict successfully")
    
    return model, new_model

def test_model(model):
    """测试模型"""
    # 创建测试输入
    test_input = torch.randn(1, 10)
    
    # 将输入移到与模型相同的设备
    if next(model.parameters()).is_cuda:
        test_input = test_input.to('cuda')
    
    # 测试前向传播
    with torch.no_grad():
        output = model(test_input)
    
    print(f"Model output: {output.item()}")
    return output

# 使用示例
if __name__ == "__main__":
    # 保存模型（仅执行一次）
    # save_model_example()
    
    # 加载模型
    model, new_model = load_model_example()
    
    # 测试模型
    print("\n=== Testing entire model ===")
    test_model(model)
    
    print("\n=== Testing model from state dict ===")
    test_model(new_model)
    
    # 检查模型设备
    print(f"\nModel device: {'CUDA' if next(model.parameters()).is_cuda else 'CPU'}")
    print(f"New model device: {'CUDA' if next(new_model.parameters()).is_cuda else 'CPU'}")

常见问题

Q: 为什么模型会包含 CUDA 设备信息？

A: 当模型在 GPU 上训练或使用时，模型的参数会存储在 CUDA 设备上，保存时会包含这些设备信息。

Q: 使用 map_location 会影响模型性能吗？

A: 不会影响模型的结构和权重，只是将模型加载到指定设备。如果在 CPU 上运行，推理速度会比 GPU 慢。

Q: 为什么推荐保存模型状态字典而不是整个模型？

A: 保存状态字典更灵活，不依赖于模型类的定义，并且加载时可以更容易地处理设备映射。

Q: 如何在不同设备间迁移模型？

A: 使用 model.to(device) 可以将模型移到指定设备，例如 model.to('cuda') 或 model.to('cpu')。

Q: 没有 GPU 的环境下可以训练模型吗？

A: 可以，但训练速度会慢很多。对于大型模型，建议使用有 GPU 的环境。

总结

遇到 RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. 错误时，主要需要：

1. 使用 map_location='cpu' 参数加载模型
2. 检查 CUDA 可用性并动态选择加载设备
3. 推荐使用 torch.save(model.state_dict(), path) 保存模型
4. 确保环境配置正确
5. 对于复杂模型，可能需要单独处理特定层

通过以上解决方案，大部分情况下都能成功在 CPU 环境中加载在 GPU 上保存的模型，或者在不同 CUDA 设备间迁移模型。