PyTorch 实现手写数字识别：全连接网络 + CNN 卷积网络（MNIST 数据集实战）

丨白色风车丨2026-06-13 9:33

一、前言

MNIST 手写数字数据集是深度学习入门经典数据集，包含 60000 张训练集、10000 张测试集灰度手写数字图片（0~9）。本文基于 PyTorch 分别实现全连接神经网络（FC） 和 卷积神经网络（CNN） 完成手写数字分类，完整包含数据集加载、数据可视化、模型搭建、训练、测试全流程，适合深度学习新手入门学习。

环境依赖

python

运行

复制代码 
torch
torchvision
torchaudio
matplotlib

二、完整项目结构与整体思路

  1. 导入相关库，加载并解析 MNIST 官方数据集
  2. 数据可视化，查看样本图片与标签
  3. 使用 DataLoader 批量封装数据，提升训练效率
  4. 自动适配 CPU/GPU/MPS 设备，加速运算
  5. 搭建全连接神经网络，完成训练与评估
  6. 搭建卷积神经网络 CNN，完成训练与评估
  7. 对比两种网络在手写数字任务上的效果

三、代码实现与逐行讲解

3.1 库导入与数据集加载

首先导入 PyTorch 核心库、视觉工具库、数据加载工具，并下载加载 MNIST 训练集和测试集。

python

运行

复制代码 
# 导入深度学习相关库
import torch
import torchvision
import torchaudio
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor
from matplotlib import pyplot as plt

# 加载训练数据集：60000张手写数字图片
training_data = datasets.MNIST(
    root="data",        # 数据集存放路径
    train=True,         # 标记为训练集
    download=True,      # 自动在线下载数据集
    transform=ToTensor()# 将PIL图像转为Tensor张量（模型可识别格式）
)

# 加载测试数据集：10000张手写数字图片
test_data = datasets.MNIST(
    root="data",
    train=False,        # 标记为测试集
    download=True,
    transform=ToTensor()
)

# 打印训练集样本总数
print(f"训练集样本数量：{len(training_data)}")

核心参数说明

  • root：数据集本地存储目录，首次运行会自动创建 data 文件夹；
  • train：True 加载训练集，False 加载测试集；
  • ToTensor()：将像素值归一化到 [0,1]，同时转换为 PyTorch 张量。

3.2 数据集可视化

通过 matplotlib 展示前 9 张手写数字图片，直观查看数据样式：

python

运行

复制代码 
# 创建画布，展示9张样本图片
figure = plt.figure(figsize=(6, 6))
for i in range(9):
    img, label = training_data[i]
    # 3行3列子图
    figure.add_subplot(3, 3, i + 1)
    plt.title(f"Label: {label}")
    plt.axis("off")  # 隐藏坐标轴
    # squeeze() 去除维度为1的通道，灰度图使用灰度配色
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

运行代码会弹出窗口，展示 9 张 28×28 手写数字灰度图，并标注对应真实标签。

3.3 数据加载器 DataLoader

数据集单张读取效率低，DataLoader 可以分批打包数据（batch），降低内存占用、加速训练：

python

运行

复制代码 
# 批次大小：每一批加载64张图片
batch_size = 64

# 训练集、测试集批量加载器
train_dataloader = DataLoader(training_data, batch_size=batch_size)
test_dataloader = DataLoader(test_data, batch_size=batch_size)

# 查看数据维度格式 [N, C, H, W]
for X, y in test_dataloader:
    print(f"数据形状 [批次, 通道, 高度, 宽度]：{X.shape}")
    print(f"标签形状 & 数据类型：{y.shape} {y.dtype}")
    break

输出解释

  • N：批次大小（64）；C：通道数（灰度图 = 1）；H/W：图片高宽（28×28）；
  • 标签 y 为一维张量，代表每张图片对应的数字类别（0~9）。

3.4 设备自动适配（CPU/GPU/MPS）

优先使用 NVIDIA CUDA 显卡，其次苹果 M 系列芯片 MPS，最后使用 CPU：

python

运行

复制代码 
# 自动判断可用设备
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"当前使用设备：{device}")

模型和数据都需要迁移到对应设备，才能使用硬件加速。

3.5 方案一：全连接神经网络（FC）实现分类

全连接网络（多层感知机 MLP）是最基础的神经网络，将 28×28 图片展平为一维向量，通过多层线性层提取特征。

3.5.1 定义全连接网络模型

python

运行

