深度学习系统学习系列【8】之设计卷积神经网络架构（Pytorch版本）

网络层间排列规律

• 卷积神经网络中最常见的是卷积层后接Pooling 层，目的是减少下一次卷积输入图像大小，然后重复该过程，提出图像中的高维特征，最后通过全连接层得到输出。因此最常见的卷积神经网络结构：
  $$\begin{gathered} INPUT\to [CONV]\to [POOL]\to [FC]\to OUTPUT \\ CONV为卷积层，POOL为Pooling层，FC为全连接层，[x]为重复x层多次 \end{gathered}$$
• $INPUT\to FC$：实现一个简单的线性分类器。
• $INPUT\to CONV$：实现一个滤波操作，如高斯模糊、中值滤波等。
• $INPUT\to [CONV\to POLL]\times 2\to FC\to OUTPUT$：经过两次卷积层接 Pooling 层，实现小规模的卷积神经分类网络。
• $INPUT\to [CONV\to CONV\to POOL]\times 3\to [FC]\times 2\to OUTPUT$：多次连续两个卷积层后接一个Pooling层，该思路适用于深层次网络（如VGGNet、ResNet 等)。因为在执行 Pooling 操作前，多次小卷积核卷积可以有效减少网络权重参数，从输入数据中学习到更高维特征。

卷积神经网络（CNN）参数设计规范

输入层设计

• 尺寸要求：输入矩阵边长应可被2多次整除，常用尺寸为32、64、96、224、384或512。
  • 理由：便于卷积层和池化层计算（如每次池化输出特征图尺寸减半），减少数据丢失。
  • 示例：AlexNet经典输入尺寸为224×224。

卷积层设计

1. 卷积核尺寸

   • 优先使用小卷积核（如1×1、3×3、5×5），深层网络应减小卷积核尺寸。
   • 理由：
     • 避免因感知区域过大导致特征提取困难。
     • 小卷积核参数更少，计算效率更高。

2. 步长（Stride）

   • 步长建议设为1，空间下采样由池化层完成。
   • 理由：小步长能保留更多特征信息。

3. 填充（Padding）

   • 使用Same Padding（零填充）保持输入/输出尺寸一致。
   • 填充规则：
     • 卷积核尺寸K=3 → padding=1
     • K=5 → padding=2
     • 通用公式：padding=(K-1)/2

池化层设计

• 推荐参数：2×2窗口，步长为2的Max Pooling。
• 注意事项： 避免窗口尺寸＞3，否则易导致信息丢失。 下采样过于激进会显著降低模型性能。

全连接层设计

• 层数限制：不超过3层（通常为2个全连接层+1个输出层）。
• 理由：过多全连接层易引发过拟合和梯度消失。

• 核心原则：通过小卷积核、适度填充和保守下采样平衡特征提取与计算效率，避免信息丢失。

可视化手写字体特征网络

MNIST手写字体数据库

• MNIST 数据集是一个手写体数据集。，数据集中每一个样本都是$09$的手写数字，该数据集由 4 部分组成：训练图片集、 训练标签集、测试图片集和测试标签集。训练集中有60000个样本，测试集中有 10000 个样本，每张照片均由 28×28 大小的灰度图像组成。
• 为了便于存储和下载，官方对MNIST数据集的图片进行集中处理，将每一张图片拉伸成为$(1,784)$的向量表示。
  

 # ---------------------------
    # 数据预处理和加载
    # ---------------------------
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST 均值和标准差
    ])

    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    val_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

    train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
    val_loader = DataLoader(dataset=val_dataset, batch_size=64, shuffle=False)

LeNet-5 模型介绍

• LeNet-5 是由 Yann LeCun 等人于 1998 年提出的经典卷积神经网络（CNN），主要用于手写数字识别（如 MNIST 数据集）。它是早期 CNN 的代表，奠定了现代卷积神经网络的基础架构。
• LeNet-5 由 7 层组成（2 个卷积层 + 2 个池化层 + 3 个全连接层），结构如下：

关键特点

1. 小卷积核（5×5）

   • 受限于当时的计算能力，采用较小的卷积核提取局部特征。
   • 现代 CNN（如 VGG、ResNet）更多使用 3×3 卷积核。

