当前位置：首页 > news >正文

基于 PyTorch 的手写数字识别

news 2025/10/14 7:33:20

一、任务背景

利用 PyTorch 框架，基于 MNIST 数据集实现手写数字识别，通过构建神经网络模型，完成数据加载、预处理、模型训练与结果可视化等流程，直观了解神经网络的工作原理与实现过程。

二、核心步骤与代码解析

（一）准备数据

导入库与模块：导入numpy、torch相关库，包括数据加载（torchvision.datasets.MNIST）、预处理（torchvision.transforms）、数据迭代（torch.utils.data.DataLoader）、神经网络构建（torch.nn）及优化器（torch.optim）等模块。
定义超参数：设置训练批次大小（train_batch_size = 64）、测试批次大小（test_batch_size = 128）、学习率（learning_rate = 0.01）、训练轮数（num_epochs = 20），这些参数影响模型训练效率与效果。
数据预处理与加载：使用transforms.Compose组合ToTensor（将图像转为张量）和Normalize（标准化张量）操作，对 MNIST 数据集进行预处理。通过DataLoader创建训练和测试数据迭代器，shuffle=True使训练数据打乱，增强泛化性，测试数据不打乱以稳定评估。

（二）可视化源数据

利用matplotlib.pyplot库，从测试数据迭代器中取出一批数据，通过subplot和imshow函数可视化部分手写数字图像，并标注真实标签，直观感受数据样貌。

（三）构建模型

定义Net类继承nn.Module，在__init__方法中，使用nn.Sequential组合线性层（nn.Linear）和批归一化层（nn.BatchNorm1d）构建网络结构，包含输入层、两个隐藏层和输出层。forward方法定义前向传播过程，对隐藏层输出使用ReLU激活函数，输出层使用softmax激活函数（按行计算，使输出为各数字类别概率）。

（四）实例化模型与定义优化相关

设备与模型实例化：根据torch.cuda.is_available()判断是否使用 GPU，实例化Net模型并将其移至对应设备。
损失函数与优化器：选择CrossEntropyLoss作为损失函数，适合多分类任务；使用SGD优化器，设置学习率和动量，加速模型收敛。

（五）训练模型

训练循环：遍历训练轮数，在每轮中，先将模型设为训练模式（model.train()）。动态调整学习率（每 5 轮乘以 0.9），平衡训练速度与精度。
前向与反向传播：对训练数据批次，将图像和标签移至设备，通过模型得到输出，计算损失（criterion）。反向传播前清空梯度（optimizer.zero_grad()），然后反向传播计算梯度（loss.backward()），最后更新参数（optimizer.step()）。
记录与评估：记录每轮训练损失和准确率，同时在测试集上评估模型，记录测试损失和准确率，使用SummaryWriter记录训练损失用于后续可视化，打印每轮训练与测试的损失和准确率，监控模型训练过程。