深度学习——基于PyTorch的MNIST手写数字识别详解

知舟不叙2025-06-14 12:16

文章目录

一、引言

手写数字识别是计算机视觉和深度学习领域的经典入门项目。本文将详细介绍如何使用PyTorch框架构建一个完整的MNIST手写数字识别系统，从数据加载、模型构建到训练和测试的全过程。

二、MNIST数据集简介

MNIST数据集包含70,000张手写数字图像：

  • 60,000张用于训练
  • 10,000张用于测试
  • 图像为灰度图，28×28像素
  • 数字已居中处理，减少了预处理工作
python 复制代码 
import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

三、数据加载与预处理

1. 下载数据集

python 复制代码 
training_data = datasets.MNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor(),
)

test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor(),
)
  • root: 指定数据集存储路径
  • train: True表示训练集，False表示测试集
  • download: 自动下载数据集
  • transform=ToTensor(): 将图像转换为PyTorch张量

2. 数据可视化

python 复制代码 
from matplotlib import pyplot as plt

figure = plt.figure()
for i in range(9):
    img, label = training_data[i+59000]
    figure.add_subplot(3,3,i+1)
    plt.title(label)
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

这段代码展示了训练集中的最后9张图片，帮助我们直观理解数据。

四、数据加载器(DataLoader)