2. 平均池化（而非 Max Pooling）

   • LeNet-5 使用 2×2 平均池化（计算窗口内像素均值）。
   • 现代 CNN 普遍采用 Max Pooling（保留最显著特征）。

3. 激活函数：Tanh / Sigmoid

   • 原始 LeNet-5 使用 Tanh 或 Sigmoid 激活函数。
   • 现代 CNN 通常使用 ReLU（计算更快，缓解梯度消失）。

4. 全连接层占比大

   • 后 3 层均为全连接层（C5、F6、输出层），参数量较大。
   • 现代 CNN 倾向于减少全连接层（如用全局平均池化替代）。

代码实现（Pytorch版）

• 在 PyTorch 中，不像 Keras 那样内置 model.summary() 方法来查看模型结构和参数信息。可以使用第三方库 torchinfo来实现类似功能。

pip install torchinfo
summary(model, input_size=(1, 1, 28, 28))

项目结构和文件

• data：保存MNIST数据集
• leNet5_checkpoints：保存训练模型
• test_images：保存用于测试的图片
• LetNet5_stu：模型定义、训练和保存
• predict_single_digits：简单使用图片测试训练结果。

源码地址

• 本次项目案例地址：letNet5_stu

模型定义和概览

import torch
import torch.nn as nn
import torch.nn.functional as functional
from torchinfo import summary
class LeNet5(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet5, self).__init__()
        # 卷积层
        """
        第一个层卷积：conv1
        in_channels=1：输入图像的通道数。例如，MNIST手写数字图像是灰度图，所以通道数为1。
        out_channels=6：输出通道数（即卷积核数量），表示这一层会提取6个特征图。
        kernel_size=5：卷积核大小为 5x5。
        padding=2：在输入图像四周填充 2 层像素，使得输出特征图与输入图像的空间尺寸一致（即保持尺寸不变）。
        """
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)  # 修改 padding 以保持尺寸
        """
        第二个层卷积：conv2
        in_channels=6：上一层输出了 6 个特征图，因此当前层的输入通道数为 6。
        out_channels=16：本层使用 16 个卷积核，输出 16 个特征图。
        kernel_size=5：卷积核大小为 5x5。
        没有指定 padding，默认为 0，因此输出尺寸会比输入小（如输入 14x14 → 输出 10x10）。
        """
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)

        # 池化层和激活函数显式定义
        """
        定义一个 ReLU（Rectified Linear Unit）激活函数层
        作用：
            1. 在卷积层之后使用 ReLU 可以增强模型的非线性表达能力
        将 ReLU 实例化为类成员变量，因此可以在 forward 函数中通过 self.relu() 调用 
            2. 同时也方便torchinfo打印输出信息
        """
        self.relu = nn.ReLU()
        """
        定义一个二维最大池化（Max Pooling）层
        最大池化是一种下采样操作，它从每个池化窗口中取最大值作为输出。
        作用：
            1. 减少特征图的空间尺寸（高度、宽度），从而减少计算量和参数数量。
            2. 提高模型对小范围平移的鲁棒性。
        如：输入尺寸为 28x28 的特征图经过 kernel_size=2, stride=2 的 MaxPool 后，输出尺寸变为 14x14    
        """
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # 全连接层
        """
        第一个全连接层（Fully Connected Layer）
        作用：将卷积提取到的高维特征映射为一个长度为 120 的向量
        16 * 5 * 5：输入特征的数量。由前面卷积和池化操作后输出的特征图展平后的维度。
            16 是上一层输出的通道数（即特征图数量）
            5 * 5 是每个特征图的空间尺寸（高度 × 宽度），来源于经过多次池化后的结果
        120：该层输出的神经元数量，是 LeNet-5 的设计结构中规定的
        """
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 根据图像尺寸调整输入大小
        """
        定义第二个全连接层
        作用：进一步压缩或抽象特征信息，提升模型表达能力
        120：输入来自上一层 (fc1) 的输出
        84：输出神经元数量，也是 LeNet-5 结构中预设的中间层节点数
        """
        self.fc2 = nn.Linear(120, 84)
        """
        定义最后一个全连接层，用于分类输出
        作用：输出每个类别的预测得分，通常配合 Softmax 激活函数得到概率分布
        84：输入来自上一层 (fc2) 的输出。
        mum_classes：类别总数，默认为 10，适用于如 MNIST 手写数字识别任务。

        """
        self.fc3 = nn.Linear(84, num_classes)

        # 可选：加载预训练权重的方法可以自行实现或使用torch.load

    def forward(self, x):
        # 第一层卷积 + 池化
        x = self.relu(self.conv1(x))
        x = self.pool(x)

        # 第二层卷积 + 池化
        x = self.relu(self.conv2(x))
        x = self.pool(x)

        # 展平
        """
        将输入张量 x 从卷积层输出的形状（如 [batch_size, channels, height, width]）展平为二维张量，以便输入到全连接层。
        -1：自动推导 batch size 大小。
        16 * 5 * 5：这是经过前面卷积和池化操作后每个样本的特征总数。
            16 是通道数（由 conv2 输出）。
            5 * 5 是每个通道的空间尺寸（高度 × 宽度），来源于经过多次池化后的结果。
        将四维张量 [batch_size, 16, 5, 5] 转换为二维张量 [batch_size, 400]，其中 400 = 16×5×5。
        """
        x = x.view(-1, 16 * 5 * 5)

        # 全连接层
        """
        功能：第一个全连接层 + ReLU 激活函数。
        self.fc1 是定义在 __init__ 中的线性层：nn.Linear(16*5*5, 120)。
        输入维度：400（即 16*5*5）。
        输出维度：120。
        作用：将展平后的特征向量映射到更高层次的表示空间。
        """
        x = functional.relu(self.fc1(x))
        """
        功能：第二个全连接层 + ReLU 激活函数。
        self.fc2：nn.Linear(120, 84)。
        输入维度：120。
        输出维度：84。
        作用：进一步压缩和抽象特征信息。
        """
        x = functional.relu(self.fc2(x))
        """
        最后一个全连接层，用于最终分类输出。
        self.fc3：nn.Linear(84, num_classes)，默认 num_classes=10。
        输入维度：84。
        输出维度：类别数量（如 MNIST 数据集为 10 类）。
        作用：输出每个类别的预测得分。通常配合 Softmax 函数得到概率分布。
        """
        x = self.fc3(x)

        return x

