手写数字识别代码详解

文件目录如下，其中数据集data目录运行时在与手写数字识别同级目录自动生成，具体文件内代码见下文

一、conf.py文件

"""
项目配置
"""
import torch

train_batch_size = 128   # 训练批次大小，表示每次训练神经网络时每次使用的图像张量和标签张量的数量为128
test_batch_size = 1000   # 测试批次大小，即测试时为1000
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")   # 判断当前系统是否支持cuda来加速运算，后续用于设置模型和数据在何处运行

二、dataset.py文件

"""
准备数据集
"""
import torch
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST
import torchvision
import conf

def mnist_dataset(train): #准备minist的dataset数据集，输入参数train，表示是否需要加载训练数据集，True则加载，否则加载测试数据集
    func = torchvision.transforms.Compose([   # 将多个数据处理函数组合在一起，以便对数据进行连续的转换操作，以下用来两个转换操作
        torchvision.transforms.ToTensor(),   # 将PIL image或numpy.ndarray的数据类型转为Tensor，并进行归一化处理（0~1之间）
        torchvision.transforms.Normalize(   # 对每个通道进行标准化，即先减均值再除以标准差
            mean=(0.1307,),
            std=(0.3081,)
        )]
    )

    # 1. 准备Mnist数据集
    return  MNIST(root="../data/mnist", train=train, download=True, transform=func)   # root为存储MNIST数据集的目录，train表示是否加载训练数据集，download表示是否从远程下载MNIST数据，transform表示对数据集进行的数据处理操作，最后返回处理后的MNIST数据集

def get_dataloader(train=True):   # 获取数据加载器，train=True表示是否是训练集
    mnist = mnist_dataset(train)   # 调用函数返回MNIST数据集，并将其赋值给mnist变量
    batch_size = conf.train_batch_size if train else conf.test_batch_size  # 根据是否为训练集确定批次大小
    return DataLoader(mnist,batch_size=batch_size,shuffle=True)  # 返回数据加载器对象，DataLoader是pytorch中用于数据加载的工具类，用于从给定的数据集中读取数据并返回迭代器
    # 此处将 MNIST 数据集传递给 DataLoader 对象，并指定批次大小和是否打乱数据顺序。这样，就可以在训练或测试神经网络模型时使用该迭代器来获取数据

if __name__ == '__main__':
    for (images,labels) in get_dataloader():   # 遍历数据加载器中的所有图像和标签数据，将每次循环得到的图像数据赋值给images变量，标签数据赋值给labels变量
        print(images.size())  # 打印图像数据的大小，即图形张量的形状，具体为形如(batch_size,channel,height,width)的4维张量，依次为批次大小、图像通道数、图像高度和宽度
        print(labels.size())   # 打印标签数据的大小，即标签张量的形状，具体为形如(batch_size,)的1维张量，表示批次大小
        break  # 查看第一个批次数据后停止循环，即输出第一个批次的图像和标签数据的大小

三、models.py文件

"""定义模型"""

import torch.nn as nn
import torch.nn.functional as F

class MnistModel(nn.Module):   # 定义神经网络模型，继承pytorch的nn.Module类
    def __init__(self):  # 初始化函数，用于定义网络结构和模型参数
        super(MnistModel,self).__init__()   # 继承nn.Module类的属性和方法
        self.fc1 = nn.Linear(1*28*28,100)   # 定义全连接层，包含100个神经元，即输出特征数量为100
        self.fc2 = nn.Linear(100,10)  # 定义全连接层，包含10个神经元，即输出特征数量为10，输入特征的数量为100

    def forward(self, image):   # 前向传播函数，用于进行模型推断操作，image是一个手写数字图像，
        image_viwed = image.view(-1,1*28*28) #[batch_size,1*28*28]  #形态转换，以便能够输入到全连接层
        fc1_out = self.fc1(image_viwed) #[batch_size,100]    # 将数据输入到第一个全连接层self.fc1中
        fc1_out_relu = F.relu(fc1_out) #[batch_siz3,100]   # ReLU激活函数
        out = self.fc2(fc1_out_relu) #[batch_size,10]   # 输入到第二个全连接层
        return F.log_softmax(out,dim=-1)  # 输出层使用该函数对得到的结果进行对数概率归一化处理，以便能够进行后续的损失计算和反向传播操作

四、train.py文件

"""
进行模型的训练
"""
from dataset import get_dataloader
from models import MnistModel
from torch import optim
import torch.nn.functional as F
import conf
from tqdm import tqdm
from test2 import eval
import numpy as np
import torch
import os

#1. 实例化模型，优化器，损失函数
model = MnistModel().to(conf.device)  # 创建一个MnistModel的实例，并将其移动到指定设备
optimizer = optim.Adam(model.parameters(),lr=1e-3)   # 使用Adam优化算法创建优化器，拥有优化神经网络模型中的参数，前者返回模型中所有需要训练的权重和偏置张量，后者指定学习率为0.001
                                                    # Adam优化算法是一种自适应梯度下降优化算法，可在高维空间中快速、稳定的优化神经网络模型

# if os.path.exists("./models/model.pkl"):
#     model.load_state_dict(torch.load("./models/model.pkl"))
#     optimizer.load_state_dict(torch.load("./models/optimizer.pkl"))


