MNIST手写数字识别（pytorch）

介绍

MNIST手写数据集共包含70000张手写数字图片以及对应标签，测试集部分为60000张，剩余10000张为测试集，在pytorch中可直接进行下载。
输入图片大小为64x1x28x28，首先经过一次卷积图片大小变为64x6x30x30，最大池化后为64x6x15x15。经过第二次卷积后图片大小为64x16x11x11，最大池化后大小为64x16x5x5。

python代码

import torch
import torch.nn as nn
from torchvision import datasets, transforms
import torch.utils.data

train_dataset = datasets.MNIST(root="C:/Users/lxw/Desktop/MNIST/data_set",
                               train=True,
                               transform=transforms.transforms.ToTensor(),
                               download=True)
test_dataset = datasets.MNIST(root="C:/Users/lxw/Desktop/MNIST/data_set",
                              train=False,
                              transform=transforms.transforms.ToTensor(),
                              download=True)
batch_size = 64
train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=1,
                                          shuffle=True)


class LeNet(nn.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = nn.Sequential(nn.Conv2d(1, 6, 3, 1, 2),
                                   nn.ReLU(),
                                   nn.MaxPool2d(2, 2)
                                   )
        self.conv2 = nn.Sequential(nn.Conv2d(6, 16, 5),
                                   nn.ReLU(),
                                   nn.MaxPool2d(2, 2))
        self.fc1 = nn.Sequential(nn.Linear(16 * 5 * 5, 120),
                                 nn.BatchNorm1d(120),
                                 nn.ReLU())
        self.fc2 = nn.Sequential(nn.Linear(120, 120),
                                 nn.BatchNorm1d(120),
                                 nn.ReLU(),
                                 nn.Linear(120, 10))

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size()[0], -1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x


net = LeNet()
my_loss = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=0.001)
for epoch in range(1):
    i = 0
    for data in train_loader:
        inputs, labels = data
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = my_loss(outputs, labels)
        loss.backward()
        optimizer.step()
        i = i + 1
        if i % 100 == 0:
            print(i, "\t", loss.item())

net.eval()
total = 0
prt = 0
for data in test_loader:
    test_input, test_label = data
    test_output = net(test_input)
    total += 1
    if torch.argmax(test_output, dim=1) == test_label:
        prt += 1
print("accuracy=\t", prt / total)

运行结果

结果显示，卷积神经网络识别的精确度高于98%。