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

库名	核心作用	关键模块 / 类
`torch`	深度学习框架核心，提供张量操作、自动求导	`torch.Tensor`、`torch.autograd`、`torch.optim` 这三个模块构成了 PyTorch 深度学习框架的核心骨架。
`torchvision`	计算机视觉专用库，含数据集、图像变换	`datasets`（MNIST/CIFAR10）内置音频数据集、`transforms`音频变换、`models预训练音频模型`
`torch.nn`	神经网络构建核心，提供层、损失函数	`nn.Module`（模型基类）、`nn.Linear`（全连接层）、`nn.CrossEntropyLoss`
`torch.utils.data`	数据加载与批处理	`Dataset`（数据集抽象类）、`DataLoader`（批处理迭代器）
torchaudio	PyTorch 官方音频处理库	functional：底层音频处理函数，io：音频文件读写
`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，解决梯度消失，除非是做二分类的概率输出（最后一层），否则隐藏层几乎不再使用 Sigmoid，ReLU 及其变体（Leaky ReLU 等）是首选。
优化器：Adam替代SGD，自适应学习率，收敛更快，Adam 是目前深度学习中最常用的优化器，因为它能让模型训练得更快，且对新手更友好。
数据打乱：DataLoader设shuffle=True（仅训练集），提升泛化能力
学习率调整：0.001 > 0.01（更稳定的训练），当你把优化器换成 Adam 时，必须把学习率降下来，通常设置为 0.001 或 0.0001。
正则化：添加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