一、环境准备

使用 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(6477, 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% 的手写数字,效果优秀。

Logo

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念,把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起,为开发者提供从开发、训练到部署的一站式体验。

更多推荐