引言

在深度学习与计算机视觉中，卷积神经网络（CNN）是最基础、最核心的模型之一。而MNIST手写数字识别，则是每个入门者必经的“Hello World”项目。它不仅数据规模适中，还能直观展示CNN如何从像素中自动提取特征并完成分类。

本文将使用PyTorch，从零构建一个CNN模型，完整走通数据加载、模型定义、训练与测试的全流程。我们不仅会给出可直接运行的代码，还会详细解释每一层的维度变化、核心API的含义，并通过最终的98.54%准确率结果，验证CNN在图像任务中的强大能力。读完本文，你将彻底搞懂：

• 为什么CNN比全连接网络更适合图像？

• 卷积、池化、ReLU各自的作用是什么？

• 数据在CNN中是如何流转并最终输出分类结果的？

环境与依赖

• Python 3.8+

• PyTorch 1.10+

• torchvision

---

任务描述

我们使用经典的MNIST手写数字数据集，完成以下任务：

1. 自动下载并加载60000张训练图片、10000张测试图片；

2. 定义一个由3个卷积模块和1个全连接输出层组成的CNN；

3. 训练模型10个轮次（epoch），观察损失下降；

4. 在测试集上评估模型，输出准确率和平均损失。

实现步骤详解

1. 导入库并检查设备

import torch
from torch import nn
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

# 自动选择设备：CUDA（NVIDIA） > MPS（Apple） > CPU
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {device}")

运行结果（示例）：

Using device: cpu

2. 加载MNIST数据集

# 训练集：60000张，transform=ToTensor() 将PIL图像转为张量，并归一化到[0,1]
training_data = datasets.MNIST(
    root='data',
    train=True,
    download=True,
    transform=ToTensor(),
)

# 测试集：10000张
test_data = datasets.MNIST(
    root='data',
    train=False,
    download=True,
    transform=ToTensor(),
)

# 创建DataLoader，每批32张，训练时打乱顺序
train_dataloader = DataLoader(training_data, batch_size=32, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=32)

说明：ToTensor() 将原始像素值从0~255缩放到0~1，是神经网络训练前的标准预处理。

3. 定义CNN模型

我们使用 nn.Sequential 将卷积、激活、池化打包成模块，使代码更简洁。

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 模块1：提取底层特征（边缘、线条）
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 64, kernel_size=5, stride=1, padding=2),  # 输入1通道，输出64通道
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)   # 尺寸减半
        )
        # 模块2：提取中层特征（笔画组合）
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 32, 5, 1, 2),
            nn.ReLU(),
            nn.Conv2d(32, 16, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        # 模块3：提取高层特征（数字整体形态）
        self.conv3 = nn.Sequential(
            nn.Conv2d(16, 32, 5, 1, 2),
            nn.ReLU(),
        )
        # 全连接输出层
        self.out = nn.Sequential(
            nn.Linear(32 * 7 * 7, 10),  # 输入1568维，输出10类
            nn.ReLU()
        )

    def forward(self, x):
        x = self.conv1(x)   # [batch, 64, 14, 14]
        x = self.conv2(x)   # [batch, 16, 7, 7]
        x = self.conv3(x)   # [batch, 32, 7, 7]
        x = x.view(x.size(0), -1)  # 展平：[batch, 1568]
        output = self.out(x)
        return output

model = CNN().to(device)

维度变化详解

层	输入尺寸	操作	输出尺寸
输入	[32, 1, 28, 28]	-	-
conv1	[32, 1, 28, 28]	卷积(1→64, 5×5, padding=2)	[32, 64, 28, 28]
ReLU	[32, 64, 28, 28]	非线性激活	[32, 64, 28, 28]
MaxPool	[32, 64, 28, 28]	2×2池化，步长2	[32, 64, 14, 14]
conv2第一层	[32, 64, 14, 14]	卷积(64→32)	[32, 32, 14, 14]
conv2第二层	[32, 32, 14, 14]	卷积(32→16)	[32, 16, 14, 14]
MaxPool	[32, 16, 14, 14]	2×2池化	[32, 16, 7, 7]
conv3	[32, 16, 7, 7]	卷积(16→32)	[32, 32, 7, 7]
view展平	[32, 32, 7, 7]	保留batch，合并其余	[32, 1568]
全连接	[32, 1568]	线性映射	[32, 10]

4. 定义损失函数与优化器

loss_fn = nn.CrossEntropyLoss()          # 多分类交叉熵损失
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)  # 随机梯度下降

