PyTorch-“识数”的AI

回顾

我们第一章，学会了如何创建一个简单的神经网络，并训练它。但是，这个模型能干的事情过于草率，只能预测一个简单的线性函数。

接下来，我们将扩大训练数据集，并创建一个能够识别数字的模型。

获取数据集

说明

在第一章，我们自己建立了一个测试数据集，并训练了神经网络。但是，这个数据集的规模很小。

但问题是自行创建高质量数据集成本很高（企业级项目尤其如此，涉及数据采集、清洗、标注等诸多环节），即使获取到现成数据集，其格式、尺寸、数值范围也可能与模型要求不一致——这时可以用 transforms 进行格式对齐和数值标准化。
所以，我们将使用一个现成的数据集-MNIST。

MNIST

MNIST是一个手写数字数据集，它有60000张训练图片和10000张测试图片。每张图片是 28 × 28 的0−9的手写数字图片组成。
可以发现，所有图片的像素都被限制在了28×28 像素上。那么，根据第一章的张量知识点，就很明显能分析出来模型是如何认识图片并训练的。

核心逻辑：像素灰度值 → 二维像素矩阵
所以假设3*3的图片，你直接从MNIST获取到的图片最终展示的结果是：

黑 白 黑
白 白 白
黑 黑 白

对应二维矩阵

$$\begin{bmatrix} 0 & 255 & 0 \\ 255 & 255 & 255 \\ 0 & 0 & 255 \end{bmatrix}$$

除了MNIST之外，还有许多数据集，比如CIFAR-10、CIFAR-100、Fashion-MNIST等等。

获取数据集

好的，我们先写代码，然后讲解里面都干了什么

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")

# ========== 1. 定义预处理管道 ==========
transform = transforms.Compose([
    transforms.ToTensor(),               # 操作1：图片转张量 + 像素归一化 [0,1]
    transforms.Normalize((0.1307,), (0.3081,)) # 操作2：标准化减均值除标准差
])

# ========== 2. 加载MNIST数据集 ==========
train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

# ========== 3. 构建数据加载器（分批读取） ==========
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

# ========== 4. 查看一批数据 ==========
imgs, labels = next(iter(train_loader))
print("批次图片张量shape：", imgs.shape)  # [64, 1, 28, 28]
print("标签shape：", labels.shape)        # [64]

在这段代码里，我们获取了MNIST数据集（第一次运行会下载数据集），并构建了数据加载器。
还使用了预处理管道对数据进行预处理。我们简单的讲一下每段代码都干了什么。

数据集加载

train_dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root="./data", train=False, download=True, transform=transform)

在这里，我们定义了数据路径，是否为训练模型，是否下载数据源，以及数据预处理管道。
是否下载数据源的默认值为False，为True的时候，如果数据集已经下载过了，那么就会使用本地数据集。
接下来，我们讲解一下预处理管道都干了什么。

预处理管道

transform = transforms.Compose([
    transforms.ToTensor(),               # 操作1：图片转张量 + 像素归一化 [0,1]
    transforms.Normalize((0.1307,), (0.3081,)) # 操作2：标准化减均值除方差
])

在这个预处理管道中，我们定义了两个动作：

• 归一化
• 标准化

归一化

归一化一共做了两个动作

1. 格式转换 + 维度调整
   硬盘里读取的 MNIST 是 PIL 灰度图，形状 (H=28, W=28)
   转换成 PyTorch 张量格式 [C=1, H=28, W=28]
   通道放最前面，是 CNN 网络规定的输入格式。

2. 像素缩放：0~255 → 0~1
   原始像素是整数 uint8，范围 0(黑) ~ 255(白)
   内部自动除以 255：pixel = pixel / 255.0
   输出浮点张量，值域固定 [0.0, 1.0]

这里说明一下C=1代表的是颜色的通道数量，因为我们只有黑白，所以通道数为1，现图片多是三通道或四通道（红绿蓝/红绿蓝+透明度），那么展现的张量就类似于

[
    [28*28的红色像素矩阵]
    [28*28的绿色像素矩阵]
    [28*28的蓝色像素矩阵]
]

那归一化的行为有什么意义呢？