# 示例输入 (batch_size=1, channel=1, height=28, width=28)
# x = torch.randn(1, 1, 28, 28).to(device)
# y = model(x)
# print(y)

if  __name__ == '__main__':
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 创建模型实例
    model = LeNet5().to(device)
    # 打印模型结构
    print(model)
    summary(model, input_size=(1, 1, 28, 28))

LeNet5(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (relu): ReLU()
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
LeNet5                                   [1, 10]                   --
├─Conv2d: 1-1                            [1, 6, 28, 28]            156
├─ReLU: 1-2                              [1, 6, 28, 28]            --
├─MaxPool2d: 1-3                         [1, 6, 14, 14]            --
├─Conv2d: 1-4                            [1, 16, 10, 10]           2,416
├─ReLU: 1-5                              [1, 16, 10, 10]           --
├─MaxPool2d: 1-6                         [1, 16, 5, 5]             --
├─Linear: 1-7                            [1, 120]                  48,120
├─Linear: 1-8                            [1, 84]                   10,164
├─Linear: 1-9                            [1, 10]                   850
==========================================================================================
Total params: 61,706
Trainable params: 61,706
Non-trainable params: 0
Total mult-adds (M): 0.42
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.05
Params size (MB): 0.25
Estimated Total Size (MB): 0.30
==========================================================================================

LeNet5网络训练和模型保存

• LeNet5_stu.py

import os
import torch
import torch.nn as nn
import torch.nn.functional as functional
from torch import optim
from torchinfo import summary
from torchvision import datasets, transforms
from torch.utils.data import DataLoader


class LeNet5(nn.Module):
    def __init__(self, num_classes=10):
        super(LeNet5, self).__init__()
        # 卷积层
        """
        第一个层卷积：conv1
        in_channels=1：输入图像的通道数。例如，MNIST手写数字图像是灰度图，所以通道数为1。
        out_channels=6：输出通道数（即卷积核数量），表示这一层会提取6个特征图。
        kernel_size=5：卷积核大小为 5x5。
        padding=2：在输入图像四周填充 2 层像素，使得输出特征图与输入图像的空间尺寸一致（即保持尺寸不变）。
        """
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5, padding=2)  # 修改 padding 以保持尺寸
        """
        第二个层卷积：conv2
        in_channels=6：上一层输出了 6 个特征图，因此当前层的输入通道数为 6。
        out_channels=16：本层使用 16 个卷积核，输出 16 个特征图。
        kernel_size=5：卷积核大小为 5x5。
        没有指定 padding，默认为 0，因此输出尺寸会比输入小（如输入 14x14 → 输出 10x10）。
        """
        self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)

        # 池化层和激活函数显式定义
        """
        定义一个 ReLU（Rectified Linear Unit）激活函数层
        作用：
            1. 在卷积层之后使用 ReLU 可以增强模型的非线性表达能力
        将 ReLU 实例化为类成员变量，因此可以在 forward 函数中通过 self.relu() 调用 
            2. 同时也方便torchinfo打印输出信息
        """
        self.relu = nn.ReLU()
        """
        定义一个二维最大池化（Max Pooling）层
        最大池化是一种下采样操作，它从每个池化窗口中取最大值作为输出。
        作用：
            1. 减少特征图的空间尺寸（高度、宽度），从而减少计算量和参数数量。
            2. 提高模型对小范围平移的鲁棒性。
        如：输入尺寸为 28x28 的特征图经过 kernel_size=2, stride=2 的 MaxPool 后，输出尺寸变为 14x14    
        """
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)

        # 全连接层
        """
        第一个全连接层（Fully Connected Layer）
        作用：将卷积提取到的高维特征映射为一个长度为 120 的向量
        16 * 5 * 5：输入特征的数量。由前面卷积和池化操作后输出的特征图展平后的维度。
            16 是上一层输出的通道数（即特征图数量）
            5 * 5 是每个特征图的空间尺寸（高度 × 宽度），来源于经过多次池化后的结果
        120：该层输出的神经元数量，是 LeNet-5 的设计结构中规定的
        """
        self.fc1 = nn.Linear(16 * 5 * 5, 120)  # 根据图像尺寸调整输入大小
        """
        定义第二个全连接层
        作用：进一步压缩或抽象特征信息，提升模型表达能力
        120：输入来自上一层 (fc1) 的输出
        84：输出神经元数量，也是 LeNet-5 结构中预设的中间层节点数
        """
        self.fc2 = nn.Linear(120, 84)
        """
        定义最后一个全连接层，用于分类输出
        作用：输出每个类别的预测得分，通常配合 Softmax 激活函数得到概率分布
        84：输入来自上一层 (fc2) 的输出。
        mum_classes：类别总数，默认为 10，适用于如 MNIST 手写数字识别任务。

        """
        self.fc3 = nn.Linear(84, num_classes)

        # 可选：加载预训练权重的方法可以自行实现或使用torch.load

    def forward(self, x):
        # 第一层卷积 + 池化
        x = self.relu(self.conv1(x))
        x = self.pool(x)

        # 第二层卷积 + 池化
        x = self.relu(self.conv2(x))
        x = self.pool(x)

        # 展平
        """
        将输入张量 x 从卷积层输出的形状（如 [batch_size, channels, height, width]）展平为二维张量，以便输入到全连接层。
        -1：自动推导 batch size 大小。
        16 * 5 * 5：这是经过前面卷积和池化操作后每个样本的特征总数。
            16 是通道数（由 conv2 输出）。
            5 * 5 是每个通道的空间尺寸（高度 × 宽度），来源于经过多次池化后的结果。
        将四维张量 [batch_size, 16, 5, 5] 转换为二维张量 [batch_size, 400]，其中 400 = 16×5×5。
        """
        x = x.view(-1, 16 * 5 * 5)

        # 全连接层
        """
        功能：第一个全连接层 + ReLU 激活函数。
        self.fc1 是定义在 __init__ 中的线性层：nn.Linear(16*5*5, 120)。
        输入维度：400（即 16*5*5）。
        输出维度：120。
        作用：将展平后的特征向量映射到更高层次的表示空间。
        """
        x = functional.relu(self.fc1(x))
        """
        功能：第二个全连接层 + ReLU 激活函数。
        self.fc2：nn.Linear(120, 84)。
        输入维度：120。
        输出维度：84。
        作用：进一步压缩和抽象特征信息。
        """
        x = functional.relu(self.fc2(x))
        """
        最后一个全连接层，用于最终分类输出。
        self.fc3：nn.Linear(84, num_classes)，默认 num_classes=10。
        输入维度：84。
        输出维度：类别数量（如 MNIST 数据集为 10 类）。
        作用：输出每个类别的预测得分。通常配合 Softmax 函数得到概率分布。
        """
        x = self.fc3(x)

        return x