python 复制代码 
train_dataloader = DataLoader(training_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
  • batch_size=64: 将数据分成每批64个样本
  • 优点：减少内存使用，提高训练效率

五、设备配置

python 复制代码 
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

自动检测并选择可用的计算设备（CUDA GPU、Apple MPS或CPU）。

六、神经网络模型

python 复制代码 
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 = self.hidden1(x)
        x = torch.relu(x)
        x = self.hidden2(x)
        x = torch.relu(x)
        x = self.out(x)
        return x

model = NeuralNetwork().to(device)

这是一个简单的全连接神经网络：

  1. nn.Flatten(): 将28×28图像展平为784维向量
  2. 两个隐藏层：128和256个神经元
  3. 输出层：10个神经元对应10个数字类别
  4. 使用ReLU激活函数

七、训练与测试函数

1. 训练函数

python 复制代码 
def train(dataloader, model, loss_fn, optimizer):
    model.train()
    batch_size_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_size_num % 100 == 0:
            print(f"loss: {loss.item():>7f} [number:{batch_size_num}]")
        batch_size_num += 1

关键步骤：

  • model.train(): 设置模型为训练模式
  • 前向传播计算预测值
  • 计算损失
  • 反向传播更新权重
  • 每100个batch打印一次损失

2. 测试函数

python 复制代码 
def Test(dataloader, model, loss_fn):
    size = 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
    correct /= size
    print(f"Test result: \n Accuracy:{(100*correct)}%, Avg loss:{test_loss}")

关键点：

  • model.eval(): 设置模型为评估模式
  • torch.no_grad(): 禁用梯度计算
  • 计算整体准确率和平均损失

八、模型训练与评估

python 复制代码 
loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# 初始训练和测试
train(train_dataloader, model, loss_fn, optimizer)
Test(test_dataloader, model, loss_fn)

# 多轮训练
epochs = 10
for t in range(epochs):
    print(f"epoch {t+1}\n---------------")
    train(train_dataloader, model, loss_fn, optimizer)

print("Done!")
Test(test_dataloader, model, loss_fn)
  • 使用交叉熵损失函数
  • 使用Adam优化器，学习率0.01
  • 训练10个epoch

九、关键概念解析

  1. Epoch vs Batch:

    • Epoch: 整个数据集的一次完整遍历
    • Batch: 一次前向/反向传播处理的数据量

  2. 训练模式与评估模式:

    • model.train(): 启用Dropout和BatchNorm等训练专用层
    • model.eval(): 禁用这些层，确保测试结果一致性

  3. 梯度清零的重要性:

    • PyTorch会累积梯度，必须手动清零(optimizer.zero_grad())

十、优化建议

  1. 增加模型复杂度:

    • 可以尝试卷积神经网络(CNN)提高准确率

  2. 学习率调整:

    • 使用学习率调度器动态调整学习率

  3. 数据增强:

    • 对训练图像进行旋转、平移等增强，提高模型泛化能力

  4. 早停机制:

    • 当验证集准确率不再提升时停止训练，防止过拟合

十一、完整代码

python 复制代码 
import torch
print(torch.__version__)

'''
 MNIST包含70000张手写数字图像：60000用于训练，10000用于测试
 图像是灰度的，28×28像素的，并且居中的，以减少预处理和加快运行
'''
import torch
from torch import nn    #导入神经网络模块
from torch.utils.data import DataLoader  #数据包管理工具，打包数据
from torchvision import  datasets  #封装了很多与图像相关的模型，数据集
from torchvision.transforms import ToTensor  #数据转换，张量，将其他类型的数据转换为tensor张量，numpy array

'''下载训练数据集（包含训练图片+标签）'''
training_data = datasets.MNIST( #跳转到函数的内部源代码，pycharm按下ctrl + 鼠标点击
    root="data", #表示下载的手写数字  到哪个路径。60000
    train=True, #读取下载后的数据中的训练集
    download=True, #如果你之前已经下载过了，就不用下载
    transform=ToTensor(), #张量，图片是不能直接传入神经网络模型
 )   #对于pytorch库能够识别的数据一般是tensor张量


'''下载测试数据集（包含训练图片+标签）'''
test_data = datasets.MNIST( #跳转到函数的内部源代码，pycharm按下ctrl + 鼠标点击
    root="data", #表示下载的手写数字  到哪个路径。60000
    train=False, #读取下载后的数据中的训练集
    download=True, #如果你之前已经下载过了，就不用下载
    transform=ToTensor(), #Tensor是在深度学习中提出并广泛应用的数据类型
 )   #Numpy数组只能在CPU上运行。Tensor可以在GPU上运行。这在深度学习应用中可以显著提高计算速度。
print(len(training_data))

'''展示手写数字图片，把训练集中的59000张图片展示'''
from matplotlib import pyplot as plt
figure = plt.figure()
for i in range(9):
    img,label = training_data[i+59000] #提取第59000张图片

    figure.add_subplot(3,3,i+1) #图像窗口中创建多个小窗口，小窗口用于显示图片
    plt.title(label)
    plt.axis("off")  #plt.show(I) 显示矢量
    plt.imshow(img.squeeze(),cmap="gray") #plt.imshow()将Numpy数组data中的数据显示为图像，并在图形窗口中显示
    a = img.squeeze()  #img.squeeze()从张量img中去掉维度为1的，如果该维度的大小不为1，则张量不会改变
plt.show()

'''创建数据DataLoader（数据加载器）'''
# batch_size:将数据集分为多份，每一份为batch_size个数据
#       优点：可以减少内存的使用，提高训练速度

train_dataloader = DataLoader(training_data,batch_size=64)
test_dataloader = DataLoader(test_data,batch_size=64)


'''判断当前设备是否支持GPU，其中mps是苹果m系列芯片的GPU'''
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")   #字符串的格式化，CUDA驱动软件的功能：pytorch能够去执行cuda的命令
# 神经网络的模型也需要传入到GPU，1个batch_size的数据集也需要传入到GPU，才可以进行训练


''' 定义神经网络  类的继承这种方式'''
class NeuralNetwork(nn.Module): #通过调用类的形式来使用神经网络，神经网络的模型，nn.mdoule
    def __init__(self): #python基础关于类，self类自己本身
        super().__init__() #继承的父类初始化
        self.flatten = nn.Flatten() #展开,创建一个展开对象flatten
        self.hidden1 = nn.Linear(28*28,128) #第1个参数：有多少个神经元传入进来，第2个参数：有多少个数据传出去
        self.hidden2 = nn.Linear(128,256) #第1个参数：有多少个神经元传入进来，第2个参数：有多少个数据传出去
        self.out = nn.Linear(256,10) #输出必须和标签的类别相同，输入必须是上一层的神经元个数
    def forward(self,x):   #前向传播，你得告诉它 数据的流向 是神经网络层连接起来，函数名称不能改
        x = self.flatten(x)  #图像进行展开
        x = self.hidden1(x)
        x = torch.relu(x)   #激活函数，torch使用的relu函数
        x = self.hidden2(x)
        x = torch.relu(x)
        x = self.out(x)
        return x
model = NeuralNetwork().to(device) #把刚刚创建的模型传入到GPU
print(model)

def train(dataloader,model,loss_fn,optimizer):
    model.train() #告诉模型，我要开始训练，模型中w进行随机化操作，已经更新w，在训练过程中，w会被修改的
# pytorch提供2种方式来切换训练和测试的模式，分别是：model.train() 和 mdoel.eval()
# 一般用法是：在训练开始之前写上model.train(),在测试时写上model.eval()
    batch_size_num = 1
    for X,y in dataloader:              #其中batch为每一个数据的编号
        X,y = X.to(device),y.to(device) #把训练数据集和标签传入cpu或GPU
        pred = model.forward(X)         # .forward可以被省略，父类种已经对此功能进行了设置
        loss = loss_fn(pred,y)          # 通过交叉熵损失函数计算损失值loss
        # Backpropagation 进来一个batch的数据，计算一次梯度，更新一次网络
        optimizer.zero_grad()           # 梯度值清零
        loss.backward()                 # 反向传播计算得到每个参数的梯度值w
        optimizer.step()                # 根据梯度更新网络w参数

        loss_value = loss.item()        # 从tensor数据种提取数据出来，tensor获取损失值
        if batch_size_num %100 ==0:
            print(f"loss: {loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num += 1

def Test(dataloader,model,loss_fn):
    size = len(dataloader.dataset)  #10000
    num_batches = len(dataloader)  # 打包的数量
    model.eval()        #测试，w就不能再更新
    test_loss,correct =0,0
    with torch.no_grad():       #一个上下文管理器，关闭梯度计算。当你确认不会调用Tensor.backward()的时候
        for X,y in dataloader:
            X,y = X.to(device),y.to(device)
            pred = model.forward(X)
            test_loss += loss_fn(pred,y).item() #test_loss是会自动累加每一个批次的损失值
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()
            a = (pred.argmax(1) == y) #dim=1表示每一行中的最大值对应的索引号，dim=0表示每一列中的最大值对应的索引号
            b = (pred.argmax(1) == y).type(torch.float)
    test_loss /= num_batches #能来衡量模型测试的好坏
    correct /= size  #平均的正确率
    print(f"Test result: \n Accuracy:{(100*correct)}%, Avg loss:{test_loss}")

loss_fn = nn.CrossEntropyLoss()  #创建交叉熵损失函数对象，因为手写字识别一共有十种数字，输出会有10个结果
#
optimizer = torch.optim.Adam(model.parameters(),lr=0.01) #创建一个优化器，SGD为随机梯度下降算法
# # params：要训练的参数，一般我们传入的都是model.parameters()
# # lr:learning_rate学习率，也就是步长
#
# # loss表示模型训练后的输出结果与样本标签的差距。如果差距越小，就表示模型训练越好，越逼近真实的模型
train(train_dataloader,model,loss_fn,optimizer) #训练1次完整的数据。多轮训练
Test(test_dataloader,model,loss_fn)

epochs = 10
for t in range(epochs):
    print(f"epoch {t+1}\n---------------")
    train(train_dataloader,model,loss_fn,optimizer)
print("Done!")
Test(test_dataloader,model,loss_fn)

十二、总结

本文详细介绍了使用PyTorch实现MNIST手写数字识别的全过程，包括：

  1. 数据加载与预处理
  2. 神经网络模型构建
  3. 训练与测试流程
  4. 关键概念解析

这个项目虽然简单，但包含了深度学习项目的基本要素，是学习PyTorch和计算机视觉的优秀起点。通过调整模型结构和超参数，可以进一步提高识别准确率。

相关推荐
NAGNIP8 小时前
一文搞懂深度学习中的通用逼近定理！
人工智能·算法·面试
冬奇Lab10 小时前
一天一个开源项目（第36篇）：EverMemOS - 跨 LLM 与平台的长时记忆 OS，让 Agent 会记忆更会推理
人工智能·开源·资讯
冬奇Lab10 小时前
OpenClaw 源码深度解析（一）：Gateway——为什么需要一个"中枢"
人工智能·开源·源码阅读
AngelPP13 小时前
OpenClaw 架构深度解析：如何把 AI 助手搬到你的个人设备上
人工智能
宅小年13 小时前
Claude Code 换成了Kimi K2.5后，我再也回不去了
人工智能·ai编程·claude
九狼14 小时前
Flutter URL Scheme 跨平台跳转
人工智能·flutter·github
ZFSS14 小时前
Kimi Chat Completion API 申请及使用
前端·人工智能
天翼云开发者社区15 小时前
春节复工福利就位！天翼云息壤2500万Tokens免费送，全品类大模型一键畅玩！
人工智能·算力服务·息壤
知识浅谈15 小时前
教你如何用 Gemini 将课本图片一键转为精美 PPT
人工智能
Ray Liang15 小时前
被低估的量化版模型，小身材也能干大事
人工智能·ai·ai助手·mindx
热门推荐
01GitHub 镜像站点02【OpenClaw 本地实战 Ep.3】突破瓶颈：强制修改 openclaw.json 解锁 32k 上下文记忆03OpenClaw 使用和管理 MCP 完全指南04OpenClaw + 飞书（Feishu）环境搭建指南05Clawdbot部署教程：解决‘gateway token missing’授权问题的完整步骤06Window 10部署openclaw报错node.exe : npm error code 12807Claude Code + GLM4.7 避坑指南：解决 Unable to connect to Anthropic services08AI 规范驱动开发“三剑客”深度对比：Spec-Kit、Kiro 与 OpenSpec 实战指南09AI Agent 平台横评：ZeroClaw vs OpenClaw vs Nanobot10OpenClaw优化飞书API 额度已耗尽问题