开篇之前说明一下：本文纯粹是技术交流和探讨，所用数据为非公开数据集，仅限于学习，不可用以商业和其他用途。

 一、项目目标

• 通过构建一个简单的 CNN 神经网络，实现对 数字验证码（如 “7384”） 的识别。
• 熟悉图像分类流程：数据生成、图像预处理、网络设计、训练评估与预测。
• 掌握卷积神经网络在图像识别中的基本应用方法。

二、整体思路

实现数字验证码识别目标的整体思路是：(1)验证码图像数据准备-->(2)数据集制作（为了简单，将每张验证码图像中的4个数字切割开来，形成每张图只有1个数字的训练集，并且文件名的第一个字符为训练图的数字，即标签。）-->(3)编写数据处理方法、模型、训练方法、测试方法-->(4)训练和保存模型->(5)测试。

三、实现步骤

步骤一：验证码图像数据准备

本例准备了1143张验证码图像，切割成4572张单个数字的训练图像。

步骤二：数据集制作

如上文所属，本例的训练集为4572张单个数字的训练图像，并且图像文件名的第一个字符为标签，如下图所示：

步骤三：编写实现代码

        1.编写数据加载模块

        构建一个用于数字验证码识别任务的数据加载模块，并通过 PyTorch 提供的数据接口 Dataset 和 DataLoader 对图像数据进行封装和加载。

        具体过程如下：

---

        （1）导入必要的库：

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

        这些库的功能：

        torch.utils.data.Dataset: PyTorch 自定义数据集基类

        DataLoader: 用于批量加载数据、打乱顺序

        transforms: 图像预处理（如灰度、归一化）

        PIL.Image: 用于打开图像文件

        os: 文件路径操作

---

        （2）自定义数据集类 DigitDataset：

class DigitDataset(Dataset):

        这个类继承了 PyTorch 的 Dataset，是用户自定义的数据集类。

        初始化方法：

def __init__(self, data_dir):

        data_dir 是图像数据所在的目录（如 "data"）。

        图像预处理流程：

self.transform = transforms.Compose([transforms.Grayscale(),         # 转换为灰度图（1通道）transforms.ToTensor(),          # 转换为张量并归一化为 [0, 1]transforms.Normalize(0.5, 0.5)  # 再归一化为 [-1, 1]
])

        获取所有图像文件名：

self.file_list = os.listdir(data_dir)

        将文件夹中所有图片文件名收集到列表中。

---

        （3）数据集长度：

def __len__(self):return len(self.file_list)

        告诉 PyTorch 这个数据集总共有多少样本。

---

        （4） 获取某个样本：

def __getitem__(self, idx):

        PyTorch 在训练时会调用这个方法来获取第 idx 个样本。

filename = self.file_list[idx]
img_path = os.path.join(self.data_dir, filename)
image = Image.open(img_path)

         获取图像路径并打开图像。

label = int(filename[0])  # 文件名第一个字符是标签

         将文件名的**第一个字符（数字）**作为标签，例如："7abc.png" 的标签是 7。

return self.transform(image), label

         返回的是图像张量和对应的标签整数。

---

          （5） 创建数据加载器：

dataset = DigitDataset("data")
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

        加载器的作用是：使用 "data" 文件夹下的数据创建一个数据集。用 DataLoader 创建批次大小为 32 的训练集，每个 epoch 自动打乱顺序。

---

        总结一下：数据加载模块实现了一个用于数字图像分类任务（如验证码识别）的数据集加载器，图像标签从文件名中自动提取，图像被统一预处理为灰度、张量、归一化格式，方便输入到神经网络中进行训练。DigitDataset的完整代码如下：