# 自定义模型保存函数
def save_checkpoint(model, val_acc, filepath='./leNet5_checkpoints'):
    if not os.path.exists(filepath):
        os.makedirs(filepath)
    filename = f"{filepath}/model_epoch_{epoch:02d}_valacc_{val_acc:.3f}.pth"
    torch.save(model.state_dict(), filename)
    print(f"Saved model to {filename}")


# 示例输入 (batch_size=1, channel=1, height=28, width=28)
# x = torch.randn(1, 1, 28, 28).to(device)
# y = model(x)
# print(y)


if __name__ == '__main__':

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

    # 创建模型实例
    model = LeNet5().to(device)
    print(model)
    summary(model, input_size=(1, 1, 28, 28))

    # ---------------------------
    # 数据预处理和加载
    # ---------------------------
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))  # MNIST 均值和标准差
    ])

    train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
    val_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)

    train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
    val_loader = DataLoader(dataset=val_dataset, batch_size=64, shuffle=False)

    # ---------------------------
    # 损失函数和优化器
    # ---------------------------
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adadelta(model.parameters())

    # ---------------------------
    # 模型保存路径
    # ---------------------------
    checkpoint_dir = "leNet5_checkpoints"
    if not os.path.exists(checkpoint_dir):
        os.makedirs(checkpoint_dir)

    best_val_acc = 0.0

    # ---------------------------
    # 训练和验证循环
    # ---------------------------
    epochs = 10
    for epoch in range(epochs):
        # 训练阶段
        model.train()
        running_loss = 0.0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()

        # 验证阶段
        model.eval()
        correct = total = 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                _, preds = torch.max(outputs, 1)
                total += labels.size(0)
                correct += (preds == labels).sum().item()

        val_acc = correct / total
        avg_loss = running_loss / len(train_loader)
        print(f"Epoch [{epoch + 1}/{epochs}], Loss: {avg_loss:.4f}, Val Accuracy: {val_acc:.4f}")

        # 保存最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            save_path = os.path.join(checkpoint_dir, f"model_epoch_{epoch + 1:02d}_valacc_{val_acc:.3f}.pth")
            torch.save(model.state_dict(), save_path)
            print(f"Saved best model to {save_path}")

