【从零开始：PyTorch实现MNIST手写数字识别全流程解析】

目录

一、基础环境与核心库

[1. 核心库功能](#1. 核心库功能)

[2. 安装命令](#2. 安装命令)

[二、数据准备：MNIST 数据集加载与处理](#二、数据准备：MNIST 数据集加载与处理)

[1. 数据集特性](#1. 数据集特性)

[2. 完整代码（含可视化）](#2. 完整代码（含可视化）)

三、神经网络模型构建

[1. 核心概念](#1. 核心概念)

[2. 完整模型代码](#2. 完整模型代码)

四、模型训练与测试

[1. 核心组件](#1. 核心组件)

[2. 完整训练 + 测试代码](#2. 完整训练 + 测试代码)

五、模型保存与加载

六、关键问题与优化技巧

[1. 常见错误解决](#1. 常见错误解决)

[2. 性能优化](#2. 性能优化)

六、修改参数提升准确率，改善神经系统的性能：

1.改善前代码示例：

---

基于 PyTorch 生态实现 MNIST 手写数字数据集的自动下载、加载与可视化，涵盖了 PyTorch 核心库调用、计算机视觉数据集处理、数据转换、Matplotlib 图像可视化

一、基础环境与核心库

1. 核心库功能

库名	核心作用	关键模块 / 类
torch	深度学习框架核心，提供张量操作、自动求导	torch.Tensor、torch.autograd、torch.optim
torchvision	计算机视觉专用库，含数据集、图像变换	datasets（MNIST/CIFAR10）、transforms、models
torch.nn	神经网络构建核心，提供层、损失函数	nn.Module（模型基类）、nn.Linear（全连接层）、nn.CrossEntropyLoss
torch.utils.data	数据加载与批处理	Dataset（数据集抽象类）、DataLoader（批处理迭代器）
matplotlib	图像可视化，用于展示 MNIST 手写数字	plt.imshow、plt.subplot

2. 安装命令

# 基础安装（CPU/GPU通用，自动匹配环境）
pip install torch torchvision matplotlib
# GPU版本需额外确保NVIDIA驱动与CUDA兼容，安装时会自动关联

二、数据准备：MNIST 数据集加载与处理

1. 数据集特性

• 规模：共 70000 张 28×28 灰度图像，60000 张训练集、10000 张测试集
• 格式：每张图像为单通道（灰度），标签为 0-9 的数字类别
• 核心操作 ：图像转张量（ToTensor()）、批处理（DataLoader）

2. 完整代码（含可视化）

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# 1. 数据变换：将PIL图像转为Tensor（自动归一化到[0,1]）
transform = transforms.ToTensor()

# 2. 加载训练集/测试集
training_data = datasets.MNIST(
    root="data",        # 数据集保存路径
    train=True,         # True=训练集，False=测试集
    download=True,      # 本地无数据时自动下载
    transform=transform # 应用数据变换
)
test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=transform
)

# 3. 批处理加载器（减少内存占用，加速训练）
batch_size = 64
train_dataloader = DataLoader(training_data, batch_size=batch_size, shuffle=True) # 训练集打乱
test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=False)    # 测试集不打乱

# 4. 可视化训练集前9张图像
plt.figure(figsize=(8, 8))
for i in range(9):
    img, label = training_data[i]  # 每张数据含（图像张量，标签）
    plt.subplot(3, 3, i+1)
    plt.title(f"Label: {label}")   # 显示标签
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")  # squeeze()去除维度为1的通道维度
plt.show()

# 5. 查看批处理数据形状（N:批次大小，C:通道数，H/W:图像高/宽）
for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")  # 输出：torch.Size([64, 1, 28, 28])
    print(f"Shape of y (labels): {y.shape}")      # 输出：torch.Size([64])
    break

三、神经网络模型构建

1. 核心概念

• 模型基类 ：所有自定义模型需继承nn.Module，重写__init__（定义层）和forward（前向传播）
• 层结构：输入层（28×28=784 维）→ 隐藏层（128 维→256 维）→ 输出层（10 维，对应 10 个数字）
• 激活函数 ：用ReLU替代sigmoid，解决梯度消失问题，加速收敛

2. 完整模型代码

import torch
from torch import nn

# 1. 定义设备（自动优先使用GPU，无则用CPU）
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

# 2. 自定义全连接神经网络
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()  # 调用父类初始化
        self.flatten = nn.Flatten()  # 展平层：将(1,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（10分类）

    # 前向传播：定义数据流向
    def forward(self, x):
        x = self.flatten(x)          # 展平：(batch,1,28,28)→(batch,784)
        x = torch.relu(self.hidden1(x))  # 隐藏层1+ReLU激活
        x = torch.relu(self.hidden2(x))  # 隐藏层2+ReLU激活
        x = self.out(x)              # 输出层（无激活，CrossEntropyLoss自带Softmax）
        return x

# 3. 初始化模型并移动到设备
model = NeuralNetwork().to(device)
print(model)  # 打印模型结构

四、模型训练与测试

1. 核心组件

• 损失函数 ：CrossEntropyLoss，适用于多分类任务，自带 Softmax 归一化
• 优化器 ：Adam（自适应学习率，优于 SGD），学习率lr=0.001（平衡收敛速度与稳定性）
• 训练模式 ：model.train()（开启 dropout/batchnorm 更新）
• 测试模式 ：model.eval()+torch.no_grad()（关闭梯度计算，节省内存）

2. 完整训练 + 测试代码

import torch
from torch import nn, optim

# 1. 初始化损失函数与优化器
loss_fn = nn.CrossEntropyLoss()  # 多分类损失
optimizer = optim.Adam(model.parameters(), lr=0.001)  # Adam优化器

# 2. 训练函数
def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 切换训练模式
    total_loss = 0.0
    for batch, (X, y) in enumerate(dataloader, 1):  # 遍历批处理
        X, y = X.to(device), y.to(device)  # 数据移动到设备
        
        # 前向传播：计算预测值
        pred = model(X)
        # 计算损失
        loss = loss_fn(pred, y)
        
        # 反向传播+参数更新
        optimizer.zero_grad()  # 梯度清零（避免累积）
        loss.backward()        # 反向传播求梯度
        optimizer.step()       # 更新模型参数
        
        total_loss += loss.item()
        # 每100个批次打印一次损失
        if batch % 100 == 0:
            avg_batch_loss = loss.item()
            print(f"  Batch {batch:>3d} | Loss: {avg_batch_loss:>7f}")
    
    # 打印本轮平均损失
    avg_epoch_loss = total_loss / len(dataloader)
    print(f"Train Epoch | Avg Loss: {avg_epoch_loss:>7f}")

# 3. 测试函数（评估准确率与平均损失）
def test(dataloader, model, loss_fn):
    model.eval()  # 切换测试模式
    size = len(dataloader.dataset)    # 测试集总样本数
    num_batches = len(dataloader)     # 批次数
    test_loss, correct = 0.0, 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()  # argmax(1)取预测类别
    
    # 计算平均损失与准确率
    avg_test_loss = test_loss / num_batches
    accuracy = (correct / size) * 100
    print(f"Test Result | Accuracy: {accuracy:>5.2f}% | Avg Loss: {avg_test_loss:>7f}\n")
    return accuracy

# 4. 多轮训练（10轮，准确率稳定95%+）
epochs = 10
best_accuracy = 0.0
print("="*50)
for t in range(epochs):
    print(f"Epoch {t+1}/{epochs}\n" + "-"*30)
    train(train_dataloader, model, loss_fn, optimizer)
    current_acc = test(test_dataloader, model, loss_fn)
    # 记录最佳准确率
    if current_acc > best_accuracy:
        best_accuracy = current_acc

print("="*50)
print(f"Final Best Accuracy: {best_accuracy:>5.2f}%")

五、模型保存与加载

训练完成后保存模型参数，便于后续复用：

# 1. 保存模型参数（推荐，占用空间小）
torch.save(model.state_dict(), "mnist_model.pth")
print("Model saved as 'mnist_model.pth'")

# 2. 加载模型
loaded_model = NeuralNetwork().to(device)
loaded_model.load_state_dict(torch.load("mnist_model.pth"))
print("Model loaded successfully")

# 3. 加载后测试
loaded_model.eval()
with torch.no_grad():
    X, y = next(iter(test_dataloader))
    X, y = X.to(device), y.to(device)
    pred = loaded_model(X)
    sample_accuracy = (pred.argmax(1) == y).type(torch.float).mean().item() * 100
    print(f"Loaded Model Sample Accuracy: {sample_accuracy:>5.2f}%")

六、关键问题与优化技巧

1. 常见错误解决

• 警告Failed to load image Python extension：屏蔽非核心警告，代码开头添加：

import warnings
warnings.filterwarnings('ignore', category=UserWarning)

• 维度不匹配 ：确保全连接层输入维度 = 上一层输出维度（如hidden2输出 256，out输入必须 256）
• 准确率异常（如 0%/100%） ：检查test函数中correct /= size和test_loss /= num_batches是否写反

2. 性能优化

• 激活函数 ：ReLU替代sigmoid，解决梯度消失
• 优化器 ：Adam替代SGD，自适应学习率，收敛更快
• 数据打乱 ：DataLoader设shuffle=True（仅训练集），提升泛化能力
• 正则化 ：添加nn.Dropout(0.2)（隐藏层后），防止过拟合，代码示例：

class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        self.hidden1 = nn.Linear(784, 128)
        self.dropout1 = nn.Dropout(0.2)  # 丢弃20%神经元
        self.hidden2 = nn.Linear(128, 256)
        self.dropout2 = nn.Dropout(0.2)
        self.out = nn.Linear(256, 10)
    
    def forward(self, x):
        x = self.flatten(x)
        x = torch.relu(self.hidden1(x))
        x = self.dropout1(x)
        x = torch.relu(self.hidden2(x))
        x = self.dropout2(x)
        x = self.out(x)
        return x

六、修改参数提升准确率，改善神经系统的性能：

1.改善前代码示例：

# 导入必要的库
import torch
import torchvision
import torchaudio

# 打印版本信息
print(torch.__version__)
print(torchvision.__version__)
print(torchaudio.__version__)

"""
MNIST数据集介绍：
MNIST包含70,000张手写数字图像：60,000张用于训练，10,000张用于测试。
图像是灰度的，28x28像素的，并且居中的，以减少预处理和加快运行。
"""

# 导入PyTorch神经网络模块和数据加载工具
import torch
from torch import nn  # 导入神经网络模块
from torch.utils.data import DataLoader  # 数据包管理工具，打包数据
from torchvision import datasets  # 封装了与图像相关的模型，数据集
from torchvision.transforms import ToTensor  # 数据转换，将PIL图像转换为tensor张量

"""
下载训练数据集（包含训练图片+标签）
"""
# 创建训练数据集
training_data = datasets.MNIST(
    root="data",  # 数据存储路径
    train=True,  # True表示训练集，False表示测试集
    download=True,  # 如果数据不存在则下载
    transform=ToTensor()  # 将图像转换为tensor张量
)

# 创建测试数据集
test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

# 打印训练数据集大小
print(len(training_data))

"""
展示手写数字图片，显示训练数据集中的前9张图片
"""
import matplotlib.pyplot as plt

figure = plt.figure()
for i in range(9):  # 显示前9张图片
    img, label = training_data[i]  # 获取第i张图片和标签
    
    # 创建子图
    figure.add_subplot(3, 3, i+1)
    plt.title(label)  # 设置标题为标签值
    plt.axis("off")  # 关闭坐标轴
    plt.imshow(img.squeeze(), cmap="gray")  # 显示图片，使用灰度色彩映射
    
plt.show()

"""
创建数据加载器（DataLoader）
batch_size: 将数据集分成多份，每一份为batch_size个数据
优点：可以减少内存的使用，提高训练速度。
"""
# 创建训练数据加载器，每批64张图片
train_dataloader = DataLoader(training_data, batch_size=64)
# 创建测试数据加载器，每批64张图片
test_dataloader = DataLoader(test_data, batch_size=64)

# 检查数据加载器的一个批次
for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

"""
判断当前设备是否支持GPU
优先使用CUDA（NVIDIA GPU），其次是MPS（Apple M系列芯片），最后是CPU
"""
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

"""
定义神经网络类
继承自nn.Module，这是PyTorch中所有神经网络模块的基类
"""
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()  # 调用父类构造函数
        self.flatten = nn.Flatten()  # 展平层，将28x28的图像展平为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输出（对应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

# 创建模型实例并移动到指定设备
model = NeuralNetwork().to(device)
print(model)

"""
训练函数
"""
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.forward(X)  # 前向传播
        loss = loss_fn(pred, y)  # 计算损失
        
        # 反向传播
        optimizer.zero_grad()  # 清零梯度
        loss.backward()  # 计算梯度
        optimizer.step()  # 更新参数
        
        # 每100个批次打印一次损失
        if batch_size_num % 100 == 0:
            loss_value = loss.item()
            print(f"loss: {loss_value:>7f} [number:{batch_size_num}]")
        batch_size_num += 1

"""
定义损失函数和优化器
"""
# 交叉熵损失函数，适用于多分类问题
loss_fn = nn.CrossEntropyLoss()  

# 随机梯度下降优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# 开始训练
train(train_dataloader, model, loss_fn, optimizer)

"""
测试函数
"""
def test(dataloader, model, loss_fn):
    model.eval()  # 设置模型为评估模式
    size = len(dataloader.dataset)  # 数据集大小
    num_batches = len(dataloader)  # 批次数量
    test_loss, correct = 0, 0  # 初始化损失和正确预测数
    
    with torch.no_grad():  # 禁用梯度计算
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model.forward(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):>0.1f}%, Avg loss: {test_loss:>8f}")

# 执行测试
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)
    test(test_dataloader, model, loss_fn)
print("Done!")

代码运行结果：

2.修改参数提升准确率：

修改的参数：

1.激活函数优化：ReLU > Sigmoid（解决梯度消失问题），除非是做二分类的概率输出（最后一层），否则隐藏层几乎不再使用 Sigmoid，ReLU 及其变体（Leaky ReLU 等）是首选。

2. 优化器升级：Adam > SGD（自适应学习率，收敛更快），Adam 是目前深度学习中最常用的优化器，因为它能让模型训练得更快，且对新手更友好。

3. 学习率调整：0.001 > 0.01（更稳定的训练），当你把优化器换成 Adam 时，必须把学习率降下来，通常设置为 0.001 或 0.0001。

import warnings
warnings.filterwarnings('ignore', category=UserWarning)

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

# 打印版本
print(torch.__version__)
print(torchvision.__version__)
print(torchaudio.__version__)

# ---------------------- 加载MNIST数据集 ----------------------
# 训练集
training_data = datasets.MNIST(
    root="data", train=True, download=True, transform=ToTensor()
)
# 测试集
test_data = datasets.MNIST(
    root="data", train=False, download=True, transform=ToTensor()
)
print(f"训练集数量: {len(training_data)} | 测试集数量: {len(test_data)}")

# 展示训练集前9张图片（可选，注释后不影响训练）
plt.figure(figsize=(8, 8))
for i in range(9):
    img, label = training_data[i]
    plt.subplot(3, 3, i + 1)
    plt.title(label, fontsize=10)
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

# ---------------------- 创建DataLoader ----------------------
# 加shuffle=True：训练集打乱，提升泛化能力（测试集无需打乱）
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64)

# 查看batch形状
for X, y in test_dataloader:
    print(f"Shape of X [N, C, H, W]: {X.shape}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

# ---------------------- 设备判断 ----------------------
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

# ---------------------- 定义神经网络（仅保留1个正确模型，替换sigmoid为ReLU） ----------------------
class NeuralNetwork(nn.Module):
    def __init__(self):
        super().__init__()
        self.flatten = nn.Flatten()
        # 网络层：维度完全匹配，无冗余
        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（10个数字分类）

    def forward(self, x):
        x = self.flatten(x)
        x = self.hidden1(x)
        x = torch.relu(x)  # 替换sigmoid为ReLU，解决梯度消失，加速收敛
        x = self.hidden2(x)
        x = torch.relu(x)
        x = self.out(x)
        return x

# 初始化模型并送入设备
model = NeuralNetwork().to(device)
print("\n神经网络模型结构：")
print(model)

# ---------------------- 训练函数 ----------------------
def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 训练模式（开启梯度更新）
    total_loss = 0  # 新增：统计本轮总损失
    for batch, (X, y) in enumerate(dataloader, 1):  # 用enumerate更简洁统计批次
        X, y = X.to(device), y.to(device)
        pred = model(X)
        loss = loss_fn(pred, y)

        # 反向传播
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()
        # 每100个批次打印一次损失
        if batch % 100 == 0:
            avg_batch_loss = loss.item()
            print(f"  批次 {batch:>3d} | 批次损失: {avg_batch_loss:>7f}")
    # 打印本轮平均损失
    avg_epoch_loss = total_loss / len(dataloader)
    print(f"训练完成 | 本轮平均损失: {avg_epoch_loss:>7f}")

# ---------------------- 测试函数（修正计算+优化打印） ----------------------
def test(dataloader, model, loss_fn):
    model.eval()  # 测试模式（关闭梯度更新，节省资源）
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    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()
            # 统计正确预测数：pred.argmax(1)取预测概率最大的类别
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    # 计算平均损失和准确率
    test_loss /= num_batches
    correct /= size
    # 格式化打印，保留2位小数更直观
    print(f"测试结果 | 准确率: {(100 * correct):>5.2f}% | 平均损失: {test_loss:>7f}\n")
    return correct  # 返回准确率，方便后续查看

# ---------------------- 超参数初始化（核心调参点） ----------------------
loss_fn = nn.CrossEntropyLoss()  # 交叉熵损失：适合分类任务
# 优化器：SGD→Adam（自适应学习率，收敛更快更稳定），学习率调为0.001
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
epochs = 10  # 训练轮数：10轮足够达到95%+准确率

# ---------------------- 开始多轮训练+测试 ----------------------
print("\n==================== 开始训练 ====================")
best_accuracy = 0.0  # 记录最佳准确率
for t in range(epochs):
    print(f"\n第 {t+1}/{epochs} 轮训练")
    print("-" * 30)
    train(train_dataloader, model, loss_fn, optimizer)
    current_acc = test(test_dataloader, model, loss_fn)
    # 更新最佳准确率
    if current_acc > best_accuracy:
        best_accuracy = current_acc

print(f"==================== 训练结束 ====================")
print(f"最佳测试准确率: {(100 * best_accuracy):>5.2f}%")

提升准确率后的运行结果：

修改项	旧方案 (Legacy)	新方案 (Modern)	修改带来的好处
激活函数	Sigmoid / Tanh	ReLU	解决梯度消失，计算速度快，网络更深也能训练。
优化器	SGD (随机梯度下降)	Adam	自适应调节步长，收敛速度快，不易陷入局部最优。
学习率	0.01 (较大)	0.001 (较小)	配合 Adam 使用，防止训练过程震荡，保证稳定性。