PyTorch从零手写CNN实现MNIST手写数字识别（完整实战·单篇完结）

1. 项目概述

手写数字识别是深度学习图像分类入门标杆任务，常用于深度学习课程实验、模型基线测试、入门课题验证。本文基于PyTorch从零搭建纯卷积神经网络，不使用任何预训练模型、迁移学习，原生实现MNIST手写数字十分类任务。完整覆盖数据集加载、数据标准化、CNN模型搭建、训练迭代、测试验证、指标可视化、结果分析全流程，代码可直接运行、结构清晰、便于二次修改与创新拓展。

2. 数据集介绍

MNIST 是开源手写数字标准数据集，为灰度单通道图像，分辨率统一28×28，无尺寸杂乱问题。数据集包含0~9共10个数字类别，整体样本分布均衡、噪声低、质量高。

数据集详细参数：

训练集：60000张图像
测试集：10000张图像
图像通道：1通道灰度图
任务类型：单标签多分类

3. 环境依赖

项目基于通用深度学习环境，适配Windows、Linux、Mac平台，支持CPU/GPU训练。

bash 复制代码

pip install torch torchvision matplotlib numpy -i https://pypi.tuna.tsinghua.edu.cn/simple

4. 数据集预处理与加载

针对MNIST数据集特性，采用官方均值与方差做标准化处理，贴合数据集分布，加速模型收敛。不做多余数据增强，保证基线模型公平性与简洁性。通过DataLoader实现分批加载、数据集打乱，提升训练稳定性。

python 复制代码

import torch
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
import numpy as np

# 数据集标准化参数（MNIST官方均值、方差）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载训练集、测试集
train_dataset = torchvision.datasets.MNIST(
    root="./mnist_data", train=True, download=True, transform=transform
)
test_dataset = torchvision.datasets.MNIST(
    root="./mnist_data", train=False, download=True, transform=transform
)

# 数据分批加载
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False, num_workers=0)

# 输出数据集信息
print(f"训练集样本数量：{len(train_dataset)}")
print(f"测试集样本数量：{len(test_dataset)}")

# 样本可视化
def show_mnist_sample():
    data_iter = iter(train_loader)
    images, labels = next(data_iter)
    img = torchvision.utils.make_grid(images[:8])
    img = img / 2 + 0.5
    np_img = img.numpy().transpose(1, 2, 0)
    plt.figure(figsize=(8, 4))
    plt.imshow(np_img)
    plt.title("MNIST手写数字样本")
    plt.axis("off")
    plt.show()

show_mnist_sample()

5. 自定义CNN网络模型

针对小尺寸灰度图像设计轻量化CNN结构，双层卷积逐级提取数字边缘、拐点、纹理特征，配合最大池化降维减参，Dropout抑制过拟合，最后通过全连接层完成10分类输出。结构简洁、参数量小、训练速度快，适合作为图像分类基线模型。

python 复制代码

import torch.nn as nn
import torch.nn.functional as F

class MNIST_CNN(nn.Module):
    def __init__(self):
        super(MNIST_CNN, self).__init__()
        # 第一层卷积：1通道输入，32卷积核，3*3卷积
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        # 第二层卷积：32通道输入，64卷积核
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        # 池化层
        self.pool = nn.MaxPool2d(2, 2)
        # 正则化
        self.dropout = nn.Dropout(0.25)
        # 全连接层
        self.fc1 = nn.Linear(64 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        # 卷积+激活+池化
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = self.dropout(x)
        # 特征展平
        x = x.view(-1, 64 * 7 * 7)
        # 全连接映射
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 设备自动适配
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MNIST_CNN().to(device)
print("模型初始化完成")
print(model)

6. 模型训练与验证

采用多分类通用交叉熵损失函数，Adam优化器更新参数。全程分开训练、验证流程，逐轮记录训练损失、测试损失、训练精度、测试精度，保证实验数据可追溯、可对比。

python 复制代码

import torch.optim as optim

# 超参数设置
EPOCHS = 20
LR = 0.001

# 损失函数、优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)

# 指标保存列表
train_loss_list = []
train_acc_list = []
test_loss_list = []
test_acc_list = []

# 迭代训练
for epoch in range(EPOCHS):
    # 训练模式
    model.train()
    train_loss = 0.0
    train_correct = 0
    total_train = 0

    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        optimizer.zero_grad()

        outputs = model(images)
        loss = criterion(outputs, labels)

        loss.backward()
        optimizer.step()

        train_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total_train += labels.size(0)
        train_correct += (predicted == labels).sum().item()

    avg_train_loss = train_loss / len(train_loader)
    train_acc = 100 * train_correct / total_train

    # 测试模式
    model.eval()
    test_loss = 0.0
    test_correct = 0
    total_test = 0

    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total_test += labels.size(0)
            test_correct += (predicted == labels).sum().item()

    avg_test_loss = test_loss / len(test_loader)
    test_acc = 100 * test_correct / total_test

    # 保存指标
    train_loss_list.append(avg_train_loss)
    train_acc_list.append(train_acc)
    test_loss_list.append(avg_test_loss)
    test_acc_list.append(test_acc)

    print(f"第{EPOCHS+1:2d}轮 | 训练损失:{avg_train_loss:.4f} | 训练精度:{train_acc:.2f}% | 测试损失:{avg_test_loss:.4f} | 测试精度:{test_acc:.2f}%")

# 保存模型权重
torch.save(model.state_dict(), "mnist_cnn_best.pth")
print("模型权重保存成功")

7. 实验结果可视化

绘制损失曲线与准确率曲线，直观展示模型收敛状态、拟合效果与训练稳定性。

python 复制代码

plt.rcParams["font.family"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False

plt.figure(figsize=(12, 5))
# 损失曲线
plt.subplot(1, 2, 1)
plt.plot(train_loss_list, label="训练损失", color="#e74c3c")
plt.plot(test_loss_list, label="测试损失", color="#3498db")
plt.title("损失变化曲线")
plt.xlabel("迭代轮数")
plt.ylabel("损失值")
plt.legend()

# 精度曲线
plt.subplot(1, 2, 2)
plt.plot(train_acc_list, label="训练精度", color="#e74c3c")
plt.plot(test_acc_list, label="测试精度", color="#3498db")
plt.title("准确率变化曲线")
plt.xlabel("迭代轮数")
plt.ylabel("准确率(%)")
plt.legend()

plt.tight_layout()
plt.show()

8. 创新优化拓展方向

增加多尺度卷积、残差连接，提升特征提取能力
引入学习率衰减、动态学习率策略，提升后期精度
添加混淆矩阵，定量分析各类别识别准确率与错分情况
嵌入注意力机制，强化关键特征区域权重
模型轻量化剪枝、量化，实现推理加速

9. 常见问题解决

数据集下载失败：切换网络或手动放置数据集至对应目录
显存溢出：调小batch_size，自动适配CPU训练
精度过低：增大迭代轮数、调整学习率、检查数据归一化参数
中文乱码：代码内置字体适配，直接运行即可正常显示