PyTorch 手写数字识别学习笔记

最近学习了 PyTorch 框架的基础知识，并尝试使用 MNIST 手写数字数据集完成一个简单的图像识别任务。整个过程涉及神经网络结构、激活函数、梯度下降、反向传播、损失函数、数据加载和模型测试等内容。

刚开始接触深度学习时，神经网络的结构图、公式和代码往往看起来比较复杂。但把它们拆开之后，会发现模型的核心流程其实很清楚：输入数据经过一层层计算，得到预测结果；然后通过损失函数衡量预测误差，再利用反向传播计算梯度，最后更新模型参数。模型不断重复这个过程，就能够逐渐学会识别数据中的规律。

一、MNIST 手写数字数据集

MNIST 是深度学习入门中非常经典的数据集，主要用于手写数字识别。数据集包含大量手写数字图片，每张图片的大小是 28×28 像素，数字范围从 0 到 9。

因为每张图片都是灰度图，所以每个像素点通常用一个数值表示。将图片输入模型之前，需要先把它转换成张量。PyTorch 中可以使用 torchvision 提供的数据集接口：

复制代码

from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.MNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.MNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

这里的 training_data 是训练集，test_data 是测试集。训练集用于让模型学习，测试集则用于检验模型是否真正具有识别能力。

一张图片的尺寸是 28×28，因此总共有：

复制代码

28 × 28 = 784

个像素。如果使用全连接神经网络，就需要将图片展开成长度为 784 的向量。

二、DataLoader 的作用

数据集准备好之后，还需要使用 DataLoader 按批次读取数据：

复制代码

from torch.utils.data import DataLoader

batch_size = 64

train_dataloader = DataLoader(
    training_data,
    batch_size=batch_size
)

test_dataloader = DataLoader(
    test_data,
    batch_size=batch_size
)

这里的 batch_size=64 表示每次向模型输入 64 张图片，而不是一次性将所有图片都放进去。

使用 batch 训练有几个好处。第一，可以降低内存占用。第二，计算速度更快。第三，每个 batch 产生的梯度存在一定随机性，有助于模型跳出不理想的参数状态。

从 DataLoader 中读取数据时，通常会得到输入图片和对应标签：

复制代码

for X, y in train_dataloader:
    print(X.shape)
    print(y.shape)
    break

输入图片的形状一般类似于：

复制代码

[64, 1, 28, 28]

其中 64 是 batch 大小，1 表示灰度通道，28 和 28 表示图片的高和宽。标签的形状通常是：

复制代码

[64]

每个标签对应一张图片的真实数字。

三、神经网络的基本结构

一个简单的全连接网络通常包括输入层、隐藏层和输出层。

对于 MNIST 数据集来说，网络可以设计成：

复制代码

输入层：784 个特征
第一隐藏层：128 个神经元
第二隐藏层：256 个神经元
输出层：10 个神经元

最后的 10 个输出神经元分别对应数字 0 到 9。模型输出的是每个类别的分数，分数最高的类别就是模型最终的预测结果。

在 PyTorch 中，通常通过继承 nn.Module 来定义神经网络：

复制代码

import torch
from torch import nn

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.output = 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.output(x)
        return x

nn.Flatten() 会把输入图片从 [batch_size, 1, 28, 28] 展开成 [batch_size, 784]。

nn.Linear 表示全连接层。第一层将 784 个像素特征转换成 128 个特征，第二层将 128 个特征转换成 256 个特征，最后一层输出 10 个类别分数。

四、为什么需要激活函数？

如果网络中每一层都只有线性计算，那么无论堆叠多少层，最终仍然只是一个线性变换。这样的模型表达能力有限，难以学习复杂的图像特征。

激活函数的作用就是增加非线性能力，让网络能够拟合更加复杂的函数。

常见的激活函数包括 Sigmoid、Tanh、ReLU 和 Leaky ReLU。

1. Sigmoid

Sigmoid 函数的公式为：

复制代码

σ(x) = 1 / (1 + e^(-x))

它的输出范围是 0 到 1，曲线呈现出 S 形。在输入值很大或很小时，函数会逐渐饱和，导数接近于 0。

在反向传播过程中，如果多层网络都使用 Sigmoid，那么许多接近 0 的导数会连续相乘，最终导致梯度越来越小，这就是梯度消失。

2. Tanh

Tanh 函数的输出范围是 -1 到 1，并且以 0 为中心。相比 Sigmoid，它在某些场景下更加适合处理零中心数据。

不过，Tanh 在输入值较大时也容易饱和，因此仍然存在梯度消失的问题。

3. ReLU

ReLU 的公式非常简单：

复制代码

f(x) = max(0, x)

当输入大于 0 时，输出等于输入；当输入小于等于 0 时，输出为 0。

