PyTorch GPU训练模型在CPU环境预测时报未找到NVIDIA驱动错误如何解决

问题根因 #

报错是因为代码硬编码了.cuda()调用，强制要求将模型、张量加载到NVIDIA GPU上，无GPU的CPU环境没有对应驱动，就会触发该错误。

解决方案 #

• 第一步：先自动判断当前可用设备，避免硬编码设备类型
  加入如下设备判断代码：

  import torch
  # 自动检测可用设备，有GPU用GPU，无GPU自动切CPU
  device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
  

• 第二步：替换所有硬编码的.cuda()调用，统一用.to(device)方法指定设备
  把原代码里的model.cuda()、x_train = torch.tensor(...).cuda()这类写法全部替换为.to(device)
• 第三步：加载预训练权重时指定map_location参数
  如果是加载训练好的权重做预测，必须在torch.load时加入map_location参数，将权重自动映射到当前可用设备，代码如下：

  # 加载训练好的模型权重，自动映射到当前设备
  model.load_state_dict(torch.load("你的模型权重文件路径.pth", map_location=device))
  

调整后的完整兼容代码 #

import torch
import torch.nn as nn
import numpy as np
import time
import torch.nn.functional as F

# 第一步：先判断可用设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

n_epochs = 6
model = CNN_Text()
loss_fn = nn.CrossEntropyLoss(reduction='sum')
optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=0.001)
# 替换原model.cuda()
model = model.to(device)

# 替换所有张量的.cuda()调用
x_train = torch.tensor(train_X, dtype=torch.long).to(device)
y_train = torch.tensor(train_y, dtype=torch.long).to(device)
x_cv = torch.tensor(test_X, dtype=torch.long).to(device)
y_cv = torch.tensor(test_y, dtype=torch.long).to(device)

# 后续数据集、DataLoader代码无需修改
train = torch.utils.data.TensorDataset(x_train, y_train)
valid = torch.utils.data.TensorDataset(x_cv, y_cv)

train_loader = torch.utils.data.DataLoader(train, batch_size=batch_size, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid, batch_size=batch_size, shuffle=False)

train_loss = []
valid_loss = []

for epoch in range(n_epochs):
    start_time = time.time()
    model.train()
    avg_loss = 0.  
    for i, (x_batch, y_batch) in enumerate(train_loader):
        y_pred = model(x_batch)
        loss = loss_fn(y_pred, y_batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        avg_loss += loss.item() / len(train_loader)
    
    model.eval()        
    avg_val_loss = 0.
    val_preds = np.zeros((len(x_cv),len(le.classes_)))
    
    for i, (x_batch, y_batch) in enumerate(valid_loader):
        y_pred = model(x_batch).detach()
        avg_val_loss += loss_fn(y_pred, y_batch).item() / len(valid_loader)
        val_preds[i * batch_size:(i+1) * batch_size] =F.softmax(y_pred, dim=1).cpu().numpy()
    
    val_accuracy = sum(val_preds.argmax(axis=1)==test_y)/len(test_y)
    train_loss.append(avg_loss)
    valid_loss.append(avg_val_loss)
    elapsed_time = time.time() - start_time 
    print('Epoch {}/{} \t loss={:.4f} \t val_loss={:.4f}  \t val_acc={:.4f}  \t time={:.2f}s'.format(
                epoch + 1, n_epochs, avg_loss, avg_val_loss, val_accuracy, elapsed_time))

单独预测阶段的注意事项 #

预测时输入的张量也需要转到对应设备，示例如下：

# 预测阶段代码示例
model.eval()
# 输入数据转为张量后转到对应设备
input_tensor = torch.tensor(你的输入数据, dtype=torch.long).to(device)
with torch.no_grad(): # 关闭梯度计算，节省内存
    pred = model(input_tensor)
# 结果转CPU处理
pred_result = pred.cpu().numpy()

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