【2】人工智能训练 - 使用模型及保存预测结果

import torch
from torchvision import models
import csv
from sklearn.metrics import classification_report
import cv2 as cv
import os

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 选择模型
model = models.resnet18(num_classes=2)  #此处分类与训练时的保持一致
# 导入模型的参数
model.load_state_dict(torch.load("model-4.pt"))
# 将模型加载到设备上
model.to(device)

# 开始进行预测
model.eval()
# 预测集
test_len = 1390
print("id, label")
file = open("test500.csv", 'w', newline='')
fwrite = csv.writer(file)
fwrite.writerow(["id", 'label'])
pred_label = []

folder_path = 'test/'  #相对路径

with torch.no_grad():


    # for i in range(test_len):
    #     image = cv.imread("test/"+str(i)+".png")


    # for i in range(test_len):
    for filename in os.listdir(folder_path):
        image = cv.imread(os.path.join(folder_path + '/', filename))


        image = cv.resize(image, (112, 112))
        image = image.reshape(3, 112, 112)
        image = torch.tensor(image, dtype=torch.float32)
        # 保证模型的输入的数据维度一致
        images = torch.stack((image,))
        images = images.to(device)
        output = model(images)
        pred = torch.argmax(output, dim=1).item()
        # 把结果保存到定义的列表中
        pred_label.append(pred)
        print('sample '+filename, str(pred))
        # 将结果保存到csv文件中
        fwrite.writerow(['sample '+filename, pred])


# 演示一下预测的效果
#true_pred = [0]*1180+[1]*210
#print("Classification Report:\n", classification_report(true_pred, pred_label))