ReLU 在正数区域的导数是 1，可以有效缓解梯度消失问题，同时计算速度很快，因此被广泛应用在深度神经网络中。

4. Leaky ReLU

ReLU 在负数区域的输出始终是 0，可能导致某些神经元长期不更新，这种现象通常被称为"神经元死亡"。

Leaky ReLU 在负数区域保留一个很小的斜率：

复制代码

f(x) = x，x > 0
f(x) = αx，x ≤ 0

这样即使输入是负数，也能够保留少量梯度。

五、梯度消失与梯度爆炸

神经网络训练依赖反向传播，而反向传播的核心是链式法则。一个参数的梯度，通常需要乘以多层权重和激活函数导数。

如果连续相乘的数大部分小于 1，最终的梯度就会越来越小，这就是梯度消失。梯度消失会导致前面层的参数几乎不再更新，网络训练速度变慢，深层特征也很难学习。

如果连续相乘的数大部分大于 1，梯度则可能越来越大，形成梯度爆炸。梯度爆炸会使参数迅速变得非常大，损失值甚至可能变成 NaN。

常见的解决方法包括：

1.使用 ReLU 等激活函数；

2.合理初始化网络参数；

3.使用 BatchNorm；

4.选择合适的学习率；

5.使用梯度裁剪；

6.采用残差连接；

7.使用更加稳定的网络结构。

六、损失函数

模型完成前向传播后，会输出 10 个类别的预测分数。接下来需要判断预测结果与真实标签之间的差距，这就需要使用损失函数。

对于 MNIST 这种多分类任务，可以使用交叉熵损失：

复制代码

loss_fn = nn.CrossEntropyLoss()

损失越小，表示模型预测结果越接近真实标签；损失越大，表示预测结果越不准确。

使用 CrossEntropyLoss 时，模型最后一层通常不需要手动添加 Softmax，因为交叉熵损失函数内部已经完成了相关计算。

例如，模型可能输出：

复制代码

[1.2, -0.5, 4.8, 0.7, ...]

其中下标为 2 的位置分数最大，那么模型就会预测这张图片是数字 2。

七、梯度下降与优化器

损失函数只能衡量模型表现，真正负责修改参数的是优化器。

最基本的梯度下降公式为：

复制代码

θ = θ - η∇L(θ)

其中，θ 表示模型参数，η 表示学习率，∇L(θ) 表示损失函数关于参数的梯度。

学习率决定参数每次更新的幅度。如果学习率太大，模型可能跳过最优点，甚至导致损失越来越大；如果学习率太小，模型虽然能够慢慢收敛，但训练过程会比较慢。

可以使用 SGD 优化器：

复制代码

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=1e-3
)

也可以使用 Adam：

复制代码

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3
)

Adam 会根据历史梯度自动调整参数的更新步长，通常收敛速度比较快，适合用来进行基础实验。

八、训练循环

训练循环是模型学习的核心。一般包括以下几个步骤：

读取一个 batch；
进行前向传播；
计算损失；
清空旧梯度；
反向传播；
更新参数。

代码如下：

复制代码

def train(dataloader, model, loss_fn, optimizer):
    model.train()

    for X, y in dataloader:
        pred = model(X)
        loss = loss_fn(pred, y)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

model.train() 表示模型进入训练模式。

optimizer.zero_grad() 用来清空之前保存的梯度。PyTorch 默认会累积梯度，如果不清零，前面 batch 的梯度就会影响后面的参数更新。

loss.backward() 自动计算所有参数的梯度。

optimizer.step() 根据当前梯度更新模型参数。

九、模型测试

模型训练完成后，还需要使用测试集检查模型的实际表现：

复制代码

def test(dataloader, model, loss_fn):
    model.eval()

    test_loss = 0
    correct = 0

    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)

            test_loss += loss_fn(
                pred, y
            ).item()

            correct += (
                pred.argmax(1) == y
            ).type(torch.float).sum().item()

    test_loss /= len(dataloader)
    accuracy = correct / len(dataloader.dataset)

    print("测试结果：")
    print(f"准确率：{accuracy * 100:.2f}%")
    print(f"平均损失：{test_loss:.4f}")

测试时需要调用 model.eval()，表示模型进入评估模式。

torch.no_grad() 用来关闭梯度计算。测试阶段不需要修改模型参数，因此关闭梯度可以节省内存和计算资源。

pred.argmax(1) 用来找到每一行预测分数中的最大值位置，这个位置就是模型预测的类别。

十、完整训练过程

将模型、损失函数和优化器准备好之后，就可以开始训练：

复制代码

model = NeuralNetwork()

loss_fn = nn.CrossEntropyLoss()

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3
)

epochs = 10