5. 训练函数

def train(dataloader, model, loss_fn, optimizer):
    model.train()          # 设置为训练模式
    batch_num = 1
    for X, y in dataloader:
        X, y = X.to(device), y.to(device)
        pred = model(X)                # 前向传播
        loss = loss_fn(pred, y)        # 计算损失

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

        # 每100个batch打印一次当前损失
        if batch_num % 100 == 0:
            loss_value = loss.item()    # 将标量张量转为Python浮点数
            print(f"loss: {loss_value:>7f} [number:{batch_num}]")
        batch_num += 1

关键点：

• loss.item()：提取损失张量的数值，便于打印和记录，同时不再保留计算图，节省内存。

• model.train()：启用BatchNorm、Dropout等层的训练行为（本例虽未使用，但属良好习惯）。

6. 测试函数

def test(dataloader, model, loss_fn):
    total_samples = len(dataloader.dataset)   # 测试集总样本数（10000）
    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(X)
            test_loss += loss_fn(pred, y).item()
            # 统计正确预测数：argmax(1)取概率最大的类别
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

    test_loss /= num_batches
    correct /= total_samples
    print(f"Test result:\n {(100 * correct):.2f}% ,avg loss:\n{test_loss:.6f}")

关键点：

• pred.argmax(1)：沿着类别维度（维度1）取最大值索引，得到每个样本的预测类别。

• torch.no_grad()：禁用梯度计算，大幅减少内存消耗。

7. 训练与测试主循环

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

print("Done")
test(test_dataloader, model, loss_fn)

运行结果与分析

运行上述代码，输出如下（节选）：

Using device: cpu
Epoch1
-----------
loss: 2.300039 [number:100]
loss: 2.299318 [number:200]
...
loss: 0.165141 [number:1800]
Epoch2
-----------
...
Epoch10
-----------
loss: 0.000902 [number:1800]
Done
Test result:
 98.54% ,avg loss:
0.043030

结果分析：

• 损失下降趋势：第一个epoch结束时损失已降至0.165，后续每个epoch损失持续走低，最终稳定在0.001左右，表明模型收敛良好。

• 测试准确率：98.54%，在仅10个epoch、简单CNN结构下表现优秀，说明CNN能有效提取手写数字特征。

• 泛化能力：测试集与训练集独立，准确率接近99%，无明显过拟合。

关键点解析

1. 为什么图像任务首选CNN？

• 局部感受野：卷积核只关注局部区域，保留像素间空间关系。

• 参数共享：同一卷积核在整张图上复用，参数量远小于全连接网络。

• 池化下采样：降低分辨率，同时增强平移不变性。

2. nn.Sequential 的作用

将多个层打包成一个模块，forward中只需一行代码即可依次执行，提高代码可读性。若网络存在分支或跳跃连接，则需单独定义各层。

3. view(x.size(0), -1) 的理解

• x.size(0) 是batch大小。

• -1 让PyTorch自动计算剩余维度大小（即 channels × height × width）。

• 作用：将四维特征图 [batch, channels, h, w] 展平为二维 [batch, features]，以输入全连接层。

4. loss.item() 的必要性

loss 是一个包含梯度信息的标量张量，直接打印会显示 tensor(0.2345, grad_fn=<...>)。使用 .item() 提取为Python浮点数，输出更简洁，且不再保留计算图，节省内存。

5. pred.argmax(1) 详解

pred 形状为 [batch, 10]，每行是模型对10个类别的原始分数（logits）。argmax(1) 表示在维度1（类别维度）上取最大值的索引，即模型预测的类别编号。例如某样本输出 [2.3, 0.5, ...]，则预测为0类。

6. 训练与测试模式的差异

• model.train()：启用Dropout、BatchNorm等层的训练行为。

• model.eval()：固定BatchNorm统计量，关闭Dropout，确保推理结果稳定。

