深度学习——基于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和计算机视觉的优秀起点。通过调整模型结构和超参数，可以进一步提高识别准确率。

相关推荐
仙人掌_lz1 小时前
Qwen-3 微调实战：用 Python 和 Unsloth 打造专属 AI 模型
人工智能·python·ai·lora·llm·微调·qwen3
m0_678693332 小时前
深度学习笔记26-天气预测(Tensorflow)
笔记·深度学习·tensorflow
美林数据Tempodata2 小时前
大模型驱动数据分析革新：美林数据智能问数解决方案破局传统 BI 痛点
数据库·人工智能·数据分析·大模型·智能问数
硅谷秋水2 小时前
NORA：一个用于具身任务的小型开源通才视觉-语言-动作模型
人工智能·深度学习·机器学习·计算机视觉·语言模型·机器人
正儿八经的数字经3 小时前
人工智能100问☞第46问：AI是如何“学习”的？
人工智能·学习
飞哥数智坊3 小时前
别卷提示词了！像带新人一样“带”AI，产出效率翻倍
人工智能
扫地的小何尚3 小时前
全新NVIDIA Llama Nemotron Nano视觉语言模型在OCR基准测试中准确率夺冠
c++·人工智能·语言模型·机器人·ocr·llama·gpu
xiaohanbao093 小时前
day54 python对抗生成网络
网络·python·深度学习·学习
m0_575470884 小时前
n8n实战：自动化生成AI日报并发布
人工智能·ai·自动化·ai自动写作
热门推荐
01Coze扣子平台完整体验和实践（附国内和国际版对比）02【图像处理与机器视觉】XJTU期末考点03KGG转MP3工具|非KGM文件|解密音频04零代码入门 | Coze——让大模型接入自己的数据库05海康Visionmaster-常见问题排查方法-启动阶段06YOLOv8入门 | 重要性能衡量指标、训练结果评价及分析及影响mAP的因素【发论文关注的指标】07扣子（coze）实战|我用扣子搭建了一个自动分析小红薯笔记内容的AI应用|详细步骤拆解08从零安装 LLaMA-Factory 微调 Qwen 大模型成功及所有的坑09【SpeedAI科研小助手】2分钟极速解决知网维普重复率、AIGC率过高，一键全文降！文件格式不变，公式都保留的！10DeepSeek各版本说明与优缺点分析