前言

前面几期的内容,我们把神经网络大部分知识点学习了,这一讲我们将用卷积神经网络作为例子,带着大家分析代码部分还有剩下的原理部分,我相信这一讲会很有趣!我们采用的数据集是MNIST手写数据集;如果前面几期的内容没有看的,可以点击下面的链接!觉得内容有帮助的,希望可以关注一下博主!博主将持续更新专栏 AI筑基录 

AI筑基录——损失函数篇-CSDN博客

AI筑基录——激活函数篇-CSDN博客

AI筑基录——梯度下降篇-CSDN博客

AI筑基录——卷积篇-CSDN博客

卷积神经网络

在正式搭建卷积神经网络之前,我们首先要知道,一个神经网络究竟可以用来做什么?简单的说神经网络的任务就是学习输入和输出之间的关系,我们给神经网络大量带有答案的数据,神经网络通过不断调整内部参数,逐渐学会如何根据输入得到正确输出;可以用于分类任务,回归任务,生成任务;

数据集

网络或者模型都是通过训练才拥有能力的,所以我们的整体的流程分为训练测试;这些训练或者测试都需要数据,那么这样的“养分”我们称为数据集;所以我们可以看出,训练集应该是分为测试集和训练集,实际上还可以分多一个验证集,后面我们会专门出一期内容是关于数据如何搭建,处理的文章,这边不展开描述

我们现在介绍一下 MNIST 手写训练集,7 万张 28×28 像素的手写数字灰度图,每张图片都有对应标签(标签也就是数据对应的正确答案),可谓是是深度学习领域的 "Hello World"

图片在计算机的形状:[1,28,28],分别是通道数和长宽,因为训练的过程中我们是批次输入,因此前面还多一个 B 这个维度;

环境配置

配置环境是一门基本功,但是有时候确实也是非常痛苦;但是这个很简单配置,而且有了 AI 之后基本都是它帮助我们配置即可;这边建议下载 Anaconda,库管理平台

这个任务我们只需要 PyTorch + torchvision,没有其他第三方库

安装 PyTorch:推荐用 Conda 创建独立环境,避免包冲突

# 创建环境(Python 3.10+)
conda create -n cnn_mnist python=3.11 -y
conda activate cnn_mnist

# 安装 PyTorch(GPU 版,CUDA 12.x)
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124

通常我们都是在 GPU 上训练的,训练速度会比 CPU 快;

数据集部分不用专门去官网下载,一行代码就可以直接搞定,自动下载到本地缓存

torchvision.datasets.MNIST

核心模块

卷积神经网络,卷积一定是对应的核心,我们看看代码部分(从上到下分析,几乎逐行分析):

class CNNNet(nn.Module):
    
    def __init__(self, num_classes=10):
        super().__init__()

        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)

        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)

        self.pool = nn.MaxPool2d(2, 2)

        self.dropout = nn.Dropout(0.3)

        # 28 → pool → 14 → pool → 7 → pool → 3 (int)
        # 实际: 28/2=14 → 14/2=7 → conv3 后 pool: 7/2=3,所以 128 * 3 * 3 = 1152
        self.fc1 = nn.Linear(128 * 3 * 3, 256)
        self.fc2 = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))   # (B,32,14,14)
        x = self.pool(F.relu(self.bn2(self.conv2(x))))   # (B,64,7,7)
        x = self.pool(F.relu(self.bn3(self.conv3(x))))   # (B,128,3,3)

        x = x.view(x.size(0), -1)                         # flatten
        x = F.relu(self.fc1(self.dropout(x)))
        x = self.fc2(x)
        return x

我们通常是采用类去定义模块,因为具有继承的功能;像 nn.Module 可以理解为 PyTorch 提供的“神经网络基础模板”,我们自己设计网络时,通常都需要继承它;

初始化函数__init__ 用来提前定义网络中需要使用变量,因为我们最后是做分类,因此我们要最后的类别数,一共有 10 类

卷积层,参数从左到右分别表示输入通道,输出通道,卷积核,步长,填充;比如第一层卷积中的输出通道是 32,表示的是每个卷积核都可以尝试寻找一种不同的图像特征,例如横线、竖线、边缘、拐角等,32 个卷积核处理后,就会得到 32 张特征图;

归一化,我们前面讲过,这个就是对输出通道进行归一化;原因我们解释一遍:卷积层提取出来的数据有的可能特别大,有的可能特别小,BatchNorm 会对这些数据做一次整理,避免数据差距过大;

池化层,保留明显特征,同时缩小特征图,减少后续计算量;这边采取的窗口大小是 2 ;