• torch.no_grad()：在测试阶段关闭梯度计算，加快计算速度并降低显存占用。

常见问题与改进方向

Q1：为什么使用SGD而不是Adam？
SGD是最简单的优化器，适合入门；Adam通常收敛更快，可尝试替换并观察效果。

Q2：如何提高准确率？

• 增加网络深度或宽度（如再添加卷积层）。

• 使用数据增强（随机旋转、平移等）。

• 调整学习率、使用学习率衰减。

• 尝试更深的预训练模型（如ResNet）进行微调。

Q3：遇到“CUDA out of memory”怎么办？
减小 batch_size，或使用CPU训练（本例CPU即可完成）。

完整代码

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(),
)

'''打包'''
train_dataloader = DataLoader(training_data, batch_size=32)
test_dataloader = DataLoader(test_data, batch_size=32)

'''判断当前设备是否支持GPU，其中mps是苹果的GPU'''
device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using device: {device}")

'''定义卷积神经网络cnn'''
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        self.conv1=nn.Sequential(
            nn.Conv2d(
                in_channels=1,  #输入的通道 64*1*28*28
                out_channels=64, #得到的特征图64
                kernel_size=5, #卷积和 5 * 5
                stride=1,   #步长1
                padding=2,
            ),
            nn.ReLU(),   #64*64*28*28
            nn.MaxPool2d(kernel_size=2)  #64*64*14*14
        )
        self.conv2=nn.Sequential(
            nn.Conv2d(64,32,5,1,2),#64*32*14*14
            nn.ReLU(),
            nn.Conv2d(32,16,5,1,2),#64*16*14*14
            nn.ReLU(),
            nn.MaxPool2d(2)#64*16*7*7
        )
        self.conv3=nn.Sequential(
            nn.Conv2d(16,32,5,1,2),#64*32*7*7
            nn.ReLU(),
        )
        self.out = nn.Sequential(
            nn.Linear(32*7*7,10),
            nn.ReLU()
        )
    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)

loss_fn = nn.CrossEntropyLoss()  #交叉熵损失函数

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

def train(dataloader,model,loss_fn,optimizer):
    model.train()
    batch_num = 1
    for X,y in dataloader:
        X,y = X.to(device),y.to(device)
        pred = model(X)
        loss = loss_fn(pred,y)

        optimizer.zero_grad()  # 这三是固定模版，梯度清零
        loss.backward()  # 反向传播计算得到每个参数
        optimizer.step()  # 根据选择的优化器更新w，我这个选择的是Adam

        loss_value = loss.item()  # 从tensor数据中提取出来，tensor获取损失值
        if batch_num % 100 == 0:  # 打包训练100次输出一次损失值
            print(f"loss: {loss_value:>7f} [number:{batch_num}]")
        batch_num += 1

def test(dataloader,model,loss_fn):
    zongshu = len(dataloader.dataset)
    num_batches = len(dataloader)
    model.eval()  # 测试 w就不更新了
    test_loss, count = 0, 0
    with torch.no_grad():  # 上下文管理器 ，关闭梯度计算，减少计算所用的内存消耗
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model.forward(X)  # .forward可以省略
            test_loss += loss_fn(pred, y).item()
            count += (pred.argmax(1) == y).type(torch.float).sum().item()
        test_loss /= num_batches
        count /= zongshu
        print(f"Test result:\n {(100 * count)}% ,avg loss:\n{test_loss}]")

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

总结

本文通过一个完整的PyTorch代码示例，详细讲解了CNN从数据加载、模型定义到训练测试的全过程。我们不仅给出了可直接运行的代码，还对每一层的维度变化、关键API（如view、item、argmax）进行了深入剖析，并结合最终98.54%的测试准确率验证了CNN在图像分类任务上的强大能力。

CNN中的卷积与池化，则是现代计算机视觉的基石。掌握本文内容，你将具备进一步学习更复杂模型（如ResNet、目标检测）的能力。

---

完整代码：已内嵌于文中，复制即可运行。
参考资料：

• PyTorch官方教程：https://pytorch.org/tutorials/

• MNIST数据集：http://yann.lecun.com/exdb/mnist/