一、项目背景与数据集

MNIST 是计算机视觉领域的“Hello World”数据集，包含：

• 训练集：60000 张 28×28 灰度手写数字图片，标签 0~9
• 测试集：10000 张 28×28 灰度手写数字图片，标签 0~9

我们的目标：训练一个卷积神经网络，让它能准确识别图片中的手写数字。

二、环境准备与数据加载

首先导入必要的库，并使用 torchvision 轻松下载和加载 MNIST 数据：

import torch
from torch import nn
from torch.utils.data import DataLoader
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(),
)

ToTensor() 会将图片像素值从 [0,255] 缩放到 [0.0, 1.0]，并转换为张量格式。

三、可视化数据样本

随机查看前 9 张图片及其标签，感受数据形态：

from matplotlib import pyplot as plt

figure = plt.figure()
for i in range(9):
    img, label = training_data[i]
    figure.add_subplot(3, 3, i+1)
    plt.title(label)
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

squeeze() 用于去掉通道维度（1×28×28 → 28×28），方便 imshow 显示。

四、创建数据加载器（DataLoader）

将数据按批次打包，每批 64 张，便于梯度下降训练：

train_dataloader = DataLoader(training_data, batch_size=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}")   # [64, 1, 28, 28]
    print(f"Shape of y: {y.shape} {y.dtype}")    # [64] int64
    break

• N：批次大小（64）
• C：通道数（灰度图为 1）
• H, W：高和宽（28×28）

五、定义 CNN 网络结构

我们构建一个三层卷积网络，结构如下：

在动手写代码之前，先补充一个贯穿本节的关键概念——池化层（Pooling Layer）。

池化层通常紧跟在卷积层之后，作用是对特征图进行下采样（降维），它本身没有需要学习的参数。它的核心价值有两点：

1. 降低计算量：把特征图的宽高缩小（例如 28×28 → 14×14），后续层的参数量和计算量随之大幅减少。
2. 提取主要特征、增强鲁棒性：在局部区域内只保留最显著的信息，让网络对微小的平移、旋转不那么敏感，从而缓解过拟合。

最常用的是最大池化（MaxPooling），它会取每个小窗口内的最大值作为输出。例如 MaxPool2d(kernel_size=2) 会把一个 2×2 的区域压缩成 1 个值，因此特征图的宽高各缩小一半：

输入 2×2 区域（kernel_size=2 的窗口）：
[1, 3]
[5, 2]
→ 最大池化输出：5

在本节的网络里，conv1 和 conv2 末尾各接了一个 MaxPool2d(2)，所以特征图尺寸经历了两次减半：28×28 → 14×14 → 7×7。这也是后面全连接层输入维度是 64 * 7 * 7 的原因。

层	操作	输入→输出尺寸
conv1	Conv2d(1→16, 5×5, pad=2) + ReLU + MaxPool(2×2)	28×28→14×14
conv2	Conv2d(16→32, 5×5, pad=2) + ReLU + Conv2d(32→32) + ReLU + MaxPool(2×2)	14×14→7×7
conv3	Conv2d(32→64, 5×5, pad=2) + ReLU	7×7→7×7
fc	Linear(64×7×7 → 10)	输出 10 个类别的 logits

代码实现：

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=16, kernel_size=5, stride=1, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
            nn.Conv2d(32, 32, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(32, 64, 5, 1, 2),
            nn.ReLU(),
        )
        self.out = nn.Linear(64 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = x.view(x.size(0), -1)   # 展平
        output = self.out(x)
        return output

model = CNN().to(device)
print(model)

关键点：经过两次池化（2×2），28×28 → 14×14 → 7×7，最终展平为 64*7*7 维特征向量。

六、定义训练与测试函数

训练函数

def train(dataloader, model, loss_fn, optimizer):
    model.train()
    batch_idx = 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()

        if batch_idx % 100 == 0:
            print(f"loss {loss.item():>7f} [number:{batch_idx}]")
        batch_idx += 1

测试函数

def test(dataloader, model, loss_fn):
    size = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()
    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:.2f}%, Avg loss: {test_loss:.6f}")

七、选择损失函数与优化器

• 损失函数：CrossEntropyLoss —— 专为多分类设计，内部融合了 Softmax。
• 优化器：Adam —— 自适应学习率，比 SGD 收敛更快。

loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

八、启动训练主循环

完整遍历训练集 10 次（epochs=10），每次输出当前损失：

epochs = 10
for t in range(epochs):
    print(f"Epoch {t+1}\n-------------------")
    train(train_dataloader, model, loss_fn, optimizer)
print("Done!")
test(test_dataloader, model, loss_fn)

九、运行结果与预期

在 10 个 epoch 后，测试集准确率通常可达 99% 左右，损失降至 0.05 以下。部分输出示例：

Epoch10
-------------------
loss0.007136 [number:100]
loss0.000854 [number:200]
loss0.000482 [number:300]
loss0.005706 [number:400]
loss0.002470 [number:500]
loss0.001377 [number:600]
loss0.000040 [number:700]
loss0.000061 [number:800]
loss0.008632 [number:900]
Done!
Test result:
 Accuracy:99.17%,Avg loss:0.04120070406551051

十、总结与扩展

要点	说明
数据	MNIST 简单易用，适合新手
网络	卷积层提取局部特征，池化层降低维度
训练	Adam 优化器 + 交叉熵损失，收敛快
结果	10 轮后准确率 >98%，验证 CNN 有效性

可以尝试的改进方向：

• 增加 Dropout 层防止过拟合
• 使用 Batch Normalization 加速训练
• 调整卷积核数量或网络深度
• 尝试数据增强（旋转、平移等）