这边有一个点注意一下,后者卷积层的输入一定和前者卷积层的输出的通道数大小一致,否则会报错;

Dropout

这个前面我们没有说过,这个是神经元随机丢弃;如果每次训练都依赖固定的几个神经元,模型可能只是记住了训练数据,而没有真正学会识别数字。Dropout 会随机关闭一部分神经元,迫使网络学习更加稳定、更加通用的特征;像代码里表示有 30% 的神经元会被丢弃;

全连接层

全连接层负责将卷积后的特征综合起来,展开为一行,作出最终判断;

self.fc2 = nn.Linear(256, num_classes)

这个就是把卷积到的最后特征直接输出成 10 个logit得分,哪个数字对应的分数最高,模型就更倾向于认为图片中写的是哪个数字;

前向传播 forward 定义了数据进入网络以后,具体按照什么顺序向前传播;比如这边就是按照卷积,归一,池化作为一个 block 重复 3 次为顺序;最后展开(为全连接做准备),最后经过全连接得到分数;

训练模块

训练模块任务就是让模型把整个训练集学习一遍,这边有一个概念叫做 EPOCH,当这个值为 10 的时候,表示的是这个训练集被网络看了 10 次;

def train_one_epoch(model, loader, optimizer, criterion, device):
    """
    使用整个训练集训练模型一轮。

    返回:
        平均训练损失
        训练集准确率
    """

    # 切换到训练模式
    # Dropout 和 BatchNorm 会按照训练状态工作
    model.train()

    # 记录整轮训练的总损失、正确数量和样本总数
    total_loss = 0
    correct = 0
    total = 0

    # 每次从 DataLoader 中取出一批图片和标签
    for x, y in loader:

        # 将图片和标签移动到 CPU 或 GPU
        x = x.to(device)
        y = y.to(device)

        # 清空上一批数据计算出的梯度
        optimizer.zero_grad()

        # 前向传播:模型根据图片进行预测
        out = model(x)

        # 比较预测结果和真实标签,计算损失
        loss = criterion(out, y)

        # 反向传播:计算每个参数的梯度
        loss.backward()

        # 根据梯度更新模型参数
        optimizer.step()

        # 当前 loss 是这一批图片的平均损失
        # 乘以批次大小,得到这一批图片的损失总和
        total_loss += loss.item() * x.size(0)

        # 找到每张图片分数最高的类别,并统计预测正确数量
        predictions = out.argmax(dim=1)
        correct += (predictions == y).sum().item()

        # 累加已经处理过的图片数量
        total += x.size(0)

    # 返回整个训练集的平均损失和准确率
    average_loss = total_loss / total
    accuracy = 100.0 * correct / total

    return average_loss, accuracy

函数接受了 5 个参数,分别是模型,加载器,优化器,损失函数,设备;

之所以需要专门设置训练模式,是因为模型中有一些层在训练和测试时表现不同;代码的注释写的算是比较相信的,因此我就简单的说一下:

循环的 x、y 表示就是图片和标签,我们放到 GPU 中进行加载,接着就是把上一轮的梯度清零,接着就是进行模型预测,损失函数对比后进行反向传播,最后用优化器跟新参数;结合我们之前的原理部分,这个代码还是很清晰理解的;最后就是对简单的结果进行计算;

推理模块

推理模块的任务就是看模型训练的效果怎么样,一般是在验证集或者是测试集上推理;

@torch.no_grad()
def evaluate(model, loader, criterion, device):
    """
    在验证集或测试集上评估模型。

    返回:
        平均损失
        分类准确率
    """

    # 将模型切换到评估模式
    # Dropout 停止随机关闭神经元
    # BatchNorm 使用训练时保存的统计信息
    model.eval()

    # 记录总损失、预测正确数量和样本总数
    total_loss = 0
    correct = 0
    total = 0

    # 每次取出一批图片和标签
    for x, y in loader:

        # 将数据移动到 CPU 或 GPU
        x = x.to(device)
        y = y.to(device)

        # 前向传播,得到十个类别分数
        out = model(x)

        # 计算当前批次的损失
        loss = criterion(out, y)

        # 累计当前批次的损失总和
        total_loss += loss.item() * x.size(0)

        # 找到每张图片分数最高的类别
        predictions = out.argmax(dim=1)

        # 统计预测正确的图片数量
        correct += (predictions == y).sum().item()

        # 累计已经处理的图片数量
        total += x.size(0)

    # 计算整个数据集的平均损失
    average_loss = total_loss / total

    # 计算整个数据集的准确率
    accuracy = 100.0 * correct / total

    return average_loss, accuracy

