资讯动态

PyTorch入门实战:MNIST手写数字识别(全连接神经网络详解)

发布时间:2026/8/23 10:30:41 来源:尧图企业网站定制
引言MNISTMixed National Institute of Standards and Technology手写数字数据集是深度学习领域的经典入门数据集包含0~9共10个类别的手写数字灰度图片每张图片大小为28×28像素。本文将从零开始使用PyTorch搭建一个简单的全连接神经网络又称多层感知机MLP完成MNIST分类任务。文章将详细介绍数据加载、模型定义、训练与测试的完整流程并附上代码和运行结果适合深度学习初学者参考。1. 环境配置与依赖Python 3.8PyTorch 1.10torchvisionmatplotlib2. 导入必要的库import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets from torchvision.transforms import ToTensor import matplotlib.pyplot as plt3. 加载MNIST数据集torchvision.datasets.MNIST提供了便捷的数据下载接口。通过transformToTensor()将原始PIL图像转换为PyTorch张量并将像素值从 [0,255] 归一化到 [0,1] 区间同时将形状从 (H, W) 调整为 (C, H, W)其中 C1灰度图。# 训练集60000张图片 training_data datasets.MNIST( rootdata, # 数据保存目录 trainTrue, # 加载训练集 downloadTrue, # 若本地无数据则自动下载 transformToTensor() # 转换为张量 ) # 测试集10000张图片 test_data datasets.MNIST( rootdata, trainFalse, downloadTrue, transformToTensor() )4. 可视化部分样本为了直观了解数据我们从训练集中取最后9张图片索引59000~59008进行展示。注意img的形状是 (1, 28, 28)需要调用squeeze()去掉维度1才能用matplotlib正常显示灰度图。figure plt.figure(figsize(8, 8)) for i in range(9): img, label training_data[i 59000] # 取后9张 figure.add_subplot(3, 3, i 1) # 3×3子图 plt.title(label) # 标题显示真实标签 plt.axis(off) # 关闭坐标轴 plt.imshow(img.squeeze(), cmapgray) # 显示灰度图像 plt.show()运行后得到下图可以看到手写数字的样式和对应的标签图1MNIST训练集中部分样本及标签从后往前取5. 创建DataLoaderDataLoader负责将数据集分批、打乱可选、并行加载。这里设置batch_size32即每个批次包含32张图片及其标签。train_dataloader DataLoader(training_data, batch_size32, shuffleTrue) test_dataloader DataLoader(test_data, batch_size32) # 查看一个batch的形状 for X, y in test_dataloader: print(fShape of X [N, C, H, W]: {X.shape}) # [32, 1, 28, 28] print(fShape of y: {y.shape} {y.dtype}) # [32] torch.int64 break输出示例Shape of X [N, C, H, W]: torch.Size([32, 1, 28, 28]) Shape of y: torch.Size([32]) torch.int646. 定义神经网络模型我们构建一个简单的全连接网络包含两个隐藏层。网络结构如图所示图2全连接神经网络结构示意图输入层28×28 784个像素通过nn.Flatten()展平为一维向量。隐藏层1全连接层输入784输出128后接ReLU激活函数。隐藏层2全连接层输入128输出256后接ReLU激活函数。输出层全连接层输入256输出10对应10个数字类别未使用softmax因为损失函数内部包含。class NeuralNetwork(nn.Module): def __init__(self): super().__init__() self.flatten nn.Flatten() # 展平层28*28 - 784 self.hidden1 nn.Linear(28*28, 128) # 输入784 - 128 self.hidden2 nn.Linear(128, 256) # 128 - 256 self.out nn.Linear(256, 10) # 256 - 10 def forward(self, x): x self.flatten(x) x self.hidden1(x) x torch.relu(x) # ReLU激活 x self.hidden2(x) x torch.relu(x) x self.out(x) return x为什么需要 Flatten全连接层nn.Linear要求输入是二维的[batch_size, features]而原始图像数据是三维的[batch_size, 1, 28, 28]因此必须将每个样本的像素按顺序排成一列。6.1 选择设备自动检测GPUCUDA或Apple MPS若使用Mac否则使用CPU。device cuda if torch.cuda.is_available() else mps if torch.backends.mps.is_available() else cpu print(fUsing device: {device}) model NeuralNetwork().to(device) print(model)输出示例Using device: cuda NeuralNetwork( (flatten): Flatten(start_dim1, end_dim-1) (hidden1): Linear(in_features784, out_features128, biasTrue) (hidden2): Linear(in_features128, out_features256, biasTrue) (out): Linear(in_features256, out_features10, biasTrue) )7. 定义损失函数与优化器损失函数多分类问题使用交叉熵损失CrossEntropyLoss它内部包含了softmax操作因此模型输出层不需要额外添加激活函数。优化器选择Adam优化器学习率设为0.01。Adam结合了动量与自适应学习率的优点通常比SGD收敛更快。loss_fn nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.01)8. 训练函数train函数完成一个epoch遍历一次所有训练数据的参数更新。流程包括设置模型为训练模式model.train()对Dropout、BatchNorm等层生效。遍历DataLoader获取批次数据并送入设备。前向传播计算预测值和损失。反向传播梯度清零、计算梯度、更新参数。每100个batch打印一次损失值。def train(dataloader, model, loss_fn, optimizer): model.train() # 设置为训练模式 batch_num 1 for X, y in dataloader: X, y X.to(device), y.to(device) # 前向传播 pred model(X) loss loss_fn(pred, y) # 反向传播 optimizer.zero_grad() # 梯度清零 loss.backward() # 计算梯度 optimizer.step() # 更新参数 # 每100个batch输出一次损失 if batch_num % 100 0: print(floss: {loss.item():7f} [batch: {batch_num}]) batch_num 19. 测试函数测试函数用于评估模型在测试集上的表现并计算平均损失和准确率。注意使用model.eval()将模型切换为评估模式关闭Dropout等训练专用操作。使用torch.no_grad()上下文管理器禁用梯度计算节省内存和计算时间。pred.argmax(1)取每个样本预测概率最大的类别索引因为输出层有10个神经元对应logits。累加正确预测数最后除以总样本数得到准确率。def test(dataloader, model, loss_fn): total len(dataloader.dataset) num_batches len(dataloader) model.eval() # 设置为评估模式 test_loss, correct 0, 0 with torch.no_grad(): for X, y in dataloader: X, y X.to(device), y.to(device) pred model(X) test_loss loss_fn(pred, y).item() correct (pred.argmax(1) y).type(torch.float).sum().item() test_loss / num_batches accuracy correct / total print(fTest accuracy: {100*accuracy:.2f}%, Avg loss: {test_loss:.4f})10. 训练与评估10.1 先训练一个epoch并测试print(Training for 1 epoch:) train(train_dataloader, model, loss_fn, optimizer) test(test_dataloader, model, loss_fn)10.2 训练10个epoch并观察效果epochs 10 for epoch in range(epochs): print(f\nEpoch {epoch1}\n-------------------------------) train(train_dataloader, model, loss_fn, optimizer) print(Training completed!) test(test_dataloader, model, loss_fn)11. 运行结果分析经过10个epoch的训练测试集准确率通常可以达到97%以上平均损失在0.1左右。以下是一次典型训练过程的输出部分可以看到随着训练的进行损失值逐渐下降准确率稳步提升。最终在测试集上达到了约95%的准确率表明模型已经较好地学会了手写数字的分类。12. 完整代码将上述所有代码片段整合即可得到完整的训练脚本。为了方便读者以下是整理后的完整代码可直接复制保存为.py文件运行import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import datasets from torchvision.transforms import ToTensor import matplotlib.pyplot as plt # 1. 加载数据集 training_data datasets.MNIST(rootdata, trainTrue, downloadTrue, transformToTensor()) test_data datasets.MNIST(rootdata, trainFalse, downloadTrue, transformToTensor()) # 2. 可视化可选 figure plt.figure(figsize(8, 8)) for i in range(9): img, label training_data[i59000] figure.add_subplot(3, 3, i1) plt.title(label) plt.axis(off) plt.imshow(img.squeeze(), cmapgray) plt.show() # 3. DataLoader train_loader DataLoader(training_data, batch_size32, shuffleTrue) test_loader DataLoader(test_data, batch_size32) # 4. 设备 device cuda if torch.cuda.is_available() else mps if torch.backends.mps.is_available() else cpu print(fUsing device: {device}) # 5. 模型定义 class NeuralNetwork(nn.Module): def __init__(self): super().__init__() self.flatten nn.Flatten() self.hidden1 nn.Linear(28*28, 128) self.hidden2 nn.Linear(128, 256) self.out nn.Linear(256, 10) def forward(self, x): x self.flatten(x) x torch.relu(self.hidden1(x)) x torch.relu(self.hidden2(x)) x self.out(x) return x model NeuralNetwork().to(device) print(model) # 6. 损失函数和优化器 loss_fn nn.CrossEntropyLoss() optimizer torch.optim.Adam(model.parameters(), lr0.01) # 7. 训练函数 def train(dataloader, model, loss_fn, optimizer): model.train() batch_num 1 for X, y in dataloader: X, y X.to(device), y.to(device) pred model(X) loss loss_fn(pred, y) optimizer.zero_grad() loss.backward() optimizer.step() if batch_num % 100 0: print(floss: {loss.item():7f} [batch: {batch_num}]) batch_num 1 # 8. 测试函数 def test(dataloader, model, loss_fn): total len(dataloader.dataset) num_batches len(dataloader) model.eval() test_loss, correct 0, 0 with torch.no_grad(): for X, y in dataloader: X, y X.to(device), y.to(device) pred model(X) test_loss loss_fn(pred, y).item() correct (pred.argmax(1) y).type(torch.float).sum().item() test_loss / num_batches accuracy correct / total print(fTest accuracy: {100*accuracy:.2f}%, Avg loss: {test_loss:.4f}) # 9. 开始训练 print(Training for 1 epoch:) train(train_loader, model, loss_fn, optimizer) test(test_loader, model, loss_fn) epochs 10 for epoch in range(epochs): print(f\nEpoch {epoch1}) train(train_loader, model, loss_fn, optimizer) print(Training done!) test(test_loader, model, loss_fn)13. 总结与拓展本文从零开始实现了基于PyTorch的MNIST手写数字识别涵盖了数据加载、模型构建、训练与测试等关键步骤。全连接神经网络虽然简单但完整展示了深度学习的核心流程前向传播、计算损失、反向传播、参数更新。13.1 进一步改进的方向增加模型复杂度可以尝试添加更多隐藏层或者使用卷积神经网络CNN来提取空间特征通常会显著提升准确率。调整超参数学习率、批次大小、优化器类型、激活函数等都会影响训练效果可以尝试不同的组合。数据增强对训练图像进行随机旋转、平移等变换可以提高模型的泛化能力。正则化添加Dropout层或L2正则化防止过拟合。希望这篇文章能帮助读者快速入门PyTorch为后续学习更复杂的深度学习任务打下坚实基础。如果有任何疑问或建议欢迎在评论区留言交流

读完文章,也想定制专属网站?

尧图设计师 24 小时内与您沟通定制方案

免费获取报价