手写数字识别项目

一、准备数据集

首先我们创建一个卷积模型，训练的时候就需要一个原始的数据集，那么数据集哪里来？Pytorch官网其实有一些数据集，数据集地址

我们使用到的数据集是MNIST

导入包

import torch
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor

使用数据集，所有的官方数据集都继承 torch.utils.data.Dataset，如果你没有数据集，那download = True，它会联网下载到你本地。

# label: 数据集传入的标签值
def target_transform(label):return torch.tensor(label)ds = MNIST(root='./data',  # 保存或读取数据的目录train=True,  # 是否加载训练数据集download=False,  # 是否下载数据集transform=ToTensor(),  # 用于转换图片的函数# target_transform=target_transform  # 用于转换标签的函数target_transform=lambda label: torch.tensor(label)  # 直接匿名函数转换成张量
)

测试打印数据

print(len(ds))
print(ds[0])
print(ds[0][0].shape)

二、定义模型

简单的图像识别模型的套路:卷积 -> 激活 -> 池化 -> … -> 卷积 -> 激活 -> 池化 ->展平 -> 全连接层 -> 激活-> … -> 全连接层输出,会将图片缩小的同时增加通道数,当特征图缩小到 10 以内，就结束卷积过程。之后我们会讲到LeNet5模型,这儿我们简单的定义一个模型进行训练。

from torch import nn# 卷积激活池化 模块
class ConvActivatePool(nn.Module):def __init__(self, in_channels, out_channels, kernel_size):super().__init__()# 一般卷积后会选择让图片大小保持不变 进行填充self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding='same')self.relu = nn.ReLU()# 池化在此处提取了特征的同时，让图片下采样了self.pool = nn.MaxPool2d(2)def forward(self, x):x = self.conv(x)x = self.relu(x)y = self.pool(x)return yclass NumberRecognition(nn.Module):def __init__(self):super().__init__()self.cap1 = ConvActivatePool(1, 64, 11)self.cap2 = ConvActivatePool(64, 128, 5)# 分类层self.classifier = nn.Sequential(# 展平nn.Flatten(start_dim=1),# 全连接层nn.Linear(128 * 7 * 7, 2048),nn.ReLU(),nn.Dropout(p=0.3),nn.Linear(2048, 1024),nn.ReLU(),# 输出结果为 10 分类，所以输出层全连接输出 10nn.Linear(1024, 10))# x 形状 (N, C=1, H=28, W=28)def forward(self, x):x = self.cap1(x)# N x 64 x 14 x 14x = self.cap2(x)# N x 128 x 7 x 7# 图片缩小到 10 以内，则停止卷积# 调用分类器，对图片进行分类y = self.classifier(x)return yif __name__ == '__main__':import torchmodel = NumberRecognition()x = torch.rand(16, 1, 28, 28)y = model(x)print(y.shape)

三、模型训练

3.1 导入依赖库

import math
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split, Subset
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from model import NumberRecognition

3.2 设备设置（CPU/GPU 自动选择）

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

3.3 超参数定义

EPOCH = 10          # 训练轮次：整个训练集遍历10次
LR = 1e-2           # 学习率：控制参数更新的步长（1e-2 = 0.01）
BATCH_SIZE = 10     # 批次大小：每次训练用10个样本更新一次参数
val_rate = 0.2      # 验证集比例：从训练集中划分20%作为验证集

3.4数据集准备

1.获取数据集

ds = MNIST(root='./data',        # 数据集保存路径（若不存在会自动创建）train=True,           # 加载训练集（False则加载测试集）download=False,       # 是否自动下载数据集（首次运行需设为True）transform=ToTensor(), # 对图像的变换：PIL→Tensor（0-1归一化+维度调整）target_transform=lambda label: torch.tensor(label) # 对标签的变换：int→Tensor
)

2.划分训练集与验证集

