在模型训练的最后一步写测试代码时出现问题

Question

在学习使用自己的数据训练RenNet模型的最后一步 写测试代码时出现的问题

import os
import torch
import numpy as np
import pandas as pd
from tqdm import tqdm
from PIL import Image
import torch.nn.functional as F
from torchvision import datasets
from torchvision import transforms

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

img_path = r"D:\workspace\demo\output\val\三管塔\44010650000000043119111511_20181119111521_AaoMTR.jpg"
dataset_dir = 'output'
test_path = os.path.join(dataset_dir, 'val')
test_transform = transforms.Compose([transforms.Resize(256),
                                     transforms.CenterCrop(224),
                                     transforms.ToTensor(),
                                     transforms.Normalize(
                                         mean=[0.485, 0.456, 0.406],
                                         std=[0.229, 0.224, 0.225])
                                     ])
test_dataset = datasets.ImageFolder(test_path, test_transform)

classes = ["抱杆", "不确定增高架", "单管塔", "地面支撑杆", "角钢塔", "楼面增高架", "美化树", "美化天线", "美化外罩",
           "三管塔", "市政路灯杆"]
frame = Image.open(img_path)
img = test_transform(frame)
img = torch.unsqueeze(img,dim=0)
model = torch.load('save_weights/savename')
model = model.eval().to(device)

with torch.no_grad():
    # 前向传播
    outputs = model(img)
    # 只有一张图就挤压掉batch维度
    outputs = torch.squeeze(outputs)
    # 计算图片属于4个类别的概率
    predict = torch.softmax(outputs, dim=0)
    # 得到类别索引
    predict_cla = torch.argmax(predict).numpy()

# 获取最大预测类别概率
predict_score = round(torch.max(predict).item(), 4)
# 获取预测类别的名称

predict_name = classes[predict_cla]

上边是我的代码（还没有写完）但是运行时出现的问题是

Traceback (most recent call last):
  File "D:\workspace\demo\测试.py", line 35, in <module>
    outputs = model(img)
  File "C:\Users\12301\.conda\envs\cj-env\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "D:\workspace\demo\Model.py", line 73, in forward
    out = self.conv1(x)
  File "C:\Users\12301\.conda\envs\cj-env\lib\site-packages\torch\nn\modules\module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "C:\Users\12301\.conda\envs\cj-env\lib\site-packages\torch\nn\modules\conv.py", line 457, in forward
    return self._conv_forward(input, self.weight, self.bias)
  File "C:\Users\12301\.conda\envs\cj-env\lib\site-packages\torch\nn\modules\conv.py", line 453, in _conv_forward
    return F.conv2d(input, weight, bias, self.stride,
RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

我的思路是将模型加载的GPU上 但是试了一下百度的方法解决不了
我的python版本是3.9的

Answer 1

根据报错信息可以看出，输入数据的类型为 torch.FloatTensor，而模型参数的类型为 torch.cuda.FloatTensor，两者不匹配导致出现了错误。这通常是因为在模型和输入数据之间存在设备（CPU和GPU）不匹配的情况。

在加载模型时，可以指定 map_location 参数来将模型加载到CPU上，即使模型最初在GPU上训练。可以尝试如下代码来解决这个问题：

model = torch.load('save_weights/savename', map_location=torch.device('cpu'))

此外，还需要确保输入数据的类型与模型的输入匹配，可以使用 to 方法将输入数据移动到相同的设备上：

img = img.to(device)

最后再运行代码看看是否能够成功。