首先，原始像素是整数 0~255，数值太大，不利于计算机计算，先缩小到 0~1 浮点。
这个格式转换是 PyTorch 中图像模型通用的输入约定（通道维度置于最前），后续卷积网络章节也会沿用此格式。

标准化

transforms.Normalize((0.1307,), (0.3081,))

标准化传入的参数，分别是 所有像素平均值(mean) 以及 所有像素标准差(std)

那么这两个数据传入后进行的计算是什么呢

$$out = \frac{input - mean}{std}$$

从归一化的结果中，我们发现像素值域是 [0.0, 1.0]，那么，我们按照步骤手动计算一下经过标准化后的值域：

$$0 → \frac{0 - 0.1307}{0.3081} \approx −0.424$$

$$1 → \frac{1 - 0.1307}{0.3081} \approx 2.821$$

所以经过标准化后这个数据集中的值域是 [-0.424, 2.821]。

好的，我们现在已经大致了解标准化里都干了些什么，那我们接下来说一下为什么要这么干。

首先，数据标准化的作用如下：

1. 数据中心化，梯度下降收敛更快 标准化后数据均值≈0，正负均匀分布，梯度更新更稳定，收敛速度提升几倍。
2. 避免梯度消失 / 爆炸 如果输入全为正数，反向传播时同一层所有参数的梯度符号相同，权重更新会出现 zigzag 震荡，收敛缓慢。标准化后数据零均值，正负各半，梯度更新方向更加灵活高效。；
3. 消除尺度差异，让学习率统一 如果特征数值区间差异巨大，同一个学习率无法兼顾所有参数。

那为什么经过上述的计算公式，就可以做到标准化呢，我们简单推导一下。

$$mean = \frac{\sum_{i=1}^{n} x_i}{N}$$

$$std = \sqrt{\frac{\sum_{i=1}^{n} (x_i - mean)^2}{N}}$$

其中N为像素总个数, x为像素值。

简单来说，mean就是 所有像素值相加并除以总数，像素标准差是 每个像素减去全局均值，然后全部平方求和（消除正负），最后再除以总像素数量，开平方根

标准化就是做了两个动作：

1. 减去均值 中心化，把分布中心平移到 0
2. 缩放 控制数据的波动幅度，让所有数据集的波动尺度统一

所以，标准化后数据均值为 0、标准差为 1，数据分布中心被平移、波动范围被统一缩放，有利于梯度下降。

数据加载器

好的 我们现在回到获取数据集的代码中，继续讲解数据加载器。

batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=0)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=0)

在这段代码中，两个数据集都被DataLoader包裹了，并且传入了指定的参数。我们来解读一下：

batch_size为 每次获取数据的数量 ，shuffle为 是否打乱数据集 ，num_workers为 数据加载的进程数 。

其中shuffle的意义何在呢？

1. 防止模型记住样本顺序，产生虚假拟合
   MNIST 数据集原始存放是按数字排序的：先一堆 0，再一堆 1，接着一堆 2……
   如果不打乱，训练时批次会连续出现大量相同数字。
   模型会偷懒，不学习图片笔画特征，只靠「当前第几个 batch」就能猜出标签，出现假的低损失，泛化能力极差。

2. 让每一批次数据分布均衡，梯度更稳定
   我们在训练时可能要经过多个batch，每个 batch 的数据分布本应该不同。
   一个 batch 里混合不同数字，梯度更新方向不会持续偏向某一类。
   不打乱的话，连续多个 batch 全是同一数字，梯度会单向偏移，损失震荡、收敛变慢。

3. 避免局部重复样本，提升泛化能力
   每次 epoch 顺序都不一样，模型每次看到的样本组合完全不同，不会重复学习一模一样的批次组合，减少过拟合。

所以经过上面的解释，我们可以很容易就推测出来，测试数据集是没有必要打乱的，因为测试的结果不会影响到模型，这样不打乱也减少了资源的消耗。

至此，数据集的获取，预处理，加载就就已经简单了解了。

MLP 多层感知机

多层感知机（MLP）是最简单的神经网络，它由输入层、隐藏层和输出层组成。

多层感知机的每一层是全连接（线性变换 + 激活函数）。如果不加激活函数，多层线性变换叠加后等价于单层，无法学习复杂特征。而引入 ReLU 等非线性激活函数后，网络具备了拟合任意复杂函数的能力。

