用PyTorch实现MNIST手写数字数据集的识别

利用卷积网络对MNIST数据集分类，实现0~9的手写数字识别，是卷积神经网络的入门操作。包括（1）数据加载，（2）模型搭建，（3）模型训练和保存，（4）模型调用与测试。具体代码如下

import torch
import torch.nn
import torch.utils.data
import torchvision.datasets
import torchvision.transforms
import matplotlib.pyplot as plt 

#read the data
train_dataset = torchvision.datasets.MNIST(root='./data/mnist',
                                           train=True,
                                           transform=torchvision.transforms.ToTensor(),
                                           download=True)
test_dataset = torchvision.datasets.MNIST(root='./data/mnist',
                                           train=False,
                                           transform=torchvision.transforms.ToTensor(),
                                           download=True)

batch_size=100
train_loader = torch.utils.data.DataLoader(
    dataset=train_dataset, batch_size = batch_size
    )
test_loader = torch.utils.data.DataLoader(
    dataset=test_dataset, batch_size=batch_size
)
print('len(train_loader)={}'.format(len(train_loader)))
print('len(train_loader)={}'.format(len(test_loader)))

#define the Net Structure
class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()   
        self.conv0 = torch.nn.Conv2d(1, 64, kernel_size=3, padding=1)
        self.relu1 = torch.nn.ReLU()
        self.conv2 = torch.nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.relu3 = torch.nn.ReLU()
        self.pool4 = torch.nn.MaxPool2d(stride=2, kernel_size=2)
        self.fc5 = torch.nn.Linear(128*14*14, 1024)
        self.relu6 = torch.nn.ReLU()
        self.drop7 = torch.nn.Dropout(p=0.5)
        self.fc8 = torch.nn.Linear(1024, 10)

    def forward(self, x):
        x=self.conv0(x)
        x=self.relu1(x)
        x=self.conv2(x)
        x=self.relu3(x)
        x=self.pool4(x)
        x=x.view(-1,128*14*14)
        x=self.fc5(x)
        x=self.relu6(x)
        x=self.drop7(x)
        x=self.fc8(x)
        return x
    
net = Net()
print(net)

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters())

#train the Net
num_epochs = 5
for epoch in range(num_epochs):
    for idx, (images, lables) in enumerate(train_loader):
        optimizer.zero_grad()
        preds = net(images)
        loss = criterion(preds, lables)
        loss.backward()
        optimizer.step()

        if idx %5 ==0:
            print('epoch{}, batch{}, loss={:g}'.format(
                epoch, idx, loss.item()
            ))
            
#save the trained net
torch.save(net, 'net.pkl')

#load the trained net
net1 = torch.load('net.pkl')

#test the trained net
correct=0
total=1
for images, labels in test_loader:
    preds = net(images)
    predicted = torch.argmax(preds, 1)
    total += lables.size(0)
    correct += (predicted == labels).sum().item()

accuracy = correct/total
print('accuracy of test data:{:.1%}'.format(accuracy))