目录：

一、下载MNIST数据集到本地

二、模型训练及测试图片

简介：

MNIST是深度学习和计算机视觉领域的入门级数据集。

MNIST数据集总共包含两个子数据集：一个训练数据集和一个测试数据集。它们分别包含了60K和10K的28×28的灰度图像。

MNIST文件包含四个不属于任何标准格式的文件，因此它们需要专门进行编程方可阅读，当然也可以使用PyTorch的torchvision.datasets.MINIST API接口直接读取。

此篇文章完成的是对单数字图片的识别！！！

一、下载MNIST数据集到本地

可以通过MNIST官网或者网盘资源下载，下面提供代码导入的方式：

# 下载数据集
from torchvision import datasets, transforms

train_set = datasets.MNIST("data",train=True,download=True, transform=transforms.ToTensor(),)
test_set = datasets.MNIST("data",train=False,download=True, transform=transforms.ToTensor(),)

• datasets.MNIST：是Pytorch的内置函数torchvision.datasets.MNIST，可以导入数据集
• train=True ：读入的数据作为训练集
• transform：读入我们自己定义的数据预处理操作
• download=True：当我们的根目录（root）下没有数据集时，便自动下载
   

根据提示安装好三个软件包后，运行此文件，会生成data目录。

二、模型训练及测试图片

⚠️：

• 导入代码后需存在model目录存放训练数组及test目录存放测试图片，完整文件结构如下：

• 代码最后注释为测试图片部分，在模型训练之后将以下代码部分的80-84行判断循环部分注释，为了在测试时不返回多余信息。将测试部分注释取消，以进行图片测试。

import os
import matplotlib.pyplot as plt
import torch
from PIL import Image
from torch import nn
from torch.nn import Conv2d, Linear, ReLU
from torch.nn import MaxPool2d
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader
transform = transforms.Compose({
    transforms.ToTensor()
})
train_data = MNIST(root='./data', train=True, transform=transform, download=True)
train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True)
test_data = MNIST(root="./data", train=False, transform=transform, download=True)
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=True)

train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度：{}".format(train_data_size))
print("测试数据集的长度：{}".format(test_data_size))

class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        self.conv1 = Conv2d(in_channels=1, out_channels=10, kernel_size=5, stride=1, padding=0)
        self.maxpool1 = MaxPool2d(2)
        self.conv2 = Conv2d(in_channels=10, out_channels=20, kernel_size=5, stride=1, padding=0)
        self.maxpool2 = MaxPool2d(2)
        self.linear1 = Linear(320, 128)
        self.linear2 = Linear(128, 64)
        self.linear3 = Linear(64, 10)
        self.relu = ReLU()

    def forward(self, x):
        x = self.relu(self.maxpool1(self.conv1(x)))
        x = self.relu(self.maxpool2(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = self.linear1(x)
        x = self.linear2(x)
        x = self.linear3(x)

        return x

model = MnistModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.14)

def train():
    for index, data in enumerate(train_loader):
       input, target = data 
       y_predict = model(input)
       loss = criterion(y_predict, target)
       optimizer.zero_grad() 
       loss.backward()
       optimizer.step()
       if index % 100 == 0: 
           torch.save(model.state_dict(), "./model/model.pkl") 
           torch.save(optimizer.state_dict(), "./model/optimizer.pkl")
           print("训练次数为：{}，损失值为：{}".format(index, loss.item() ))


if os.path.exists('./model/model.pkl'):
   model.load_state_dict(torch.load("./model/model.pkl"))

def test():
    correct = 0    
    total = 0  
    with torch.no_grad(): 
        for data in test_loader:
            input, target = data
            output = model(input)  
            probability, predict = torch.max(input=output.data, dim=1)   
         
            total += target.size(0) 
            correct += (predict == target).sum().item()
        print("测试准确率为：%.6f" %(correct / total))

if __name__ == '__main__':
    for i in range(15):
        print({"————————第{}轮测试开始——————".format (i + 1)})
        train()
        test()

def test_mydata():
    image = Image.open('./test/test_2.jpg')
    image = image.resize((28, 28))
    image = image.convert('L')
    transform = transforms.ToTensor()
    image = transform(image)
    image = image.unsqueeze(0) 
    with torch.no_grad():
        output = model(image)
        probabilities = torch.nn.functional.softmax(output, dim=1)
        probability, predict = torch.max(probabilities, dim=1)

    predicted_digit = predict.item()
    prob_value = probability.item()

    print("此手写图片值为：%d, 其最大概率为：%.4f" % (predicted_digit, prob_value))
    plt.title("此手写图片值为：{}".format(predicted_digit), fontname='SimHei')
    plt.imshow(image.squeeze())
    plt.show()

# if __name__ == '__main__':
#     test_mydata()

最终结果如下：