1. 输入层（Input layer）

• 仅负责接收原始数据/特征，不做任何计算，神经元数量 = 数据的特征维度（如输入是 28×28 的图片，展平后神经元数为784）。
• 无激活函数，仅传递数据。

2.隐藏层（Hidden layer）

• MLP 的核心层，可设置1层或多层（多层时才称为深度多层感知机），层数/神经元数是关键超参数。
• 完成特征提取与非线性变换：对输入层的特征加权求和后，通过激活函数输出，实现特征的升维/抽象。
• 相邻隐藏层之间为全连接关系。

3.输出层（Output Layer）

• 输出模型的最终预测结果，神经元数量由任务类型决定：

✔ 二分类任务：1个神经元（输出0/1）；

✔ 多分类任务：神经元数 = 类别数（如MNIST手写数字分类，设10个神经元）；

✔ 回归任务：1个神经元（输出连续值）。

• 激活函数按需选择（分类用Softmax/Sigmoid，回归用Identity）。

ReLU

ReLU（Rectified Linear Unit）是一种常用的激活函数，它的计算公式为：

$$ReLU(z) = max(0,z)$$

• z>0：输出本身
• z<=0：输出0

ReLU 的作用是 引入非线性 ，如果不用任何激活函数，多层 Linear 等价于单层线性变换，网络只能学习直线分割，无法识别手写数字这种复杂特征。

同时还能缓解梯度消失，适配 Adam 优化，以及更小的资源消耗（只需要计算最大值）

ReLU 的一个潜在问题是 “dying ReLU”：如果某个神经元的权重被训练到对所有输入都输出 ≤ 0，该神经元梯度恒为 0，后续再也无法学习。适当的数据归一化、合理的权重初始化（如 Kaiming 初始化）和较小的学习率可以缓解这一问题。

代码展示

我们现在用代码实现一个可以用来训练MNIST数据集的 MLP 模型。

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

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = x.view(x.size(0), -1)               # 展平: (B,1,28,28) -> (B,784)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = F.relu(self.fc3(x))
        return self.fc4(x)

1. class MLP(nn.Module)
自定义神经网络类，继承 PyTorch 基础模型父类 nn.Module，所有网络都必须继承它，才能自动管理参数、梯度、GPU、保存模型等功能。

2. __init__()

self.fc1 = nn.Linear(28 \* 28, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, 64)
self.fc4 = nn.Linear(64, 10)

这是一个4 层全连接多层感知机（3 层隐藏层 + 输出层），专门用于 MNIST 手写数字分类任务（28×28 灰度图片，10 分类：数字 0~9）。

• fc1：输入维度 784 (28×28 像素图片展平) → 隐藏层 1：256 个神经元
• fc2：256 → 隐藏层 2：128 神经元
• fc3：128 → 隐藏层 3：64 神经元
• fc4：64 → 输出层：10 神经元（对应 0-9 十类手写数字）

self.dropout = nn.Dropout(0.3)

Dropout 层：在训练时，随机将 30% 神经元 输出置 0（也就是随机关掉30%的神经元） ，训练时降低神经元之间的耦合，抑制过拟合；测试阶段自动关闭，不会随机置零。

3. forward 前向传播函数

forward 是模型核心函数，必须实现，定义输入数据如何逐层计算得到输出。

def forward(self, x):
    x = x.view(x.size(0), -1)               # 展平: (B,1,28,28) -> (B,784)

输入 x 原始形状：(batch_size, 通道数=1, 28, 28)

• x.size(0)：取 batch 大小，保留第一维不变
• -1：自动计算剩余所有像素，合并为一维向量

效果：四维图片张量 → 二维特征张量 (batch_size, 784)，全连接层只能接收二维向量输入

x = F.relu(self.fc1(x))
x = self.dropout(x)

• self.fc1(x)：数据经过第一层全连接线性变换
• F.relu()：激活函数，给网络引入非线性。如果没有 ReLU，多层全连接等价于单层，无法学习复杂特征
• self.dropout(x)：激活后过 Dropout，随机失活 30% 神经元

x = F.relu(self.fc2(x))
x = self.dropout(x)
x = F.relu(self.fc3(x))
return self.fc4(x)

后续层与第一层类似，不再赘述。

• self.fc4(x) 输出层：不使用激活函数，直接输出 10 个分数（logits）