函数接受了 4 个参数,分别是模型,加载器,损失函数,设备

与训练过程不同,评估阶段只需要进行前向传播,不需要计算梯度,也不需要更新模型参数。因此,我们使用 @torch.no_grad() 关闭梯度记录,从而减少计算量和显存占用;当所有图片测试完成后,函数会返回整个数据集的平均损失和准确率

主函数

if __name__ == "__main__":
    # 只有直接运行当前文件时,下面的代码才会执行

    print("-" * 50)
    print("CNN 手写数字识别 (MNIST) 训练与测试")
    print("-" * 50)

    # 自动选择 GPU 或 CPU
    device = torch.device(
        "cuda" if torch.cuda.is_available() else "cpu"
    )
    print(f"[Device] {device}")

    # ==================== 1. 准备数据 ====================

    # 定义图片预处理方式
    transform = transforms.Compose([
        # 将图片转换成 PyTorch 张量
        transforms.ToTensor(),

        # 对 MNIST 图片进行标准化
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 加载 MNIST 训练集
    train_ds = datasets.MNIST(
        root="data/mnist",
        train=True,
        download=True,
        transform=transform
    )

    # 加载 MNIST 测试集
    test_ds = datasets.MNIST(
        root="data/mnist",
        train=False,
        download=True,
        transform=transform
    )

    # 每次取出 128 张训练图片
    train_loader = DataLoader(
        train_ds,
        batch_size=128,
        shuffle=True,
        num_workers=2
    )

    # 每次取出 128 张测试图片
    test_loader = DataLoader(
        test_ds,
        batch_size=128,
        shuffle=False,
        num_workers=2
    )

    print(
        f"[Data] Train: {len(train_ds)} · "
        f"Test: {len(test_ds)}"
    )

    # ==================== 2. 创建模型 ====================

    # 创建 10 分类的卷积神经网络,并移动到 CPU 或 GPU
    model = CNNNet(num_classes=10).to(device)

    # 统计模型参数量
    total_params = (
        sum(p.numel() for p in model.parameters()) / 1e6
    )
    print(f"[Model] 参数量: {total_params:.2f}M")

    # 定义交叉熵损失函数
    criterion = nn.CrossEntropyLoss()

    # 定义 Adam 优化器,学习率为 0.001
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=0.001
    )

    # ==================== 3. 训练模型 ====================

    # 将整个训练集学习 5 遍
    epochs = 5

    for epoch in range(1, epochs + 1):
        # 使用训练集训练一轮
        train_loss, train_acc = train_one_epoch(
            model,
            train_loader,
            optimizer,
            criterion,
            device
        )

        # 使用测试集检查当前模型效果
        test_loss, test_acc = evaluate(
            model,
            test_loader,
            criterion,
            device
        )

        # 打印当前轮次的训练和测试结果
        print(
            f"Epoch {epoch:2d} | "
            f"Train Loss: {train_loss:.4f} "
            f"Acc: {train_acc:.2f}% | "
            f"Test Loss: {test_loss:.4f} "
            f"Acc: {test_acc:.2f}%"
        )

    # 打印最后一轮的测试准确率
    print("-" * 50)
    print(f"[Result] 最终测试准确率: {test_acc:.2f}%")
    print("-" * 50)

    # ==================== 4. 单张图片预测 ====================

    # 将模型切换到评估模式
    model.eval()

    # 取出测试集中的第一张图片和真实标签
    sample, label = test_ds[0]

    # 单张预测不需要计算梯度
    with torch.no_grad():
        # 原形状:[1, 28, 28]
        # 增加批次维度后:[1, 1, 28, 28]
        sample = sample.unsqueeze(0).to(device)

        # 得到模型输出
        out = model(sample)

        # 找到分数最高的数字类别
        pred = out.argmax(dim=1).item()

    # 输出真实答案和模型预测结果
    result = "✓" if pred == label else "✗"

    print(
        f"[Inference] 第一张图: "
        f"真实={label}, 预测={pred} {result}"
    )

代码的注释足够详细,因此这边就说几个点:

  1. 图片最后要变成 tensor 的格式,因为 MNIST 原本是一张图片,计算机需要先把它转换成 PyTorch 能够处理的张量;
  2. 数据因为是一批一批处理的,因此我们要用加载器打包;
  3. 子进程帮助读取和处理数据,一般这样的话速度会比较快,同时显卡的占用率会提高

效果

简单的运行了一下,可以看看效果

完整代码:

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader
from torchvision import datasets, transforms


class CNNNet(nn.Module):
    """简单的 CNN,用于 MNIST 手写数字分类(28x28 灰度图 → 10 类)"""
    def __init__(self, num_classes=10):
        super().__init__()

        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.bn1 = nn.BatchNorm2d(32)

        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
        self.bn2 = nn.BatchNorm2d(64)

        self.conv3 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1)
        self.bn3 = nn.BatchNorm2d(128)

        self.pool = nn.MaxPool2d(2, 2)

        self.dropout = nn.Dropout(0.3)

        # 28 → pool → 14 → pool → 7 → pool → 3 (int)
        # 实际: 28/2=14 → 14/2=7 → conv3 后 pool: 7/2=3,所以 128 * 3 * 3 = 1152
        self.fc1 = nn.Linear(128 * 3 * 3, 256)
        self.fc2 = nn.Linear(256, num_classes)

    def forward(self, x):
        x = self.pool(F.relu(self.bn1(self.conv1(x))))   # (B,32,14,14)
        x = self.pool(F.relu(self.bn2(self.conv2(x))))   # (B,64,7,7)
        x = self.pool(F.relu(self.bn3(self.conv3(x))))   # (B,128,3,3)

        x = x.view(x.size(0), -1)                         # flatten
        x = F.relu(self.fc1(self.dropout(x)))
        x = self.fc2(x)
        return x