from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import osclass DigitDataset(Dataset):def __init__(self, data_dir):self.data_dir = data_dirself.transform = transforms.Compose([transforms.Grayscale(),         # 转换为灰度图transforms.ToTensor(),          # 转为张量 [0,1]transforms.Normalize(0.5, 0.5)  # 归一化到 [-1,1]])self.file_list = os.listdir(data_dir)def __len__(self):return len(self.file_list)def __getitem__(self, idx):filename = self.file_list[idx]img_path = os.path.join(self.data_dir, filename)image = Image.open(img_path)label = int(filename[0])  # 文件名第一个字符为标签return self.transform(image), label# 创建数据集和数据加载器
dataset = DigitDataset("data")
train_loader = DataLoader(dataset, batch_size=32, shuffle=True)

---

        2.编写模型

        这里定义一个用于数字识别（0~9）分类的卷积神经网络模型 DigitNet，使用的是 PyTorch 框架中的 nn.Module。

        下面我们逐层解读它的结构和作用：

---

        📦 模型整体结构概览

DigitNet(└── Conv2d(1, 2, 3, padding=1)  -> ReLU -> MaxPool2d(2)└── Conv2d(2, 4, 3, padding=1)  -> ReLU└── Flatten -> Linear(168, 128) -> ReLU└── Linear(128, 10)
)

---

        🔍 详细分层解读

        1️⃣ 输入层

nn.Conv2d(1, 2, kernel_size=3, padding=1)

                 输入通道数 1：灰度图像

                输出通道数 2：提取两个特征图

                卷积核大小 3x3

                Padding=1 保证输入输出尺寸不变
        👉 输出尺寸仍为 13x15（本例中，训练输入图像大小为 13x15）

 

nn.ReLU()

         激活函数，增加非线性能力

 

nn.MaxPool2d(2)

         池化操作：尺寸缩小为 1/2

        输入 13x15 → 输出约为 6x7（向下取整）

---

        2️⃣ 第二个卷积层

nn.Conv2d(2, 4, kernel_size=3, padding=1)

         输入通道数 2，输出通道数 4

        卷积后尺寸仍为 6x7

        特征图变为 4 个

nn.ReLU()

---

        3️⃣ 展平层（Flatten）

nn.Flatten()

         将多维特征图转换为一维向量

        展平后的尺寸为：

6(height) * 7(width) * 4(channels) = 168

---

        4️⃣ 全连接层

nn.Linear(6*7*4, 128)

         输入维度：168

        输出维度：128，特征压缩

 

nn.ReLU()

         激活函数

 

nn.Linear(128, 10)

         输出层：10个神经元，对应 数字类别 0~9

---

        🧠 forward 方法

def forward(self, x):return self.model(x)

         将输入图像 x 传入 self.model 中按顺序执行模型结构，输出是10维的向量，表示每个数字的分类概率（一般用于交叉熵分类损失）。

---

        📈 模型输出解释

输出为：

tensor([[0.1, 0.05, ..., 0.15]])  # 共10个数

         每个值表示图像属于某个数字（0~9）的预测得分（在训练时会搭配 CrossEntropyLoss，不需要手动加 softmax）。

---

        ✅ 总结：该模型是一个轻量级的 CNN 分类器，用于识别尺寸约为 13×15 的灰度图像中的单个数字，最终输出一个 10 类的分类结果。模型 DigitNet的完整代码如下：

import torch.nn as nnclass DigitNet(nn.Module):def __init__(self):super().__init__()self.model = nn.Sequential(nn.Conv2d(1, 2, kernel_size=3, padding=1), # 13x15 -> 13x15nn.ReLU(),nn.MaxPool2d(2),                            # 6x7nn.Conv2d(2, 4, kernel_size=3, padding=1), # 6x7 -> 6x7nn.ReLU(),nn.Flatten(),nn.Linear(6*7*4, 128),nn.ReLU(),nn.Linear(128, 10))def forward(self, x):return self.model(x)

---

        3.编写训练函数

        定义一个用于训练 PyTorch 神经网络模型的函数 train_model。

---

        ✅ 一、总体功能说明

def train_model(model, train_loader, epochs=10):

        该函数的作用是：使用指定的数据加载器和模型，进行多轮（epoch）训练，并返回训练完成后的模型。

输入参数：

• model: 待训练的神经网络模型（如你之前定义的 DigitNet）

• train_loader: 数据加载器（封装了训练数据）

• epochs: 训练轮数，默认 10

输出结果：

• 返回训练完成后的 model

---

        🔍 二、代码编写逻辑

        1️⃣ 确定训练设备（CPU 或 GPU）

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

• 自动检测是否有可用的 GPU，如果有则用 cuda 加速，否则退回到 cpu。

• 将模型移到设备上（很关键！）

---

        2️⃣ 定义损失函数和优化器

criterion = nn.CrossEntropyLoss()

• 使用交叉熵损失函数，适用于多分类任务（如 0~9 数字识别）

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

• 使用 Adam 优化器（自动调整学习率），学习率设为 0.001

---

        3️⃣ 模型训练模式

model.train()

• 设置模型为“训练模式”，激活 Dropout 和 BatchNorm（如果有）

---

        4️⃣ 外层循环：遍历每个 epoch

for epoch in range(epochs):running_loss = 0.0

• running_loss：用于记录每个 epoch 的累计损失

---

        5️⃣ 内层循环：遍历每个 batch

for images, labels in train_loader:images, labels = images.to(device), labels.to(device)

• 加载一个 batch 的图像和标签，并移动到指定设备（GPU/CPU）

optimizer.zero_grad()        # 梯度清零
outputs = model(images)      # 正向传播
loss = criterion(outputs, labels)  # 计算损失
loss.backward()              # 反向传播，计算梯度
optimizer.step()             # 更新权重

这一套就是典型的 PyTorch 训练步骤：
清零 → 前向传播 → 计算损失 → 反向传播 → 更新参数

running_loss += loss.item()

• .item() 获取损失的浮点数，累加总损失

---

        6️⃣ 每个 epoch 打印平均损失

print(f"Epoch {epoch + 1}, Loss: {running_loss / len(train_loader):.4f}")

• 将这个 epoch 的平均损失打印出来，方便观察训练情况。

---

        7️⃣ 返回训练好的模型

return model

---

📈 总结：

        代码实现了一个完整的 PyTorch 训练流程，是一个标准、干净、可复用的模型训练函数模板。适合用于任何图像分类模型（只要输出是 logits 且标签是整数类别 ID）。用于训练神经网络模型的函数 train_model的完整代码如下：

import torch
import torch.nn as nn
def train_model(model, train_loader, epochs=10):device = torch.device("cuda" if torch.cuda.is_available() else "cpu")model = model.to(device)criterion = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=0.001)model.train()for epoch in range(epochs):running_loss = 0.0for images, labels in train_loader:images, labels = images.to(device), labels.to(device)optimizer.zero_grad()outputs = model(images)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()print(f"Epoch {epoch + 1}, Loss: {running_loss / len(train_loader):.4f}")return model

---

        4.编写预测函数

        实现了一个 数字验证码图像的预测函数 predict，其核心目的是：

  加载训练好的模型参数，读取指定图像，并输出预测的数字类别（0～9 之间的某个整数）。

---

        ✅ 总体功能：

def predict(model_path, image_path) -> int

        输入：

• model_path: 训练好的模型参数文件（如 .pt 或 .pth）

• image_path: 一张数字图像文件路径（如 "7abc.png"）

        输出：

• 返回图像中预测的数字（0～9 的某一个 int）

---

        🔍 编写逻辑

        1️⃣ 模型初始化与加载

model = DigitNet()
model.load_state_dict(torch.load(model_path))

• 初始化模型结构（必须和训练时一致）

• 加载模型参数（state_dict 是权重的字典）

---

        2️⃣ 设备选择

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

• 自动判断是否使用 GPU

• 模型放到 GPU 或 CPU 上

---

        3️⃣ 图像预处理

transform = transforms.Compose([transforms.Resize((15, 13)),       # 调整图像大小为 (高15, 宽13)transforms.Grayscale(),            # 转为灰度图（1 通道）transforms.ToTensor(),             # 转为 Tensor 格式并归一化到 [0,1]transforms.Normalize(0.5, 0.5)     # 再归一化到 [-1,1]，与训练一致！
])

• 这一部分非常重要，必须和训练时的图像处理完全一致，否则模型可能无法正确识别。

---

        4️⃣ 模型评估模式

model.eval()

• 设置为推理模式：关闭 Dropout、BatchNorm 的训练行为。

---

        5️⃣ 不计算梯度（节省资源）

with torch.no_grad():

• 在推理阶段关闭反向传播计算，提升效率、节省内存。

---

        6️⃣ 图像读取与转换

image = Image.open(image_path)
tensor = transform(image).unsqueeze(0).to(device)

• 使用 PIL 打开图像

• 用 transform 进行预处理

• unsqueeze(0)：在第 0 维添加 batch 维度（变为 [1, 1, 15, 13]）

---

        7️⃣ 模型推理与预测

output = model(tensor)                # 模型输出 logits (形如 [1,10])
predicted = torch.argmax(output, dim=1)  # 取最大值所在的类别索引

• torch.argmax(..., dim=1) 表示在分类维度上找到概率最大的类别（0～9）

---

        8️⃣ 返回预测结果

return predicted.item()

• .item() 将张量转换为 Python 的 int 类型

---

        📌 示例使用：

digit = predict("model.pth", "data/7abc.png")
print("预测结果为：", digit)

---

        ✅ 总结：

模块	作用
DigitNet()	构造模型结构
load_state_dict()	加载训练权重
transform()	图像预处理
model.eval()	设置为评估模式
torch.no_grad()	关闭梯度计算
argmax(output)	选出最可能的数字标签

        预测函数 predict的完整代码如下：

import torch
#from mymodel import DigitNet
from model import DigitNet
from torchvision import transforms  # 正确导入图像变换模块
from PIL import Image               # 标准PIL导入方式#切割好的单数字的验证码图像预测
def predict(model_path, image_path):# 初始化模型结构model = DigitNet()# 加载参数model.load_state_dict(torch.load(model_path))# 设备选择device = torch.device("cuda" if torch.cuda.is_available() else "cpu")print(device)model = model.to(device)# 图像预处理（需与训练时一致）transform = transforms.Compose([transforms.Resize((15, 13)),transforms.Grayscale(),transforms.ToTensor(),transforms.Normalize(0.5, 0.5)])# 执行预测model.eval()with torch.no_grad():image = Image.open(image_path)tensor = transform(image).unsqueeze(0).to(device)output = model(tensor)predicted = torch.argmax(output, dim=1)return predicted.item()

---

        步骤四：训练、模型保存和预测

        直接上代码：

from model import DigitNet
from train import train_model
from dataset import  train_loader
import torch
from predict import predictif __name__ == "__main__":# 初始化模型model = DigitNet()# 训练模型trained_model = train_model(model, train_loader, epochs=20)# 保存模型torch.save(trained_model.state_dict(), "digit_model1.pth")# 预测#print(predict('digit_model1.pth', "test/3_39031.jpeg"))

        运行以上代码，即进行模型训练，训练完成后保存模型 digit_model1.pth 。预测的时候则注释掉模型训练和保存模型的代码，取消后面预测代码的注释。

        声明

        本例附上完整的项目代码和数据，仅供学习和交流，不可作为其它用途。