ds_total_len = len(ds)          # 总样本数：MNIST训练集共60000个样本
train_len = int(ds_total_len * (1 - val_rate)) # 训练集样本数：60000×0.8=48000
val_len = ds_total_len - train_len             # 验证集样本数：60000×0.2=12000
train_ds, val_ds = random_split(ds, [train_len, val_len]) # 随机划分

3.创建 DataLoader（按批次加载数据）

# 计算总批次数（向上取整，避免最后一批样本被丢弃）
train_total_batch = math.ceil(train_len / BATCH_SIZE) # 48000/10=4800批
val_total_batch = math.ceil(val_len / BATCH_SIZE)     # 12000/10=1200批# 训练集DataLoader
train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True  # 训练集每次epoch前打乱样本顺序（避免模型记忆样本顺序，提升泛化）
)# 验证集DataLoader
val_dl = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=True  # 验证集打乱无意义（仅计算损失），建议设为False以提高效率
)

3.5模型初始化与断点续训

# 初始化自定义模型（NumberRecognition在model.py中定义，需确保输入输出维度匹配）
model = NumberRecognition()# 尝试加载历史模型参数（支持断点续训）
try:# 加载参数文件（weights_only=True）state_dict = torch.load('./weights/model.pth', weights_only=True)model.load_state_dict(state_dict) # 将参数加载到模型中print('加载模型参数成功')
except:# 若文件不存在（首次训练），打印提示print('未找到模型参数')# 将模型迁移到指定设备（CPU/GPU）
model.to(device)

3.6损失函数与优化器定义

# 损失函数：交叉熵损失（适合多分类任务，如MNIST的10类数字）
loss_fn = nn.CrossEntropyLoss()# 优化器：Adam优化器（常用优化器，结合SGD的动量和RMSprop的自适应学习率）
optimizer = torch.optim.Adam(model.parameters(),  # 需优化的参数（模型的所有权重和偏置）lr=LR,               # 学习率（与超参数一致）weight_decay=1e-4    # L2正则化（权重衰减，防止模型参数过大导致过拟合）
)

3.7训练函数（train ()）

# 全局变量：累计训练损失和批次数量（用于计算平均损失）
train_total_loss = 0.
train_count = 0def train():global train_total_loss, train_count # 声明使用全局变量print('开始训练')model.train() # 将模型设为“训练模式”（关键！启用Dropout/BatchNorm更新）# 遍历训练集DataLoader，每次取一个批次for i, (images, labels) in enumerate(train_dl):# 1. 将数据迁移到指定设备（与模型设备一致）images, labels = images.to(device), labels.to(device)# 2. 清空上一轮的梯度（PyTorch梯度会累加，不清空会导致梯度错误）optimizer.zero_grad()# 3. 前向传播：模型预测输出y_pred = model(images) # 输出形状：(BATCH_SIZE, 10)，每一行是10个类的得分# 4. 计算损失（预测值与真实标签的差距）loss = loss_fn(y_pred, labels)# 5. 累计损失和批次数量（用于后续计算平均损失）train_total_loss += loss.item() # loss是Tensor，用.item()转为Python数值train_count += 1# 6. 反向传播：计算参数梯度（自动微分核心）loss.backward()# 7. 优化器更新参数（根据梯度调整权重和偏置）optimizer.step()# 每100个批次打印一次训练进度（避免打印过于频繁）if (i + 1) % 100 == 0:avg_loss = train_total_loss / train_countprint(f'BATCH: [{i + 1}/{train_total_batch}]; loss: {avg_loss:.4f}')# 返回本轮训练的平均损失（用于epoch结束时打印）return train_total_loss / train_count

3.8验证函数（valid ()）

