pytorch 读取训练好的Lenet5模型并进行测试,显示错误结果

使用的数据集:MNIST

使用的网络结构:Lenet

mport torch.nn as nn
import torch.nn.functional as F

import torch
import torchvision

import numpy as np
import matplotlib.pyplot as plt
import cv2

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

transform = torchvision.transforms.Compose([
        torchvision.transforms.ToTensor(), # 转为Tensor
        torchvision.transforms.Normalize((0.5,), (0.5,)), # 归一化
                             ])

batch_size = 4
test_dataset = torchvision.datasets.MNIST(root='./mnist', train=False, transform=transform, download=True)
testloader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=12)

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, 5, padding=2)
        self.conv2 = nn.Conv2d(6, 16, 5)  
        self.fc1   = nn.Linear(16*5*5, 120)  
        self.fc2   = nn.Linear(120, 84)
        self.fc3   = nn.Linear(84, 10)

    def forward(self, x): 
        x = F.max_pool2d(F.relu(self.conv1(x)), (2,2)) 
        x = F.max_pool2d(F.relu(self.conv2(x)), 2) 
        x = x.view(x.size()[0], -1)   #展开成一维的
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)        
        return x

net = Net()
net.to(device)
print(net)

PATH = './mnist_net_100.pth'
pretrained_net = torch.load(PATH)
net.load_state_dict(pretrained_net)

#整个测试集上预测
correct = 0
total = 0

err_sum = 0

with torch.no_grad():
    for (images,labels) in testloader:
        images = images.to(device)
        labels = labels.to(device)
        outputs = net(images)
        _, predicted = torch.max(outputs,1)
        
        total += labels.size(0)
        correct += (predicted == labels).sum()
        
        
        for index in range(4):
            if predicted[index] != labels[index]:
                # print(index)
                err_sum = err_sum+1
                
                img = np.empty((28,28), dtype=np.float32)
                img[:,:] = images[index].cpu().numpy()/2+0.5
                img2 = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)                
                plt.imshow(img2)
                plt.show()  
                
print(err_sum)
print('10000张测试集合中的准确率为:', (correct.cpu().numpy()/total * 100))
print(correct)

 

你可能感兴趣的:(图像处理)