从零实现MNIST数据集中的识别手写数字

一、数据集介绍

MNIST数据集是一个经典的手写数字数据集，广泛用于机器学习和图像识别任务的测试和训练。MNIST数据集包含了60000个训练样本和10000个测试样本，每个样本都是一个28x28像素的灰度图像，表示一个手写数字（0到9）。

二、获取数据集

1.使用以下代码进行下载，通过修改root=‘ ’ 设置下载目录

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

2.下载成功，如下图

3.通过脚本将MNIST数据集的二进制格式转为图片格式
（1）训练集
注意： 修改以下两行代码的路径，路径分别为：train-images-idx3-ubyte 和train-labels-idx1-ubyte的路径

在这里插入代码片

data_file = r’./data/MNIST/raw/train-images-idx3-ubyte’

label_file = r’./data/MNIST/raw/train-labels-idx1-ubyte’

import numpy as np
import struct
from PIL import Image
import os

print("Processing started")

data_file = r'./data/MNIST/raw/train-images-idx3-ubyte'
# It's 47040016B, but we should set to 47040000B
data_file_size = 47040016
data_file_size = str(data_file_size - 16) + 'B'
data_buf = open(data_file, 'rb').read()
magic, numImages, numRows, numColumns = struct.unpack_from('>IIII', data_buf, 0)
datas = struct.unpack_from('>' + data_file_size, data_buf, struct.calcsize('>IIII'))
datas = np.array(datas).astype(np.uint8).reshape(numImages, 1, numRows, numColumns)

label_file = r'./data/MNIST/raw/train-labels-idx1-ubyte'
# It's 60008B, but we should set to 60000B
label_file_size = 60008
label_file_size = str(label_file_size - 8) + 'B'
label_buf = open(label_file, 'rb').read()
magic, numLabels = struct.unpack_from('>II', label_buf, 0)
labels = struct.unpack_from('>' + label_file_size, label_buf, struct.calcsize('>II'))
labels = np.array(labels).astype(np.int64)

datas_root = 'mnist_train'
if not os.path.exists(datas_root):
    os.mkdir(datas_root)

for i in range(10):
    file_name = datas_root + os.sep + str(i)
    if not os.path.exists(file_name):
        os.mkdir(file_name)

print("Processing images...")

for ii in range(numLabels):
    if ii % 1000 == 0:
        print(f"Processed {ii}/{numLabels} images...")
    img = Image.fromarray(datas[ii, 0, 0:28, 0:28])
    label = labels[ii]
    file_name = datas_root + os.sep + str(label) + os.sep + 'mnist_train_' + str(ii) + '.png'
    img.save(file_name)

print("Processing completed")

（2）测试集
注意： 修改以下两行代码的路径，路径分别为：t10k-images-idx3-ubyte 和t10k-labels-idx1-ubyte的路径

data_file = r'./data/MNIST/raw/t10k-images-idx3-ubyte'

label_file = r'./data/MNIST/raw/t10k-labels-idx1-ubyte'

import numpy as np
import struct
from PIL import Image
import os

print("Processing started")

data_file = r'./data/MNIST/raw/t10k-images-idx3-ubyte'
# It's 7840016B, but we should set to 7840000B
data_file_size = 7840016
data_file_size = str(data_file_size - 16) + 'B'
data_buf = open(data_file, 'rb').read()
magic, numImages, numRows, numColumns = struct.unpack_from('>IIII', data_buf, 0)
datas = struct.unpack_from('>' + data_file_size, data_buf, struct.calcsize('>IIII'))
datas = np.array(datas).astype(np.uint8).reshape(numImages, 1, numRows, numColumns)

label_file = r'./data/MNIST/raw/t10k-labels-idx1-ubyte'
# It's 10008B, but we should set to 10000B
label_file_size = 10008
label_file_size = str(label_file_size - 8) + 'B'
label_buf = open(label_file, 'rb').read()
magic, numLabels = struct.unpack_from('>II', label_buf, 0)
labels = struct.unpack_from('>' + label_file_size, label_buf, struct.calcsize('>II'))
labels = np.array(labels).astype(np.int64)

datas_root = 'mnist_test'
if not os.path.exists(datas_root):
    os.mkdir(datas_root)

for i in range(10):
    file_name = datas_root + os.sep + str(i)
    if not os.path.exists(file_name):
        os.mkdir(file_name)

print("Processing images...")

for ii in range(numLabels):
    if ii % 1000 == 0:
        print(f"Processed {ii}/{numLabels} images...")
    img = Image.fromarray(datas[ii, 0, 0:28, 0:28])
    label = labels[ii]
    file_name = datas_root + os.sep + str(label) + os.sep + 'mnist_test_' + str(ii) + '.png'
    img.save(file_name)

print("Processing completed")

4.下载成功的图片，会在上级目录中的mnist_train和mnist_test文件中，也可自己修改图片生成的目录

三、建立神经网络

该神经网络包含两个线性层

import torch
from torch import nn
#定义神经网络network
class Network(nn.Module):
    def __init__(self):
        super().__init__()

        #线性层1，输入层和隐藏层之间的线性层
        self.layer1 = nn.Linear(784,256)
        #线性层2，隐藏层和输出层之间的线性层
        self.layer2 = nn.Linear(256, 10)

    #前向传播
    def forward(self,x):
        x = x.view(-1,784) #使用view函数，将x展平
        x = self.layer1(x) #将x输入到layer1
        x = torch.relu(x)  #使用relu激活
        return self.layer2(x)  #输入到layer2计算结果

    """
    没有直接定义softmax层
    因为后面代码会使用CrossEntropyLoss损失函数
    这个损失函数会实现softmax的计算
    MNIST数据集，可以从torchvision.datasets中获取
    """