代码实现

现在开始实现训练代码

# -*- coding: utf-8 -*-
"""完整的训练流程: MLP 在 MNIST 上手写数字识别"""
import sys, os, torch, torch.nn as nn, torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

if sys.platform == "win32": sys.stdout.reconfigure(encoding="utf-8", errors="replace")
plt.rcParams["font.sans-serif"] = ["Microsoft YaHei", "SimHei"]
plt.rcParams["axes.unicode_minus"] = False
BASE = os.path.dirname(os.path.abspath(__file__))
OUT = os.path.join(BASE, "output")
os.makedirs(OUT, exist_ok=True)

BATCH, EPOCHS, LR = 64, 5, 0.001
DEV = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"设备: {DEV}")

print("\n加载 MNIST 数据集...")
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])
train_ds = datasets.MNIST(os.path.join(BASE, "..", "data"), train=True, download=True, transform=transform)
test_ds = datasets.MNIST(os.path.join(BASE, "..", "data"), train=False, download=True, transform=transform)
train_loader = DataLoader(train_ds, batch_size=BATCH, shuffle=True)
test_loader = DataLoader(test_ds, batch_size=BATCH, shuffle=False)
print(f"训练集: {len(train_ds):,} 张, 测试集: {len(test_ds):,} 张")

class MLP(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(28 * 28, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 64)
        self.fc4 = nn.Linear(64, 10)
        self.dropout = nn.Dropout(0.3)

    def forward(self, x):
        x = x.view(x.size(0), -1)     
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = F.relu(self.fc3(x))
        return self.fc4(x)

model = MLP().to(DEV)
print(f"模型参数: {sum(p.numel() for p in model.parameters()):,}")

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)

def train_epoch(model, loader, crit, opt):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEV), labels.to(DEV)
        opt.zero_grad()
        loss = crit(model(imgs), labels)
        loss.backward()
        opt.step()
        total_loss += loss.item()
        _, preds = model(imgs).max(1)
        total += labels.size(0)
        correct += preds.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total

@torch.no_grad()
def evaluate(model, loader, crit):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEV), labels.to(DEV)
        outputs = model(imgs)
        total_loss += crit(outputs, labels).item()
        _, preds = outputs.max(1)
        total += labels.size(0)
        correct += preds.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total

print(f"\n{'='*55}")
print(f"训练 {EPOCHS} 个 epoch...")
print(f"{'='*55}")

best_acc = 0
for epoch in range(1, EPOCHS + 1):
    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer)
    test_loss, test_acc = evaluate(model, test_loader, criterion)
    print(f"Epoch {epoch}: 训练loss={train_loss:.4f}, 训练acc={train_acc:.1f}%, 测试loss={test_loss:.4f}, 测试acc={test_acc:.1f}%")
    if test_acc > best_acc:
        best_acc = test_acc
        torch.save(model.state_dict(), os.path.join(OUT, "mnist_mlp_best.pth"))
        print(f"  -> 最佳模型已保存 ({best_acc:.1f}%)")

print(f"\n最佳准确率: {best_acc:.1f}%")

# 可视化预测
model.load_state_dict(torch.load(os.path.join(OUT, "mnist_mlp_best.pth"), weights_only=True))
model.eval()
test_imgs, test_labels = next(iter(test_loader))
test_imgs = test_imgs[:16].to(DEV)
with torch.no_grad():
    _, preds = model(test_imgs).max(1)

fig, axes = plt.subplots(4, 4, figsize=(8, 8))
for i, ax in enumerate(axes.flat):
    img = test_imgs[i].cpu().squeeze()
    ax.imshow(img, cmap="gray")
    color = "green" if preds[i] == test_labels[i] else "red"
    ax.set_title(f"预测:{preds[i].item()} 真实:{test_labels[i].item()}", color=color, fontsize=9)
    ax.axis("off")
plt.suptitle(f"MNIST 预测结果-MLP (测试准确率: {best_acc:.1f}%)", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.savefig(os.path.join(OUT, "mnist_mlp_predictions.png"), dpi=150, bbox_inches="tight")
plt.show()

CrossEntropyLoss + Adam

此次使用的损失函数为 CrossEntropyLoss ，优化器为 Adam 。

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR)