for t in range(epochs):
    print(f"第 {t + 1} 轮训练")

    train(
        train_dataloader,
        model,
        loss_fn,
        optimizer
    )

    test(
        test_dataloader,
        model,
        loss_fn
    )

print("训练完成")

epochs 表示模型将整个训练集重复学习多少次。

训练轮数太少，模型可能还没有充分学习；训练轮数太多，则可能出现过拟合。判断训练效果时，不能只看训练集准确率，还需要观察测试集准确率。

如果训练集准确率不断提高，但测试集准确率不再提升，说明模型可能已经开始记忆训练数据，而不是学习真正具有泛化能力的特征。

十一、使用 GPU 加速

如果电脑中有可用的 CUDA GPU，可以让模型和数据运行在 GPU 上：

复制代码

device = (
    "cuda"
    if torch.cuda.is_available()
    else "cpu"
)

model = NeuralNetwork().to(device)

训练时，还需要将输入数据和标签移动到相同设备：

复制代码

X = X.to(device)
y = y.to(device)

模型和数据必须处于同一个设备，否则会出现设备不匹配的问题。

对于 MNIST 这种数据量较小的任务，CPU 也可以完成训练。但当模型变得更加复杂，或者数据量明显增加时，GPU 的优势会更加明显。

十二、从全连接网络到卷积神经网络

全连接网络可以完成手写数字识别，但它没有充分利用图片的空间结构。

在图片中，相邻像素通常具有较强的关系。例如，数字的边缘、线条和轮廓都由局部像素组合而成。全连接网络会把所有像素展开成一个长向量，可能丢失一部分空间信息。

卷积神经网络则通过卷积核提取局部特征。网络前面的卷积层可以学习边缘和线条，后面的卷积层可以逐渐组合出数字的局部结构，最终识别出完整数字。

一个简单的卷积网络可以写成：

复制代码

class CNN(nn.Module):
    def __init__(self):
        super().__init__()

        self.network = nn.Sequential(
            nn.Conv2d(1, 32, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Conv2d(32, 64, 3),
            nn.ReLU(),
            nn.MaxPool2d(2),

            nn.Flatten(),
            nn.Linear(64 * 5 * 5, 10)
        )

    def forward(self, x):
        return self.network(x)

虽然卷积神经网络的结构更加复杂，但训练方式基本不变，仍然是前向传播、计算损失、反向传播和参数更新。

十三、学习过程中容易出现的问题

1. 忘记清零梯度

如果不调用：

复制代码

optimizer.zero_grad()

模型的梯度会不断累积，导致参数更新不正常。

2. 忘记切换模型模式

训练时应该使用：

复制代码

model.train()

测试时应该使用：

复制代码

model.eval()

如果忘记切换，某些网络层可能会产生不正确的结果。

3. 输入维度错误

全连接层需要二维输入。如果图片仍然保持 [batch_size, 1, 28, 28] 的形状，直接输入线性层就会报错。

因此，需要先使用：

复制代码

nn.Flatten()

将图片展开为 784 维向量。

4. 输出类别数量错误

MNIST 有 10 个类别，因此最后一层应该输出 10 个数：

复制代码

nn.Linear(256, 10)

如果输出数量不正确，模型就无法与标签对应。

5. 学习率设置不合理

学习率过大时，损失可能不断波动；学习率过小时，训练速度会非常慢。遇到这种情况，可以尝试调整学习率，或者更换优化器。

十四、这次学习的收获

通过这次学习，我对深度学习中的几个基本概念有了更加清楚的认识。

神经网络的每一层都可以看作一种特征变换。输入图片经过展平和线性层后，逐渐形成更加抽象的特征。激活函数负责加入非线性能力，使模型可以学习复杂关系。损失函数负责衡量预测结果和真实标签之间的差异。反向传播负责计算参数梯度，优化器负责根据梯度更新参数。

整个过程可以概括为：

输入数据，前向传播，得到预测结果，计算损失，反向传播，更新参数，重复训练

模型并不是一开始就知道数字长什么样，而是在大量样本和不断更新参数的过程中，逐渐学习出不同数字的特征。

十五、总结

PyTorch 提供了比较完整的深度学习开发工具。通过 Dataset 和 DataLoader，可以方便地读取数据；通过 nn.Module，可以灵活地构建模型；通过自动微分机制，可以自动计算梯度；通过优化器，可以快速完成参数更新。

MNIST 手写数字识别虽然是一个基础任务，但它包含了深度学习训练的完整流程。从数据准备到模型测试，每一步都对应着实际的机器学习思想。

这次学习让我认识到，深度学习并不是简单地调用一个模型，而是需要理解数据如何进入网络、参数如何变化、损失如何下降，以及模型为什么能够逐渐提高准确率。