pytorch测试lenet网络mnist手写体

1.源码实现

import torch
import torch.nn as nn
import torch.optim as optim
from tqdm import tqdm
from torch.utils.data import Dataset, DataLoader
from torchvision import datasets, transforms

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

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        x = x.view(-1, 16*4*4)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

tran_dataset = datasets.MNIST('mnist/', download=False, train=True, transform=transform)
test_dataset = datasets.MNIST('mnist/', download=False, train=False, transform=transform)

train_dataloader = DataLoader(tran_dataset, batch_size=256, shuffle=True)
test_dataloader = DataLoader(test_dataset, batch_size=256, shuffle=False)

device = "cpu"

lenet = LeNet().to(device)

epochs = 1
lr = 1e-4

optimizer = torch.optim.Adam(lenet.parameters(), lr=lr)
loss_fn = nn.CrossEntropyLoss()

train_acc_list = []
test_acc_list = []
train_loss_list = []

lenet.load_state_dict(torch.load("mnist.pth"))

for epoch in range(epochs):
    test_loss_epoch = []
    acc = 0

    with torch.no_grad():
        for test_data, labels in tqdm(test_dataloader):
            test_data = test_data.to(device)
            labels = labels.to(device)
            y_hat = lenet(test_data)
            test_loss = loss_fn(y_hat, labels)
            
            test_loss_epoch.append(test_loss.cpu().detach().numpy())
            right = torch.argmax(y_hat, 1) == labels
            acc += right.sum().cpu().detach().numpy()

        acc = acc / len(tran_dataset)
        test_acc_list.append(acc)

        print(f'test_loss:{sum(test_loss_epoch) / len(test_loss_epoch)}')

2.运行程序

$ python test.py

3.结果

100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 40/40 [00:04<00:00,  8.44it/s]
test_loss:0.07388359744313674