• 训练结果：

Using device: cuda
LeNet5(
  (conv1): Conv2d(1, 6, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (conv2): Conv2d(6, 16, kernel_size=(5, 5), stride=(1, 1))
  (relu): ReLU()
  (pool): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  (fc1): Linear(in_features=400, out_features=120, bias=True)
  (fc2): Linear(in_features=120, out_features=84, bias=True)
  (fc3): Linear(in_features=84, out_features=10, bias=True)
)
==========================================================================================
Layer (type:depth-idx)                   Output Shape              Param #
==========================================================================================
LeNet5                                   [1, 10]                   --
├─Conv2d: 1-1                            [1, 6, 28, 28]            156
├─ReLU: 1-2                              [1, 6, 28, 28]            --
├─MaxPool2d: 1-3                         [1, 6, 14, 14]            --
├─Conv2d: 1-4                            [1, 16, 10, 10]           2,416
├─ReLU: 1-5                              [1, 16, 10, 10]           --
├─MaxPool2d: 1-6                         [1, 16, 5, 5]             --
├─Linear: 1-7                            [1, 120]                  48,120
├─Linear: 1-8                            [1, 84]                   10,164
├─Linear: 1-9                            [1, 10]                   850
==========================================================================================
Total params: 61,706
Trainable params: 61,706
Non-trainable params: 0
Total mult-adds (M): 0.42
==========================================================================================
Input size (MB): 0.00
Forward/backward pass size (MB): 0.05
Params size (MB): 0.25
Estimated Total Size (MB): 0.30
==========================================================================================
100.0%
100.0%
100.0%
100.0%
Epoch [1/10], Loss: 0.1728, Val Accuracy: 0.9834
Saved best model to ./leNet5_checkpoints\model_epoch_01_valacc_0.983.pth
Epoch [2/10], Loss: 0.0485, Val Accuracy: 0.9876
Saved best model to ./leNet5_checkpoints\model_epoch_02_valacc_0.988.pth
Epoch [3/10], Loss: 0.0363, Val Accuracy: 0.9890
Saved best model to ./leNet5_checkpoints\model_epoch_03_valacc_0.989.pth
Epoch [4/10], Loss: 0.0275, Val Accuracy: 0.9892
Saved best model to ./leNet5_checkpoints\model_epoch_04_valacc_0.989.pth
Epoch [5/10], Loss: 0.0229, Val Accuracy: 0.9852
Epoch [6/10], Loss: 0.0185, Val Accuracy: 0.9886
Epoch [7/10], Loss: 0.0133, Val Accuracy: 0.9896
Saved best model to ./leNet5_checkpoints\model_epoch_07_valacc_0.990.pth
Epoch [8/10], Loss: 0.0117, Val Accuracy: 0.9908
Saved best model to ./leNet5_checkpoints\model_epoch_08_valacc_0.991.pth
Epoch [9/10], Loss: 0.0084, Val Accuracy: 0.9900
Epoch [10/10], Loss: 0.0085, Val Accuracy: 0.9905

单图单数字测试

• 安装cv2图像预处理：转灰度、二值化、去噪

pip install opencv-python

• leNet5_stu/test_images/single_digit.png
  

• predict_single_digits.py

import os
import torch
import cv2
import torchvision.transforms as transforms
from deep_learning_skill.leNet5_stu.LeNet5_stu import LeNet5  # 导入模型定义

# ---------------------------
# 配置参数
# ---------------------------
MODEL_PATH = 'leNet5_checkpoints/model_epoch_08_valacc_0.991.pth'
IMAGE_PATH = 'test_images/single_digit.png'  # 替换为自己的图片路径

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


# ---------------------------
# 图像预处理函数
# ---------------------------
def preprocess_image(image_path):
    # 加载图像并转为灰度图
    image = cv2.imread(image_path, 0)
    if image is None:
        raise FileNotFoundError(f"找不到图像文件：{image_path}")

    # 调整尺寸为 28x28
    image_resized = cv2.resize(image, (28, 28))

    # 图像增强对比度（可选）
    _, image_binary = cv2.threshold(image_resized, 128, 255, cv2.THRESH_BINARY_INV)

    # 转换为 Tensor 并标准化（使用 MNIST 的均值和标准差）
    transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize((0.1307,), (0.3081,))
    ])

    # 增加 batch 维度 [1, 1, 28, 28]
    image_tensor = transform(image_binary).unsqueeze(0).to(device)
    return image_tensor


# ---------------------------
# 推理函数
# ---------------------------
def predict_digit(model, image_tensor):
    model.eval()
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output, 1)
    return predicted.item()


# ---------------------------
# 主程序入口
# ---------------------------
if __name__ == '__main__':
    # 加载模型结构
    model = LeNet5().to(device)

    # 加载训练好的权重
    if not os.path.exists(MODEL_PATH):
        raise FileNotFoundError(f"找不到模型文件：{MODEL_PATH}")
    model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
    print(f"模型已加载：{MODEL_PATH}")

    # 加载并预处理图像
    if not os.path.exists(IMAGE_PATH):
        raise FileNotFoundError(f"找不到测试图片：{IMAGE_PATH}")
    image_tensor = preprocess_image(IMAGE_PATH)

    # 进行预测
    predicted_label = predict_digit(model, image_tensor)
    print(f"识别结果：这张图片中的数字是 -> {predicted_label}")

Using device: cuda
模型已加载：leNet5_checkpoints/model_epoch_08_valacc_0.991.pth
识别结果：这张图片中的数字是 -> 3