【pytorch】使用训练好后的模型权重，在验证集上输出分类的混淆矩阵并保存错误图片-Toy模板网

这篇具有很好参考价值的文章主要介绍了【pytorch】使用训练好后的模型权重，在验证集上输出分类的混淆矩阵并保存错误图片。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

在机器学习领域，混淆矩阵是一个非常有用的指标，它可以帮助我们更好地理解模型在验证集上的表现。本文介绍了如何使用pytorch框架，利用训练好后的模型权重，在验证集上输出分类的混淆矩阵，并保存错误图片的方法。

首先，我们需要准备一个pytorch框架的模型，并将模型权重载入到模型中。然后，我们可以使用pytorch的dataloader加载验证集，并使用模型进行预测。接下来，我们可以使用sklearn的confusion_matrix函数，计算出验证集上的混淆矩阵。最后，我们可以获取预测错误的图片，并将它们保存起来，以便后续分析。

总之，本文介绍了如何使用pytorch框架，利用训练好后的模型权重，在验证集上输出分类的混淆矩阵，并保存错误图片的方法。通过使用混淆矩阵，我们可以更好地理解模型在验证集上的表现，并及时发现模型的问题，从而提高模型的准确率。

import matplotlib.pyplot as plt
import torchvision
import torch
import torch.nn as nn
import torchvision.models as models
import torchvision.transforms as transforms
import numpy as np
import torch.nn.functional as F
from PIL import Image
from sklearn.metrics import confusion_matrix
import os
# 加载resnet18模型
resnet18 = models.resnet18(pretrained=False)

# 获取resnet18最后一层输出，输出为512维,最后一层本来是用作 分类的，原始网络分为1000类
# 用 softmax函数或者 fully connected 函数，但是用 nn.identtiy() 函数把最后一层替换掉，相当于得到分类之前的特征！
#Identity模块，它将输入直接传递给输出，而不会对输入进行任何变换。
resnet18.fc = nn.Identity()

# 构建新的网络，将resnet18的输出作为输入
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        #全卷积神经网络，不用调整输入大小
        self.resnet18 = resnet18
        self.fc1 = nn.Linear(512, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)
        self.fc5 = nn.Linear(10, 2)
        self.softmax = nn.Softmax(dim=1)
    def forward(self, x):
        x = self.resnet18(x)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = F.relu(self.fc3(x))
        x = F.relu(self.fc4(x))
        x = F.relu(self.fc5(x))
        x = self.softmax(x)
        x=x.view(-1,2)
        return x

# 实例化网络
net = Net()
net.load_state_dict(torch.load('./bestmodel98.4375.pth'))
# 将模型放入GPU
net = net.cuda()

def rgb2bgr(image):
    image = np.array(image)[:, :, ::-1]
    image=Image.fromarray(np.uint8(image))
    return image
transform2 = transforms.Compose([
    transforms.Resize(112),
    # rgb转bgr
    torchvision.transforms.Lambda(rgb2bgr),
    # 入的图片为PIL image 或者 numpy.nadrry格式的图片，其shape为（HxWxC）数值范围在[0,255],转换之后shape为（CxHxw）,数值范围在[0,1]
    transforms.ToTensor(),
    # 进行归一化和标准化，Imagenet数据集的均值和方差为：mean=(0.485, 0.456, 0.406)，std=(0.229, 0.224, 0.225)，
    # 因为这是在百万张图像上计算而得的，所以我们通常见到在训练过程中使用它们做标准化。
    transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
])

val_dataset = torchvision.datasets.ImageFolder(r'D:\eyeDataSet\validate',transform=transform2)
val_loader = torch.utils.data.DataLoader(val_dataset, batch_size=256, shuffle=True)

# 验证集上的表现情况：
correct = 0
total = 0
labelsAll=[]
predictedAll=[]
with torch.no_grad():
    for data in val_loader:
        images, labels = data
        images, labels = images.cuda(), labels.cuda()
        outputs = net(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        labelsAll+=labels.cpu()
        predictedAll +=predicted.cpu()

print('Accuracy of the network on the validation images: %.3f %%' % (100 * correct / total))
# 增加混淆矩阵的输出：
y_true = np.array(labelsAll)
y_pred = np.array(predictedAll)
cm = confusion_matrix(y_true, y_pred)
print(cm)
# 替换分类序号与分类名关系：
classes = {}
for index, name in enumerate(val_dataset.class_to_idx):
    classes[index] = name
for i in range(len(cm)):
    for j in range(len(cm[i])):
        print('{} 被预测为 {}: 共{}个'.format(classes[i], classes[j], cm[i][j]))

#将预测失败的图片保存进文件夹：
# 计算预测失败的图片
fail_index=[]
for i in range(len(labelsAll)):
    if labelsAll[i]!=predictedAll[i]:
        fail_index.append(i)

# 将预测失败的图片保存进某指定文件夹内，后缀名为.jpg
for i in fail_index:
    img_name=val_dataset.imgs[i][0]
    img = Image.open(img_name)
    # 检查文件夹是否存在
    if not os.path.exists(r'D:\fail_img'):
        os.makedirs(r'D:\fail_img')
    img.save(r'D:\fail_img\fail_img_'+str(i)+'.jpg')