#2. 进行循环，进行训练
def train(epoch):   # 训练神经网络模型，epoch为训练的轮数
    train_dataloader = get_dataloader(train=True)   # 调用函数获取训练数据集的数据加载器
    bar = tqdm(enumerate(train_dataloader),total=len(train_dataloader))   # 使用tqdm()函数创建进度条，enumerate(train_dataloader)表示将数据加载器train_dataloader转换成可迭代的枚举对象，total=len(train_dataloader)表示进度条的总长度
    total_loss = []
    for idx,(input,target) in bar:   # 遍历数据加载器中的所有数据，并将其赋值给input和target
        input = input.to(conf.device)   # 将输入数据转移到指定的计算设备上，即GPU或CPU
        target = target.to(conf.device)  # 将目标数据转移到指定计算设备上
        optimizer.zero_grad()  # 将梯度清零，以避免梯度累加导致错误的梯度更新
        output = model(input)  # 将输入数据input输入到定义好的神经网络模型model中，得到输出预测结果output
        loss = F.nll_loss(output,target)   # 计算损失函数，使用负对数似然损失函数（Negative Log Likelihood Loss)进行分类问题的学习，该函数会将网络的输出结果output和目标数据target代入公式中计算出损失值并返回
        loss.backward()  # 计算反向传播，计算损失对网络参数的梯度
        total_loss.append(loss.item())   # 将计算得到的损失值添加到列表total_loss中
        optimizer.step()   # 根据梯度更新参数
        #打印数据
        if idx%10 ==0 :   # 每迭代10个批次就保存一下模型和优化器参数的状态
            bar.set_description("epcoh:{} idx:{},loss:{:.6f}".format(epoch,idx,np.mean(total_loss)))
            torch.save(model.state_dict(),"./models/model.pkl")   # # 将模型的参数保存到磁盘中
            torch.save(optimizer.state_dict(),"./models/optimizer.pkl")   # 将优化器的参数保存到磁盘中

if __name__ == '__main__':
    for i in range(10):  # 开启训练轮数为10轮的循环
        train(i)  # 调用train()函数开始训练模型
        eval()   # 训练完成后，调用eval()函数对模型进行评估

五、test2.py文件

"""
进行模型的评估
"""
from dataset import get_dataloader
from models import  MnistModel
from torch import optim
import torch.nn.functional as F
import conf
from tqdm import tqdm
import numpy as np
import torch
import os

def eval():   # 测试模型在测试集上的性能表现
    # 1. 实例化模型，优化器，损失函数
    model = MnistModel().to(conf.device)   # 实例化MnistModel的对象并将其移动到指定的设备上（CPU或GPU）

    if os.path.exists("./models/model.pkl"):  # 判断模型是否已保存到本地来恢复模型权重，若不存在，则根据MnistModel类中定义的初始化函数随机生成权重
        model.load_state_dict(torch.load("./models/model.pkl"))
    test_dataloader = get_dataloader(train=False)   # 获取测试数据集的数据加载器
    total_loss = []
    total_acc = []
    with torch.no_grad():
        for input,target in test_dataloader: #2. 进行循环，进行训练，每次使用一个批次的图像输入和对于的标签，并将数据一道指定设备上
            input = input.to(conf.device)
            target = target.to(conf.device)
            #计算得到预测值
            output = model(input)  # 使用模型前向传播获取模型输出
            #得到损失
            loss = F.nll_loss(output,target)   # 计算模型预测与真实标签间的损失loss
            #反向传播，计算损失
            total_loss.append(loss.item())   # 将该批次数据的loss值添加到数组中

            #计算准确率
            ###计算预测值
            pred = output.max(dim=-1)[-1]
            total_acc.append(pred.eq(target).float().mean().item())   # 计算该批次数据的准确率并添加到数组中
    print("test loss:{},test acc:{}".format(np.mean(total_loss),np.mean(total_acc)))   # 计算平均损失和平均准确率并打印

if __name__ == '__main__':
    # for i in range(10):
    #     train(i)
    eval()

六、运行train.py文件

输出如下

epcoh:0 idx:460,loss:0.321847: 100%|██████████| 469/469 [00:22<00:00, 21.20it/s]
test loss:0.1768832817673683,test acc:0.9481000065803528
epcoh:1 idx:460,loss:0.147730: 100%|██████████| 469/469 [00:26<00:00, 17.82it/s]
test loss:0.1185844399034977,test acc:0.9652999997138977
epcoh:2 idx:460,loss:0.101420: 100%|██████████| 469/469 [00:25<00:00, 18.67it/s]
test loss:0.09583198800683021,test acc:0.9711999952793121
epcoh:3 idx:460,loss:0.079119: 100%|██████████| 469/469 [00:25<00:00, 18.42it/s]
test loss:0.08689267784357071,test acc:0.9724000036716461
epcoh:4 idx:460,loss:0.061671: 100%|██████████| 469/469 [00:29<00:00, 15.77it/s]
test loss:0.0784779790788889,test acc:0.9767000019550324
epcoh:5 idx:460,loss:0.051475: 100%|██████████| 469/469 [00:30<00:00, 15.20it/s]
test loss:0.07767344787716865,test acc:0.975299996137619
epcoh:6 idx:460,loss:0.042745: 100%|██████████| 469/469 [00:29<00:00, 16.05it/s]
test loss:0.07635272592306137,test acc:0.9771999955177307
epcoh:7 idx:460,loss:0.034893: 100%|██████████| 469/469 [00:31<00:00, 14.98it/s]
test loss:0.0841908399015665,test acc:0.9741999983787537
epcoh:8 idx:460,loss:0.029657: 100%|██████████| 469/469 [00:31<00:00, 15.10it/s]
test loss:0.09327867105603219,test acc:0.9734000027179718
epcoh:9 idx:460,loss:0.023495: 100%|██████████| 469/469 [00:34<00:00, 13.76it/s]
test loss:0.08141878321766853,test acc:0.976800000667572

学习导航：https://www.xqnav.top/