四、加载数据集

注意： 将root=‘’改成自己数据集的目录

import torchvision
from torchvision import transforms
from torchvision import datasets
from torch.utils.data import DataLoader

#加载数据集
if __name__=='__main__':
    transform = transforms.Compose([
        transforms.Grayscale(num_output_channels=1), #转换为单通道灰度图
        transforms.ToTensor() #转换为张量
    ])
    """
    使用ImageFolder函数，读取数据文件夹，构建数据集dataset
    这个函数会将保存数据的文件夹的名字，作为数据的标签，组织数据
    例如，对于名字为“3”的文件夹
    会将“3”作为文件夹中图像数据的标签，和图谱配对，用于后续的训练，使用起来非常的方便
    """
    #root 为数据集的地址
    train_dataset = datasets.ImageFolder(root='./mnist_train', transform=transform)
    test_dataset = datasets.ImageFolder(root='./mnist_test', transform=transform)

    #打印他们的长度
    print("train_dataset length", len(train_dataset))
    print("test_dataset length", len(test_dataset))

    #使用train_loader,实现小批量的数据读取
    train_loader =DataLoader(train_dataset, batch_size=64, shuffle=True)
    #打印train_loader的长度,938*64=60032,说明最后一组，没有64个数据
    print("train_loader length", len(train_loader))

    #循环遍历train_loader
    #每一次循环，都会提取出64个图像数据，作为一个小批量batch
    for batch_idx, (data, label) in enumerate(train_loader):
        if batch_idx ==3: #打印前三个batch观察
            break
        print("batch_idx", batch_idx)
        print("data.shape", data.shape)#数据的尺寸
        print("label.shape", label.shape)#图像中的数字/标签
        print(label)

五、模型的训练

模型训练完成之后会生成pth文件

import os
import torch
from torch import nn
from torch import optim
from torchvision import datasets
from torchvision import transforms
from torch.utils.data import DataLoader
from Network import Network

if __name__=='__main__':
    # 图像的预处理
    transform = transforms.Compose([
        transforms.Grayscale(num_output_channels=1),#转换为单通道灰度图
        transforms.ToTensor() #转换为张量
    ])

    #加载数据集
    train_dataset = datasets.ImageFolder(root='./mnist_train', transform=transform)
    print("train_dataset length", len(train_dataset))

    #小批量的数据读入
    train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
    print("train_loader length", len(train_loader))

    #模型优化，在使用Pytorch训练模型时，需要创建三个对象：
    model = Network() #模型本身，它就是我们设计的神经网络
    optimizer = optim.Adam(model.parameters()) #优化器，优化模型中的参数
    criterion = nn.CrossEntropyLoss() #损失函数，分类问题，使用交叉熵损失误差

    #创建存放结果的文件夹
    result = './result'
    if not os.path.exists(result):
        os.mkdir(result)

    #进入模型的迭代循环
    for epoch in range(10):#外层循环，代表了整个训练数据集的遍历次数
        # 整个训练集要循环多少轮，10/20/100都可能

        # 内存循环使用train_loader,进行小批量的数据读取
        for batch_idx, (data, label) in enumerate(train_loader):
            #内层每循环一次，就会进行一次梯度下降算法
            #5个步骤：
            output = model(data) #1.计算神经网络的前向传播结果
            loss = criterion(output,label)#2.计算output和标签之间的损失loss
            loss.backward() #3.使用后向传播计算梯度
            optimizer.step() #4.使用优化器更新参数
            optimizer.zero_grad() #5.梯度清零
            #每迭代100个小批量，就打印一次模型的损失，观察训练的过程
            if batch_idx % 100 ==0:
                print(f"Epoch: {epoch+1}/10"
                      f"Batch: {batch_idx+1}/{len(train_loader)}"
                      f"Loss: {loss.item():.4f}")
        file_name = os.path.join(result, f'mnist{epoch + 1}.pth')

        torch.save(model.state_dict(), file_name)  # 保存模型

模型训练完成之后会生成pth文件

六、模型的测试

注意：在测试之前，要加载训练之后的pth文件,修改成自己pth文件目录

 model.load_state_dict(torch.load('./result/mnist10.pth'))#加载刚训练好的模型文件

from Network import Network
from torchvision import transforms
from torchvision import datasets
import torch

if __name__=='__main__':
    # 图像的预处理
    transform = transforms.Compose([
        transforms.Grayscale(num_output_channels=1),#转换为单通道灰度图
        transforms.ToTensor() #转换为张量
    ])

    #读取测试数据集
    test_dataset = datasets.ImageFolder(root='./mnist_test', transform=transform)
    print("test_dataset length", len(test_dataset))

    model = Network() #定义神经网络模型
    model.load_state_dict(torch.load('./result/mnist10.pth'))#加载刚训练好的模型文件

    right = 0 #保存正确识别的数量
    for i,(x,y) in enumerate(test_dataset):
        output = model(x)#将其中的数据x输入到模型
        predict = output.argmax(1).item() #选择概率最大标签的作为预测结果
        #对比预测值predict和真实标签y
        if predict == y:
            right+=1
        else:
            #将识别错误的样例打印出来
            img_path = test_dataset.samples[i][0]
            print(f"wrong case: predict = {predict} y={y} img_path={img_path}")

    #计算出测试效果
    sample_num = len(test_dataset)
    acc = right * 1.0 / sample_num
    print("test accuracy = %d / %d = %.3lf"% (right, sample_num, acc))