复制代码 
# 自定义全连接神经网络，继承nn.Module
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        # 展平层：将 1*28*28 图片转为一维向量 784
        self.flatten = nn.Flatten()
        # 第一层全连接：784 -> 128 神经元
        self.hidden1 = nn.Linear(28 * 28, 128)
        # 第二层全连接：128 -> 256 神经元
        self.hidden2 = nn.Linear(128, 256)
        # 输出层：256 -> 10 （对应0-9共10个分类）
        self.out = nn.Linear(256, 10)

    # 前向传播（数据流转逻辑，函数名固定为forward）
    def forward(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.sigmoid(x)  # 激活函数
        x = self.hidden2(x)
        x = torch.sigmoid(x)
        x = self.out(x)
        return x

# 初始化模型并迁移到对应设备
fc_model = NeuralNetwork().to(device)
print("全连接网络结构：")
print(fc_model)
3.5.2 训练、测试通用函数

封装训练函数和测试函数，复用代码：

python

运行

复制代码 
# 训练函数
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(f"批次:{batch_num}  损失值: {loss.item():.6f}")
        batch_num += 1

# 测试/评估函数
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()
            # argmax(1) 获取预测类别，统计预测正确数量
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    # 计算平均损失和准确率
    test_loss /= num_batches
    acc = (correct / size) * 100
    print(f"测试集准确率: {acc:.2f}%  平均损失: {test_loss:.6f}\n")
3.5.3 损失函数、优化器 + 启动训练

python

运行

复制代码 
# 损失函数：交叉熵（分类任务标准损失）
loss_fn = nn.CrossEntropyLoss()
# 优化器：Adam 自适应学习率优化器
optimizer = torch.optim.Adam(fc_model.parameters(), lr=0.005)

# 迭代训练10轮
epochs = 10
for t in range(epochs):
    print(f"========== 第 {t+1} 轮训练 ==========")
    train(train_dataloader, fc_model, loss_fn, optimizer)
    test(test_dataloader, fc_model, loss_fn)

全连接网络特点

  • 结构简单、易理解；
  • 缺点：丢失图像空间特征，对图片类任务精度有限。

3.6 方案二：CNN 卷积神经网络实现分类

卷积神经网络（CNN）专门针对图像任务设计，依靠卷积层、池化层提取图像局部特征、纹理、轮廓，分类效果远优于全连接网络。

3.6.1 定义 CNN 卷积模型

python

运行

复制代码 
# 自定义卷积神经网络
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积块1：卷积+激活+池化
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 最大池化，下采样
        )
        # 卷积块2：双层卷积+激活+池化
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 卷积块3：单层卷积+激活
        self.conv3 = nn.Sequential(
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU()
        )
        # 全连接输出层：64*7*7 -> 10分类
        self.out = nn.Linear(64 * 7 * 7, 10)

    # 前向传播
    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        # 展平：batch维度保留，其余维度合并
        x = x.view(x.size(0), -1)
        output = self.out(x)
        return output

# 初始化CNN模型并迁移设备
cnn_model = CNN().to(device)
print("CNN卷积网络结构：")
print(cnn_model)

CNN 层参数说明

  • Conv2d：二维卷积，in_channels 输入通道，out_channels 卷积核数量；
  • kernel_size=5：5×5 卷积核；padding=2 补零，保证卷积后尺寸不变；
  • MaxPool2d：池化下采样，缩小特征图尺寸、减少参数量。
3.6.2 CNN 模型训练与测试

复用上方已定义的 train()、test() 函数，仅替换模型、优化器即可：

python

运行

复制代码 
# 重新定义优化器（CNN模型参数）
optimizer = torch.optim.Adam(cnn_model.parameters(), lr=0.001)

# 单轮训练+测试演示，可自行增加轮数
print("========== CNN 模型训练开始 ==========")
train(train_dataloader, cnn_model, loss_fn, optimizer)
test(test_dataloader, cnn_model, loss_fn)

四、运行结果分析

  1. 全连接网络

    • 训练收敛较慢，最终测试准确率一般在 96%~98% 左右；
    • 无法利用图像空间信息，复杂手写数字识别效果一般。

  2. CNN 卷积网络

    • 收敛速度更快，单轮训练即可达到 98%+ 准确率；
    • 依靠卷积提取图像局部特征，是图像分类任务的主流方案。

五、知识点总结

  1. MNIST 数据集：入门级手写数字数据集，28×28 灰度图，10 分类任务；
  2. DataLoader：批量加载数据，深度学习标准数据读取方式；
  3. 设备适配：PyTorch 支持 CUDA、MPS、CPU 自动切换，充分利用硬件加速；
  4. 全连接网络：将图像转为一维向量，结构简单，不适合图像特征提取；
  5. CNN 卷积网络：卷积 + 池化提取图像空间特征，图像分类首选；
  6. 训练流程固定模板：前向传播 → 计算损失 → 梯度清零 → 反向传播 → 参数更新。

六、拓展优化方向

  1. 增加训练轮数 epochs，进一步提升模型精度；
  2. 加入 Dropout 层防止过拟合；
  3. 调整学习率、批次大小、网络层数做调参实验；
  4. 保存训练好的模型权重，实现离线预测单张图片。
相关推荐
DianSan_ERP1 小时前
架构师视角：电商大促高并发下的订单API限流与防漏单架构演进
java·运维·网络·安全·微服务·架构·自动化
code monkey.1 小时前
【Linux之旅】HTTP 协议解析：从请求格式到构建 Web 服务器
linux·服务器·网络·http
十正1 小时前
aiohttp.TCPConnector 连接池原理详解
网络·python·tcp·aiohttp
福建佰胜张工1 小时前
3HNA006722-001 O-RING：ABB 喷涂机器人流体系统核心密封件技术解析
网络·人工智能·机器人
vortex52 小时前
Linux 传统设计哲学：通过调用名区分行为的艺术
linux·运维·网络
志栋智能2 小时前
超自动化安全的实施路径：从单点场景到体系化建设
运维·网络·安全·自动化
network_tester2 小时前
TSN交换机研发测试怎么做？一套可落地的“信而泰仪器 + 康芯源服务”方案解读
网络·网络协议·tcp/ip·车载系统·汽车·信息与通信·信号处理
2401_868534782 小时前
高并发架构终极总结
网络
woohuwan2 小时前
网络变压器PCB布局布线全攻略：位置、差分对与接地隔离
网络
热门推荐
01HTTP 与 HTTPS 的区别：从原理到实战详解02《置身钉内》原文-可播放阅读03【AI】2026 年具身智能模型和世界模型总结042026 AI 编程工具终极实战指南：Cursor vs Claude Code vs Copilot，开发者该怎么选？05Claude Code、Codex、Cursor三分天下：2026年AI编程Agent生态全景剖析06AI科技热点日报 | 2026年6月1日072026年6月AI行业全景：从百模大战到Agent元年，这30天发生了什么？08GitHub 镜像站点09AI一周事件 · 2026-06-03 至 2026-06-09102026 年 AI 编程工具终极横评：Cursor vs Claude Code vs Copilot vs Windsurf