损失函数CrossEntropyLoss

CrossEntropyLoss 专门用于多分类任务，衡量模型预测分布和真实标签之间的差距，训练时用来计算 loss（损失值），反向传播更新权重。

CrossEntropyLoss 内有两步操作：

1. LogSoftmax：对模型原始输出 logits 做归一化 + 对数
2. NLLLoss（负对数似然）：根据真实标签取出对应概率的负数，作为损失

其中LogSoftmax计算公式为

$$p_{i} = \frac{e^{z_{i}}}{ {\textstyle \sum_{k=0}^{C-1}} e^{z_{k}}}$$

其中i代表计算的结果下标，C代表最后一层输出神经元数量

简单来说就是将模型跑出来的结果归一化后，获取正确标签对应的数值作为损失，最终 loss 数值就越小，代表模型预测越准。

优化器Adam

优化器负责反向传播后更新模型所有权重参数，核心目标：沿着梯度下降方向最小化 CrossEntropyLoss。
Adam 优化器融合动量与自适应学习率，依据历史梯度自动为模型每个权重调整更新幅度，以此降低损失。

optimizer = torch.optim.Adam(model.parameters(), lr=LR)

Adam 的核心：为网络里每一个参数单独维护两套历史梯度缓存，自动控制参数更新步长，解决 SGD 固定学习率缺陷。

给每一个参数 w，单独创建两个变量，初始值全 0：

• m：一阶动量（梯度滑动平均值，平滑震荡）
• v：二阶动量（梯度平方滑动平均值，自适应学习率）

Adam先给每个参数的m和v赋初始值（初始值为0），迭代更新参数 ，随后对动量进行偏置校正，防止初始值为0时的偏差。最后再通过动量更新权重。

原生 Adam 的weight_decay会和梯度 g 耦合，正则化效果差，容易在 MNIST 训练后期过拟合，后续会介绍 AdamW 解决这个问题。

训练模型

训练代码如下：

def train_epoch(model, loader, crit, opt):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEV), labels.to(DEV)
        opt.zero_grad()
        loss = crit(model(imgs), labels)
        loss.backward()
        opt.step()
        total_loss += loss.item()
        _, preds = model(imgs).max(1)
        total += labels.size(0)
        correct += preds.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total

循环内第一行表示通知模型为训练模式，这样DropOut层才会起作用。

第二层循环就是通过遍历训练数据集进行训练。内部能看见也是标准的五步训练循环，只是第二步和第三步合并在一起了

评估模型

评估代码如下：

@torch.no_grad()
def evaluate(model, loader, crit):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    for imgs, labels in loader:
        imgs, labels = imgs.to(DEV), labels.to(DEV)
        outputs = model(imgs)
        total_loss += crit(outputs, labels).item()
        _, preds = outputs.max(1)
        total += labels.size(0)
        correct += preds.eq(labels).sum().item()
    return total_loss / len(loader), 100. * correct / total

这里面的@torch.no_grad()可以在评估测试时不计算梯度、不存储计算图，大幅节省显存、提速；如果不加，测试阶段依旧缓存梯度，浪费 GPU 内存。

模型保存与加载逻辑

torch提供了模型保存与加载的方法，

## 保存
torch.save(model.state_dict(), os.path.join(OUT, "mnist_mlp_best.pth"))
## 加载
model.load_state_dict(torch.load(os.path.join(OUT, "mnist_mlp_best.pth"), weights_only=True))

在这个章节中save方法只保存了权重参数，并没有保存模型的网络结构，这样的好处是文件体积会更小。如果想完整保存模型 + 权重，可以使用torch.save(model, path)，这样文件体积更大，且跨环境加载时可能因代码不一致而失败。

在加载阶段，调用torch.load方法可以读取保存的模型文件，weights\_only=True是强制限制加载逻辑：只允许读取张量权重，拒绝加载任何可执行代码、自定义对象。

model.load_state_dict (weight_dict)作用就是将读取到的权重参数赋给模型，返回的就是模型本身。

那么至此，我们就完成了整个MNIST的训练过程。并同时学习到了：

• 数据集的加载与处理
• MLP的搭建与训练
• 模型保存与加载
• 新的损失函数与优化器

接下来我们可以尝试学习卷积网络进行更高级有效的模型训练。