使用的数据集: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)