【深度学习】PyTorch CNN 手写数字识别（卷积神经网络）

文章目录

• 一、环境准备
• 代码整体流程
• 二、导入核心库
• 库介绍
• [三、MNIST 数据集加载与可视化](#三、MNIST 数据集加载与可视化)
• 训练集
• 参数讲解
• 测试集
• 数据集可视化（显示手写数字图片）
• [4. DataLoader使用](#4. DataLoader使用)
• 查看数据形状
• [5. 设备配置](#5. 设备配置)
• [6. 定义神经网络（CNN）](#6. 定义神经网络（CNN）)
• [七. 损失函数与优化器基础](#七. 损失函数与优化器基础)
• [八. 模型训练与测试函数实现](#八. 模型训练与测试函数实现)
• 训练函数
• 测试函数
• 代码运行

---

一、环境准备

使用 PyTorch 框架搭建卷积神经网络（CNN），识别 MNIST 数据集中的手写数字（0-9）。

简单来说：让电脑学会看懂手写的 0-9 数字，就像人认数字一样。

在运行代码前，必须安装依赖库，打开电脑的命令提示符（CMD）/ 终端，执行以下命令：

# 安装PyTorch（核心框架）
pip install torch torchvision
# 安装绘图库（用于显示图片）
pip install matplotlib

1. 导入工具库 → 2. 下载数据集 → 3. 可视化数据 → 4. 打包数据 → 5. 配置运行设备 → 6. 搭建 CNN 模型 → 7. 定义训练 / 测试函数 → 8. 训练模型 → 9. 测试准确率

代码整体流程

二、导入核心库

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

库介绍

1. import torch

   作用：导入 PyTorch 核心库，所有深度学习操作的基础。

   类比：就像写作文要先拿笔，torch就是深度学习的 "笔"。

2. from torch import nn

   作用：导入 PyTorch 的神经网络模块（Neural Network）。

   核心功能：提供搭建 CNN、全连接层、卷积层等所有网络层的工具。

   关键：我们搭建的所有 AI 模型，都依赖这个模块。

3. from torch.utils.data import DataLoader

   作用：数据加载器，批量管理数据集。

   通俗理解：把大量图片打包成 "小包裹"，一次性喂给电脑，避免内存不足。

4. from torchvision import datasets

   作用：导入官方封装好的标准数据集。

   包含内容：MNIST 手写数字、CIFAR10 图像分类、ImageNet 等。

   优势：不用自己收集数据，一键下载使用。

5. from torchvision.transforms import ToTensor

   作用：数据格式转换工具，将图片转为张量（Tensor）。

   核心原因：神经网络只能识别张量，不能直接识别图片。

张量是 PyTorch 中唯一的数据格式，可以理解为升级版的数组。

普通图片 / 数字 要先转换格式为张量，再被神经网络识别。

三、MNIST 数据集加载与可视化

训练集

training_data = datasets.MNIST(
    root="data",        # 数据保存路径
    train=True,         # 标记为训练集
    download=True,      # 自动下载数据
    transform=ToTensor(), # 转为张量
)

参数讲解

1. root="data"
   定义：将 MNIST 数据集下载到项目文件夹下的data文件夹中。
   数据量：训练集包含60000 张手写数字图片。
2. train=True
   定义：标记当前数据集为训练集（用于教模型学习）。
   对应：train=False 为测试集（用于检验模型学习效果）。
3. download=True
   定义：自动检测是否已下载数据，未下载则联网下载，已下载则跳过。
   优势：避免重复下载，节省时间。
4. transform=ToTensor()
   定义：将 PIL 格式的图片自动转换为张量。
   必须步骤：神经网络不支持直接处理图片。

测试集

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

数据集可视化（显示手写数字图片）

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

figure.add_subplot(3, 3, i+1)：将画布分为3 行 3 列，共 9 个小区域，依次放置图片。

plt.axis("off")：关闭坐标轴，让图片更美观。

运行结果：

4. DataLoader使用

当60000 张图片一次性传入电脑，内存会爆满，DataLoader 将大数据集拆分为小批量（batch）。

batch_size=64，即每 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}")
    print(f"Shape of y: {y.shape} {y.dtype}")
    break

输出结果：

Shape of X [N, C, H, W]: torch.Size([64, 1, 28, 28])
Shape of y: torch.Size([64]) torch.int64

Shape of X: torch.Size(64, 1, 28, 28)：

N=64：批量大小（64 张图片）

C=1：通道数（灰度图 = 1，彩色图 = 3）

H=28，W=28：图片大小（28×28 像素）

Shape of y: torch.Size(64) torch.int64：

64 个数字标签，数据类型为整数。

5. 设备配置

device = "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")

优先选择 CUDA（NVIDIA 显卡），其次选择 MPS（苹果 M 系列芯片），最后选择 CPU。

6. 定义神经网络（CNN）

CNN：全称为卷积神经网络（Convolutional Neural Network）。

用于处理图像数据（图片识别、目标检测等）

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 第一层卷积
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 16, 5, 1, 2),
            nn.ReLU(),
            nn.MaxPool2d(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)

参数：nn.Conv2d(输入通道, 输出通道, 卷积核大小, 步长, 填充)

1. 图片通道数（灰度图 = 1）
2. 提取的特征图数量
3. 提取图片特征的工具（5×5）
4. 卷积核移动的距离
5. 保持图片大小不变

池化层：nn.MaxPool2d(2)

压缩图片大小，减少计算量。2×2 池化：图片宽高直接缩小一半。

全连接层：nn.Linear(647 7, 10)

输入：64×7×7=3136 个特征

输出：10 个分类（0-9 共 10 个数字

七. 损失函数与优化器基础

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

损失函数（Loss Function）：计算模型预测值与真实值的误差。作用：训练中让损失值越来越小，模型越来越准。

优化器（Optimizer）：根据损失值自动调整模型参数。lr=0.001：学习率（参数调整的步长，新手固定即可）。model.parameters()：需要优化的模型所有参数。

八. 模型训练与测试函数实现

训练函数

def train(dataloader, model, loss_fn, optimizer):
    model.train()  # 开启训练模式
    for batch_num, (X, y) in enumerate(dataloader):
        X, y = X.to(device), y.to(device)  # 数据传入设备
        pred = model(X)                    # 模型预测
        loss = loss_fn(pred, y)            # 计算损失

        # 反向传播（核心：更新参数）
        optimizer.zero_grad()  # 清空梯度
        loss.backward()        # 计算梯度
        optimizer.step()       # 更新参数

        print(f"loss: {loss.item():.6f}  [number:{batch_num}]")
    print("Done !")

测试函数

def test(dataloader, model, loss_fn):
    model.eval()   # 开启测试模式
    total_samples = len(dataloader.dataset)
    total_loss = 0
    correct_cnt = 0

    with torch.no_grad():  # 测试不计算梯度
        for X, y in dataloader:
            X, y = X.to(device), y.to(device)
            pred = model(X)
            total_loss += loss_fn(pred, y).item()
            correct_cnt += (pred.argmax(1) == y).sum().item()

    avg_test_loss = total_loss / len(dataloader)
    acc = correct_cnt / total_samples * 100
    print("Test result:")
    print(f" Accuracy: {acc:.2f}%, Avg loss: {avg_test_loss}")

代码运行

train(train_dataloader, model, loss_fn, optimizer)
test(test_dataloader, model, loss_fn)

控制台持续打印损失值，损失值逐渐下降，说明模型在学习。

运行结果：

loss: 2.302164  [number:0]
loss: 2.290111  [number:1]
loss: 2.305712  [number:2]
loss: 2.265401  [number:3]
loss: 2.272631  [number:4]
.
.
.
loss: 0.297400  [number:933]
loss: 0.006951  [number:934]
loss: 0.002091  [number:935]
loss: 0.426608  [number:936]
loss: 0.007932  [number:937]
Done !
Test result:
 Accuracy: 97.84%, Avg loss: 0.06308832455717123

模型能正确识别 97.84% 的手写数字，效果优秀。