def train_one_epoch(model, loader, optimizer, criterion, device):
    model.train()
    total_loss, correct, total = 0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)

        optimizer.zero_grad()
        out = model(x)
        loss = criterion(out, y)
        loss.backward()
        optimizer.step()

        total_loss += loss.item() * x.size(0)
        correct += (out.argmax(1) == y).sum().item()
        total += x.size(0)

    return total_loss / total, 100.0 * correct / total


@torch.no_grad()
def evaluate(model, loader, criterion, device):
    model.eval()
    total_loss, correct, total = 0, 0, 0
    for x, y in loader:
        x, y = x.to(device), y.to(device)
        out = model(x)
        loss = criterion(out, y)

        total_loss += loss.item() * x.size(0)
        correct += (out.argmax(1) == y).sum().item()
        total += x.size(0)

    return total_loss / total, 100.0 * correct / total


if __name__ == "__main__":
    print("-" * 50)
    print("CNN 手写数字识别 (MNIST) 训练与测试")
    print("-" * 50)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"[Device] {device}")

    # ---------- 数据 ----------
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    train_ds = datasets.MNIST("data/mnist", train=True, download=True, transform=transform)
    test_ds = datasets.MNIST("data/mnist", train=False, download=True, transform=transform)

    train_loader = DataLoader(train_ds, batch_size=128, shuffle=True, num_workers=2)
    test_loader = DataLoader(test_ds, batch_size=128, shuffle=False, num_workers=2)

    print(f"[Data] Train: {len(train_ds)} · Test: {len(test_ds)}")

    # ---------- 模型 ----------
    model = CNNNet(num_classes=10).to(device)
    total_params = sum(p.numel() for p in model.parameters()) / 1e6
    print(f"[Model] 参数量: {total_params:.2f}M")

    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    # ---------- 训练 ----------
    epochs = 5
    for epoch in range(1, epochs + 1):
        train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, criterion, device)
        test_loss, test_acc = evaluate(model, test_loader, criterion, device)
        print(f"Epoch {epoch:2d} | "
              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.2f}% | "
              f"Test Loss: {test_loss:.4f} Acc: {test_acc:.2f}%")

    print("-" * 50)
    print(f"[Result] 最终测试准确率: {test_acc:.2f}%")
    print("-" * 50)

    # ---------- 单张预测验证 ----------
    model.eval()
    sample, label = test_ds[0]
    with torch.no_grad():
        pred = model(sample.unsqueeze(0).to(device)).argmax(1).item()
    print(f"[Inference] 第一张图: 真实={label}, 预测={pred} {'✓' if pred == label else '✗'}")

总结

因为这一块的代码注释确实是真的很详细,如果我再继续赘述确实显得比较无聊;我们可以看到一个简单的神经网络的脉络是比较简单的;希望大家可以做到手搓这样的卷积网络,难度不大的;我们在最后结果可以发现实际上这样的一个简单的网络参数量还是很多的,有兴趣的同学可以看看我之前主页的mapping net,这篇文章就是有少量的参数去逼近大参数的效果;

AI筑基录 的专栏会慢慢更新完,带大家有体系的理解神经网络的由来!制作不易,喜欢博主文章的可以点赞收藏关注,这些都是我持续更新的动力!

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