def valid():# 局部变量：累计验证损失和批次数量（每轮验证重新初始化，避免与训练混淆）val_total_loss = 0.val_count = 0print('开始验证')model.eval() # 将模型设为“评估模式”（关键！禁用Dropout/BatchNorm更新）# 禁用梯度计算（验证阶段无需反向传播，节省内存和时间）with torch.no_grad():# 遍历验证集DataLoaderfor i, (images, labels) in enumerate(val_dl):# 1. 数据迁移到指定设备images, labels = images.to(device), labels.to(device)# 2. 前向传播（无梯度计算）y_pred = model(images)# 3. 计算验证损失loss = loss_fn(y_pred, labels)val_total_loss += loss.item()val_count += 1# 每100个批次打印验证进度if (i + 1) % 100 == 0:avg_loss = val_total_loss / val_countprint(f'BATCH: [{i + 1}/{val_total_batch}]; loss: {avg_loss:.4f}')# 返回本轮验证的平均损失return val_total_loss / val_count

3.9主训练循环（多轮训练与验证）

# 遍历所有训练轮次
for epoch in range(EPOCH):print(f'\nEPOCH: [{epoch + 1}/{EPOCH}]') # 打印当前轮次（从1开始更直观）# 1. 训练本轮并获取训练平均损失train_loss = train()# 2. 验证本轮并获取验证平均损失val_loss = valid()# 3. 打印本轮训练结果print(f'EPOCH END; train loss: {train_loss:.4f}; val loss: {val_loss:.4f}')# 训练结束后，保存最终模型参数（覆盖原有文件）
torch.save(model.state_dict(), './weights/model.pth')
print('\n模型参数已保存至 ./weights/model.pth')

四、模型训练完整代码

import math
import torch
from torch import nn
from torch.utils.data import DataLoader, random_split, Subset
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
from model import NumberRecognitiondevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')EPOCH = 10
LR = 1e-2
BATCH_SIZE = 10
val_rate = 0.2ds = MNIST(root='./data', train=True, download=False, transform=ToTensor(),target_transform=lambda label: torch.tensor(label))ds_total_len = len(ds)
train_len = int(ds_total_len * (1 - val_rate))
val_len = len(ds) - train_len
train_ds, val_ds = random_split(ds, [train_len, val_len])train_total_batch = math.ceil(train_len / BATCH_SIZE)
val_total_batch = math.ceil(val_len / BATCH_SIZE)train_dl = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
val_dl = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=True)model = NumberRecognition()
try:state_dict = torch.load('./weights/model.pth', weights_only=True)model.load_state_dict(state_dict)print('加载模型参数成功')
except:print('未找到模型参数')model.to(device)loss_fn = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=1e-4)train_total_loss = 0.
train_count = 0def train():global train_total_loss, train_countprint('开始训练')model.train()for i, (images, labels) in enumerate(train_dl):# 3. 将数据放到设备上images, labels = images.to(device), labels.to(device)optimizer.zero_grad()y = model(images)loss = loss_fn(y, labels)train_total_loss += loss.item()train_count += 1loss.backward()optimizer.step()if (i + 1) % 100 == 0:print(f'BATCH: [{i + 1}/{train_total_batch}]; loss: {train_total_loss / train_count}')return train_total_loss / train_countdef valid():val_total_loss = 0.val_count = 0print('开始验证')model.eval()with torch.no_grad():for i, (images, labels) in enumerate(val_dl):images, labels = images.to(device), labels.to(device)y = model(images)loss = loss_fn(y, labels)val_total_loss += loss.item()val_count += 1if (i + 1) % 100 == 0:print(f'BATCH: [{i + 1}/{val_total_batch}]; loss: {val_total_loss / val_count}')return val_total_loss / val_countfor epoch in range(EPOCH):print(f'EPOCH: [{epoch + 1}/{EPOCH}]')train_loss = train()val_loss = valid()print(f'EPOCH END; train loss: {train_loss}; val loss: {val_loss}')torch.save(model.state_dict(), './weights/model.pth')

五、总结流程

1. 加载 MNIST 公开手写数字数据集（训练集）
2. 划分训练集与验证集（用于监控过拟合）
3. 加载自定义的数字识别模型（NumberRecognition），支持断点续训（加载历史参数）
4. 定义训练 / 验证流程，使用交叉熵损失和 Adam 优化器训练模型
5. 训练完成后保存模型参数，便于后续推理或继续训练。