ARTICLE DETAIL

资讯详情

深耕郑州网站建设与运营推广的一线实战洞察。

手写数字识别Python实战:MNIST模型训练与推理脚本封装

手写数字识别Python实战:MNIST模型训练与推理脚本封装 简介这份资源面向Python初学者与深度学习入门者聚焦手写数字识别这一计算机视觉经典任务帮助读者理解卷积神经网络如何自动提取图像特征并完成分类。包内共13个文件以10张jpg手写数字样本图、1个py主程序、1个hdf5训练权重文件为主另含系统隐藏文件压缩包约6.49MB体积轻便便于快速下载与本地运行。目前已有289人学习下载适合作为课程实验或自学练手项目。资源围绕MNIST数据集展开涵盖数据加载与预处理、CNN模型定义、编译训练、测试评估与预测全流程并涉及数据增强、正则化、Dropout及学习率调度等优化思路配套的hdf5权重可直接加载推理jpg样本可用于直观验证识别效果py脚本则展示了从建模到结果可视化的完整实现路径帮助读者在动手实践中掌握TensorFlow、Keras、NumPy与Matplotlib的协同使用并为进一步尝试ResNet、VGG等更复杂模型打下基础。1. 手写数字识别 Python 实战从 MNIST 到本地可跑通的推理脚本很多人第一次接触图像分类都是从手写数字识别开始的。原因很直接MNIST 数据集干净、体量小、类别只有 0 到 9一张图 28×28 灰度用一台普通笔记本的 CPU 就能在几分钟内跑完一轮训练。但真正把它落到“能用”的程度坑并不在模型结构而在数据读取、环境配置、输入预处理和推理接口这几处。我见过太多人训练准确率刷到 99%结果拿一张自己写的数字图片去预测输出全是 3 或者 8这就是典型的“训练能跑、落地翻车”。这篇内容面向两类人一是刚配好 Python 环境、想找一个完整项目练手的入门者二是已经会调库、但想把 MNIST 模型封装成一个可复用推理脚本的工程师。整条路径我会按“数据怎么来 → 模型怎么搭 → 训练怎么稳 → 推理怎么接真实图片 → 坑在哪”推进代码全部用 PyTorch 写环境用 VS Code 或 PyCharm 都能跟。你不需要 GPU也不需要下载任何来路不明的“免费 Python 源码大全”所有依赖都是公开库。2. 数据与环境的准备MNIST 读取、Python 安装与依赖锁定2.1 MNIST 数据到底长什么样为什么它适合做第一个识别项目MNIST 全称是 Modified National Institute of Standards and Technology database训练集 60000 张测试集 10000 张每张图是 28×28 像素的灰度图像素值 0 到 255标签是 0 到 9 的整数。它的图像已经做过居中、归一化尺寸处理所以省掉了目标检测里最烦的标注和裁剪环节。这也是它适合入门的原因你只需要关心“分类”这一件事。但要注意一个容易被忽略的点MNIST 的图片是黑底白字。如果你自己用画图工具写一个白底黑字的数字直接喂给模型分布完全反了预测必然乱。这个坑后面第 5 章会专门讲怎么处理。数据加载用torchvision.datasets.MNIST最省事它会自动下载并缓存。国内网络下载偶尔慢可以手动把四个压缩包放到./data/MNIST/raw/目录下文件名保持train-images-idx3-ubyte这类原始命名downloadFalse就能直接读。2.2 Python 环境配置安装、虚拟环境与依赖版本不管你用 Windows 还是 Linux第一步都是确认 Python 版本。PyTorch 目前对 3.9 到 3.11 支持最稳3.12 部分旧版本轮子还没跟上。用python --version看一眼如果是 3.8 以下建议升级。# 创建独立虚拟环境避免污染全局包 python -m venv mnist_env # Windows 激活 mnist_env\Scripts\activate # Linux / macOS 激活 source mnist_env/bin/activate # 安装依赖版本锁定避免 API 变动 pip install torch2.1.0 torchvision0.16.0 --index-url https://download.pytorch.org/whl/cpu pip install numpy1.26.4 pillow10.2.0 matplotlib3.8.2这里选 CPU 版本的 torch 是因为 MNIST 用 CPU 训练一轮不到 20 秒没必要折腾 CUDA。如果你后续要换更复杂的模型再把--index-url换成对应 CUDA 版本的源即可。pillow是用来做真实图片读取和缩放的matplotlib用来可视化预测结果调试阶段很有用。提示VS Code 里配置 Python 环境时按CtrlShiftP输入Python: Select Interpreter选中mnist_env里的解释器否则终端装好了、编辑器里还是报找不到 torch。2.3 用 DataLoader 把数据喂进模型batch、shuffle 与归一化参数数据管道写错后面训练再久也白搭。核心就三件事转成 Tensor、归一化、分批。import torch from torchvision import datasets, transforms from torch.utils.data import DataLoader # 训练集变换转Tensor 标准化 train_transform transforms.Compose([ transforms.ToTensor(), # 像素从0-255压到0-1 transforms.Normalize((0.1307,), (0.3081,)) # MNIST全局均值和标准差 ]) test_transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) train_set datasets.MNIST(root./data, trainTrue, downloadTrue, transformtrain_transform) test_set datasets.MNIST(root./data, trainFalse, downloadTrue, transformtest_transform) train_loader DataLoader(train_set, batch_size64, shuffleTrue, num_workers0) test_loader DataLoader(test_set, batch_size256, shuffleFalse, num_workers0)Normalize里的0.1307和0.3081是 MNIST 训练集统计出来的全局均值和标准差不是随便填的。用这两个值能让输入分布接近标准正态收敛更快。batch_size训练用 64 是经典配置显存或内存吃紧就降到 32测试用 256 是为了快不影响结果。num_workers在 Windows 上建议设 0设大了容易卡在子进程启动上这是血泪经验。3. 模型搭建与训练CNN 结构选择、损失函数与训练循环3.1 为什么用两层卷积而不是全连接参数量与感受野的取舍最朴素的做法是把 28×28 拉平成 784 维接全连接层。这样也能到 97% 左右但有两个问题一是参数量大784×512 就是 40 万参数二是丢掉了空间结构相邻像素的关系没被利用。卷积核在局部窗口上滑动天然适合图像。我一般用两层卷积加两层全连接的结构参数量控制在 100 万以内测试集准确率能稳定到 99% 以上。第一层卷积用 32 个 3×3 核第二层用 64 个 3×3 核中间各接一次 2×2 最大池化把 28×28 降到 14×14 再到 7×7。最后拉平成 64×7×73136 维接 128 维全连接再输出 10 类。import torch.nn as nn import torch.nn.functional as F class MnistCNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Conv2d(1, 32, kernel_size3, padding1) self.conv2 nn.Conv2d(32, 64, kernel_size3, padding1) self.pool nn.MaxPool2d(2, 2) self.fc1 nn.Linear(64 * 7 * 7, 128) self.fc2 nn.Linear(128, 10) self.dropout nn.Dropout(0.25) def forward(self, x): x self.pool(F.relu(self.conv1(x))) # 28x28 - 14x14 x self.pool(F.relu(self.conv2(x))) # 14x14 - 7x7 x x.view(-1, 64 * 7 * 7) # 拉平 x F.relu(self.fc1(x)) x self.dropout(x) return self.fc2(x)padding1是为了让 3×3 卷积后尺寸不变否则每卷一次掉 2 像素两层下来只剩 24×24后面全连接维度就得重算。Dropout(0.25)放在全连接之后抑制过拟合MNIST 上不加也能跑但加了测试集更稳。3.2 训练循环的四个关键参数学习率、优化器、轮数与日志训练代码本身不长但参数设错就会看到 loss 不降或者准确率震荡。import torch.optim as optim device torch.device(cuda if torch.cuda.is_available() else cpu) model MnistCNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) for epoch in range(5): model.train() running_loss 0.0 for 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() # 每轮结束在测试集上验证 model.eval() correct total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) preds model(images).argmax(dim1) correct (preds labels).sum().item() total labels.size(0) print(fEpoch {epoch1}, Loss {running_loss/len(train_loader):.4f}, fTest Acc {correct/total:.4f})学习率1e-3是 Adam 的常用起点MNIST 上 3 轮就能到 99%。轮数设 5 是留余量实际第 3 轮后提升很小。optimizer.zero_grad()必须放在反向传播前忘了这行梯度会累加loss 直接起飞。model.eval()和torch.no_grad()在验证时都要加前者关掉 dropout后者省内存。3.3 训练完把权重存下来state_dict 与加载方式训练完不保存下次还得重跑。保存用state_dict只存参数不存结构文件小、加载灵活。# 保存 torch.save(model.state_dict(), mnist_cnn.pth) # 加载时先实例化同结构模型再灌参数 model MnistCNN() model.load_state_dict(torch.load(mnist_cnn.pth, map_locationcpu)) model.eval()map_locationcpu是给在 GPU 上训练、CPU 上推理的场景用的不加会报找不到 cuda 设备。文件大概 3MB 左右方便随项目一起分发。4. 推理接口封装把模型变成能接收真实图片的函数4.1 真实图片和 MNIST 的分布差异黑底白字与尺寸问题训练完直接拿手机拍一张纸上的数字预测大概率错。原因有三个一是照片是白底黑字MNIST 是黑底白字二是照片有背景噪声、光照不均三是尺寸不是 28×28。推理前必须做对齐处理。处理顺序我一般这样灰度化 → 二值化 → 反色如果背景是白的→ 裁剪到数字外接框 → 缩放到 20×20 → 放到 28×28 画布居中。这套流程和 MNIST 原始制作方式接近能显著提升真实图片的命中率。4.2 用 PIL 做预处理的完整函数from PIL import Image, ImageOps import numpy as np def preprocess_image(img_path): # 读成灰度图 img Image.open(img_path).convert(L) # 自适应二值化大于阈值变黑小于变白 arr np.array(img) # 判断背景如果四角偏白说明是白底黑字需要反色 corners [arr[0,0], arr[0,-1], arr[-1,0], arr[-1,-1]] if np.mean(corners) 127: arr 255 - arr # 二值化 arr np.where(arr 100, 0, 255).astype(np.uint8) # 找非零区域裁剪 ys, xs np.nonzero(arr) if len(xs) 0: raise ValueError(图片里没找到有效笔画) arr arr[ys.min():ys.max()1, xs.min():xs.max()1] # 缩放到20x20再贴到28x28中心 img Image.fromarray(arr).resize((20, 20), Image.LANCZOS) canvas Image.new(L, (28, 28), 0) canvas.paste(img, ((28-20)//2, (28-20)//2)) # 转Tensor并归一化 tensor transforms.ToTensor()(canvas) tensor transforms.Normalize((0.1307,), (0.3081,))(tensor) return tensor.unsqueeze(0) # 加batch维度np.where(arr 100, 0, 255)这个阈值不是固定的光照强的照片可以调到 150。四角判断背景色是个经验做法比固定反色稳。裁剪后缩放到 20×20 再居中贴到 28×28是复现 MNIST 原始预处理的关键一步少了这步准确率会掉一截。4.3 推理函数与批量预测def predict(model, img_path): model.eval() tensor preprocess_image(img_path) with torch.no_grad(): logits model(tensor) prob torch.softmax(logits, dim1) pred logits.argmax(dim1).item() conf prob[0, pred].item() return pred, conf # 用法 model MnistCNN() model.load_state_dict(torch.load(mnist_cnn.pth, map_locationcpu)) label, confidence predict(model, my_digit.png) print(f预测结果: {label}, 置信度: {confidence:.4f})softmax把 logits 转成概率方便看置信度。置信度低于 0.6 的时候我一般会提示用户重拍而不是硬给一个结果。这个阈值可以根据业务容忍度调。5. 避坑与排查训练能跑但推理翻车的五个真实原因5.1 现象训练准确率 99%自己写的数字全预测错原因白底黑字没反色或者没做居中缩放。MNIST 的像素分布和你随手拍的图差异太大模型没见过这种输入。解决用 4.2 的预处理函数重点检查四角背景判断和 20×20 居中这两步。可以先拿一张 MNIST 测试集里的图走一遍预处理确认输出和原始图一致再换真实图片。5.2 现象loss 一直是 2.3 左右不降原因optimizer.zero_grad()漏了或者标签和输出维度对不上。CrossEntropyLoss 要求输出是 [batch, 10] 的 logits标签是 [batch] 的 0-9 整数。如果标签被 one-hot 成 [batch, 10]loss 会算错但不报错。解决检查标签形状print(labels.shape)应该是torch.Size([64])。另外确认zero_grad在loss.backward()之前。5.3 现象Windows 上 DataLoader 卡住不动原因num_workers大于 0 时Windows 的子进程 spawn 机制和 PyTorch 的迭代器容易死锁。解决把num_workers设为 0或者把训练代码包在if __name__ __main__:里。这是平台差异不是代码逻辑问题Linux 上同样的代码设 4 都没事。5.4 现象加载模型时报 Unexpected key(s) in state_dict原因保存和加载的模型结构不一致比如保存时用了nn.DataParallel加载时没用key 会多一个module.前缀。解决要么加载时也用DataParallel包一层要么手动去掉前缀state_dict {k.replace(module., ): v for k, v in state_dict.items()}。我一般训练脚本里就不加DataParallel单卡没必要。5.5 现象预测结果在几个数字之间反复横跳原因输入图片噪声太大或者二值化阈值不合适导致笔画断裂或粘连。解决把预处理中间结果用 matplotlib 画出来看确认二值化后的图是清晰的数字轮廓。阈值不要写死 100可以先用 Otsu 方法自动算或者简单点取图像均值的 0.6 倍作为阈值。6. 进阶技巧用测试时增强和置信度过滤把真实场景准确率再提一档模型训练完只是起点真实场景的图片千奇百怪单次推理不稳。我常用的两个技巧是测试时增强TTA和置信度过滤。TTA 的思路是对同一张输入图做几种轻微变换分别推理再把概率平均。MNIST 上有效的变换是左右平移 1 像素和旋转 ±5 度。代码不复杂def predict_tta(model, img_path): model.eval() tensor preprocess_image(img_path) # [1,1,28,28] probs [] with torch.no_grad(): # 原图 probs.append(torch.softmax(model(tensor), dim1)) # 左右各平移1像素 for shift in [-1, 1]: shifted torch.roll(tensor, shiftsshift, dims3) probs.append(torch.softmax(model(shifted), dim1)) avg_prob torch.stack(probs).mean(dim0) pred avg_prob.argmax(dim1).item() conf avg_prob[0, pred].item() return pred, conftorch.roll在宽度维度上滚动超出边界的像素会绕回来对 MNIST 这种居中图影响很小。三次推理取平均比单次稳代价是推理时间变三倍CPU 上也就几十毫秒完全可接受。置信度过滤是另一道保险。我一般设两个阈值高于 0.9 直接输出0.6 到 0.9 之间输出但标记“建议复核”低于 0.6 直接返回“无法识别”。这样在批量处理表单数字时能把明显有问题的样本挑出来人工看避免错误结果直接进下游。还有一个容易被忽略的点定期用新收集的真实图片做一轮微调。哪怕只有几十张把最后一层全连接的学习率调小到 1e-4训练两三轮模型对你自己场景的适应度会明显提升。这比重新训练整个网络划算得多。我自己踩过最深的坑是早期图省事推理时直接resize((28,28))而不做居中结果模型对写偏的数字完全没辙。后来老老实实按 MNIST 原始预处理走一遍同样的权重真实图片准确率从六成出头提到九成以上。模型没变变的只是输入对齐。所以如果你也卡在“训练很好、推理很烂”先别急着换网络结构把预处理那几行代码逐像素对一遍大概率就找到问题了。希望帮到你。本文还有配套的精品资源点击获取
返回列表