手写数字识别实战:从PyTorch模型训练到FastAPI部署与批量导出
发布时间:2026/9/7 21:00:06 作者:尧图编辑部 阅读量:1,286

这次我们来看一个很经典的深度学习实战项目手写数字识别。别急着觉得简单很多教程只讲到“训练一个 MNIST 模型”就停了真正上手时卡在数据预处理、图片尺寸不一致、推理接口封装、批量任务落地的人不在少数。这篇文章不绕弯直接按照“数据预处理 - 模型训练 - 状态验证 - API 封装 - 批量导出”的顺序把一个完整的手写数字识别项目从头到尾串起来并给出可以直接运行的代码。手写数字识别本身是计算机视觉入门最标准的任务但它足够承载深度学习里不少核心问题图片如何统一尺寸、上采样和下采样怎么选、归一化怎么做、卷积网络如何设计、训练好的模型如何部署成接口、批量图片怎么处理。项目基于 PyTorch 和 MNIST 数据集核心思路可以迁移到身份证号码识别、票据数字识别、验证码识别等场景。更关键的是这个项目门槛不高手头没有独立显卡也能在 CPU 上完成训练和推理集成 FastAPI 之后还能对外提供服务。1. 核心能力速览能力项说明项目类型基于 PyTorch 的图片分类实战项目主要功能手写数字 0-9 识别、图片批量识别、API 调用、结果导出数据集MNIST每张图片 28x28 灰度图训练环境Python 3.8 以上PyTorch 2.xCPU 可完成训练GPU 可选模型结构两层卷积 全连接分类网络推理方式单张图片推理脚本FastAPI 接口推理批量文件夹推理是否支持 API支持使用 FastAPI / Uvicorn是否支持批量支持可遍历文件夹批量识别并输出表格输出形式终端输出、JSON 返回、Excel/WPS 表格适合读者深度学习初学者、计算机视觉入门者、需要做 OCR 预研的开发者表格里的环境属于通用建议实际版本以本机安装为准。MNIST 数据集的公开性很强训练集 60000 张测试集 10000 张不需要额外购买数据。整体项目占用的磁盘空间不大数据集下载后约几十 MB模型权重文件通常只有几 MB。2. 适用场景与使用边界手写数字识别项目最适合的人群是刚接触 PyTorch 的开发者。通过这个项目可以快速理解卷积网络的基本结构、数据集的加载方式、模型训练与保存流程以及一个模型从训练到部署的完整链路。对于准备计算机二级考试或者做 WPS/Excel 数据处理相关工作的人来说它也是一个很好的自动化扩展点识别结果可以直接导出到表格再和文档批量填充、报表生成等流程结合。这个项目的识别边界也必须说清楚。MNIST 模型只适合相对规整的手写数字图片背景简单、数字主体清晰、方向没有大幅度旋转识别效果才稳定。对于复杂自然场景中的数字、倾斜字体、艺术字、多人手写混排或者包含大量噪声的图片直接套用这个模型并不合适需要引入更复杂的预处理和更大的训练数据。再强调一下合规使用手写数字识别、OCR 类能力如果用在身份证号、银行卡号、试卷答案等真实数据上必须确认数据来源合法并且只用于授权范围内的处理。涉及到个人信息、隐私数据、版权素材要在合规前提下使用不要随意采集和保存他人敏感数据。3. 环境准备与前置条件先准备环境。操作系统以 Windows 10/11 或 Linux 都可以macOS 也能跑差别不大。建议使用 conda 创建独立环境避免依赖冲突。基础依赖包括 Python、PyTorch、TorchVision、Pillow、NumPy、OpenCV、pandas、openpyxl、FastAPI 和 Uvicorn。# 创建并激活 Python 环境这里以 conda 为例 conda create -n mnist python3.9 -y conda activate mnist # 安装 PyTorch 系列 # CPU 版本可以直接用 pip 安装 pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu # 安装图像处理与接口服务依赖 pip install pillow numpy opencv-python pandas openpyxl fastapi uvicorn requests如果你本机有 NVIDIA 显卡并且已经配置好 CUDA可以把 torch 安装命令换成对应的 GPU 版本。无论 CPU 还是 GPU项目代码本身是通用的PyTorch 会根据本机环境自动选择设备。硬件方面CPU 训练 MNIST 完全可以接受。28x28 的小图两层卷积网络参数量不大普通家用 CPU 训练几个 epoch 通常几分钟到十几分钟具体时间取决于 CPU 核心数和线程配置。GPU 训练会更快但对显存要求很低哪怕是入门的低显存显卡也没有压力。更稳妥的判断是先用 CPU 把完整流程跑通再考虑 GPU 加速。4. 数据预处理统一尺寸、归一化与插值策略手写数字识别的第一步不是建模型而是把图片处理好。MNIST 原始图片是 28x28 的灰度图但实际使用中我们拿到的图片不可能都是这个尺寸。有些是手机拍出来的大图有些是截图有些是从 Excel 或 WPS 表格里导出的缩略图。所以数据预处理要解决三个核心问题上采样、下采样、归一化。上采样是把小图放大到目标尺寸下采样是把大图缩小到目标尺寸。这个过程必须用插值算法完成。插值算法看似不起眼实际对识别效果影响很大。最近邻插值最快但放大后锯齿感严重双线性插值平滑自然是默认选择双三次插值细节稍好但计算量大一些LANCZOS 插值在图片缩放质量上更高但速度也更慢。插值方式特点适用情况INTER_NEAREST速度快边缘锯齿明显只需要快速预览时使用INTER_LINEAR平滑速度快失真较小通用首选INTER_CUBIC更平滑细节保留更好图片需要高质量缩放时INTER_LANCZOS4质量高计算量较大对质量敏感、图片量不大时下面给出一套预处理代码它同时兼容训练和推理阶段。输入图片会先被转为灰度图然后统一缩放到 28x28再转为 NumPy 数组并除以 255 归一化到 0 到 1 之间。import numpy as np from PIL import Image def preprocess_image(image_path, target_size(28, 28)): 读取图片并统一到目标尺寸返回形状为 (1, 1, H, W) 的归一化数组。 # 读取灰度图 img Image.open(image_path).convert(L) # 使用双线性插值进行上采样或下采样 img img.resize(target_size, Image.BILINEAR) # 转为 NumPy 数组并归一化 arr np.array(img, dtypenp.float32) / 255.0 # 增加通道维和 batch 维方便直接输入模型 arr arr.reshape(1, 1, target_size[0], target_size[1]) return arr如果你更习惯用 OpenCV 处理也可以用 cv2.resize 完成同样的事。这里把两种方案对比一下PIL 的 resize 接口简单read 图片时不太容易因为路径中文报错OpenCV 读取速度快但 BGR 与 RGB 通道需要额外注意。因为手写数字是灰度图通道问题影响不大主要看项目里其它功能是否依赖 OpenCV。如果你后续要继续做形态学处理、轮廓提取、透视变换建议统一使用 OpenCV。import cv2 def preprocess_image_cv2(image_path, target_size(28, 28)): # 读取灰度图 img cv2.imread(image_path, cv2.IMREAD_GRAYSCALE) # 使用双线性插值调整尺寸 img cv2.resize(img, target_size, interpolationcv2.INTER_LINEAR) # 归一化并升维 arr img.astype(np.float32) / 255.0 return arr.reshape(1, 1, target_size[0], target_size[1])这里有一个关键点要提醒训练时如果使用了 torchvision 的 transforms.Normalize推理阶段必须使用相同的均值和标准差否则模型会“看不懂”输入数据。MNIST 常用的均值和标准差是 (0.1307, 0.3081)。上面这份预处理代码是直接除以 255和 torchvision 默认加载方式不完全一样但在大多数情况下仍然能给出不错的识别结果。为了让效果更稳定建议在有训练和推理代码之间保持完全一致的预处理逻辑不要把两套思路混着用。5. 模型训练CNN 结构与训练脚本数据准备好之后进入模型训练阶段。这里选择的是一个精简的卷积神经网络结构上参考了 LeNet 的设计思想两个卷积块提取局部特征每个卷积块包含卷积、ReLU 激活和最大池化最后接两层全连接分类器。这种结构对 MNIST 来说是够用的训练速度快代码直观后续也容易改造成其它模型。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms # 训练数据预处理 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)) ]) # 加载 MNIST 数据集第一次运行会自动下载 train_dataset datasets.MNIST(./data, trainTrue, downloadTrue, transformtransform) test_dataset datasets.MNIST(./data, trainFalse, downloadTrue, transformtransform) train_loader DataLoader(train_dataset, batch_size64, shuffleTrue) test_loader DataLoader(test_dataset, batch_size256, shuffleFalse) # 定义卷积网络 class CNN(nn.Module): def __init__(self): super().__init__() self.conv1 nn.Sequential( nn.Conv2d(1, 32, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2) ) self.conv2 nn.Sequential( nn.Conv2d(32, 64, kernel_size3, padding1), nn.ReLU(), nn.MaxPool2d(2) ) self.fc1 nn.Linear(64 * 7 * 7, 128) self.dropout nn.Dropout(0.25) self.fc2 nn.Linear(128, 10) def forward(self, x): x self.conv1(x) x self.conv2(x) x x.view(x.size(0), -1) x torch.relu(self.fc1(x)) x self.dropout(x) x self.fc2(x) return x device torch.device(cuda if torch.cuda.is_available() else cpu) model CNN().to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr1e-3) # 训练一个 epoch 并返回平均损失 def train_one_epoch(): model.train() total_loss 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() total_loss loss.item() return total_loss / len(train_loader) # 训练 5 个 epoch输出损失变化 for epoch in range(5): avg_loss train_one_epoch() print(fEpoch {epoch1}/5, Loss: {avg_loss:.4f}) # 保存模型权重 torch.save(model.state_dict(), mnist_cnn.pth)训练脚本里有两个值得注意的细节。第一个是model.eval()和model.train()的切换虽然这个模型里 Dropout 只在训练时生效但推理前不切换到 eval 模式会导致输出不稳定所以推理脚本里一定记得加model.eval()。第二个是torch.no_grad()推理阶段不需要计算梯度加上它既能省内存也能提升速度。训练完成后建议在测试集上验证一下模型准确率。下面这段代码计算模型在测试集上的整体准确率判断模型是否训练成功。def evaluate(): model.eval() correct 0 total 0 with torch.no_grad(): for images, labels in test_loader: images, labels images.to(device), labels.to(device) outputs model(images) preds outputs.argmax(dim1) correct (preds labels).sum().item() total labels.size(0) return correct / total acc evaluate() print(fTest Accuracy: {acc:.4f})这个结构在 MNIST 测试集上通常能达到 99% 左右实际值会受随机种子、epoch 数量、学习率等因素影响。如果你的结果和这个量级差距很大优先检查数据预处理和归一化参数是否一致。6. 功能测试与效果验证模型训练完第一步验证不是直接上接口而是先跑单张图片推理。从 MNIST 测试集里挑一张图片或者自己在画图软件里写一个数字保存为 PNG都可以作为测试素材。实现推理时要把前面写的预处理函数和模型加载逻辑结合起来。import torch from PIL import Image import numpy as np def infer_single(image_path, model, device): # 预处理 img Image.open(image_path).convert(L) img img.resize((28, 28), Image.BILINEAR) arr np.array(img, dtypenp.float32) / 255.0 tensor torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(device) # 推理 model.eval() with torch.no_grad(): output model(tensor) pred output.argmax(dim1).item() prob torch.softmax(output, dim1).max().item() return pred, prob # 模型加载 model CNN() model.load_state_dict(torch.load(mnist_cnn.pth, map_locationdevice)) model.to(device) pred, prob infer_single(test_digit.png, model, device) print(f预测结果: {pred}, 置信度: {prob:.4f})判断识别成功有两个维度。第一个是预测类别正确第二个是置信度足够高。如果置信度在 0.9 以上说明模型对这个样本很有把握如果置信度只有 0.5 左右说明样本和训练数据分布存在偏差比如笔画太粗、数字旋转、背景有噪点这些情况需要人工复核。单张验证通过后再做批量验证。批量推理的核心是遍历一个文件夹里所有图片逐张识别然后把结果集中输出。这里顺便接上一个很实用的需求把识别结果导出到 Excel/WPS 表格方便后续做数据处理。import os import pandas as pd def batch_predict(image_dir, model, device): results [] for name in sorted(os.listdir(image_dir)): if not name.lower().endswith((.png, .jpg, .jpeg, .bmp)): continue path os.path.join(image_dir, name) pred, prob infer_single(path, model, device) results.append({ 文件名: name, 预测值: pred, 置信度: round(prob, 4) }) df pd.DataFrame(results) df.to_excel(predict_results.xlsx, indexFalse) print(df) return df batch_predict(test_images, model, device)输出文件用 pandas 写入 xlsxWPS 表格可以直接打开。这样批量识别结果就和办公场景串起来了图片数字识别、结果入表、后续报表生成或模板填充整个过程完全可以自动化。7. 接口 API 与批量任务单机脚本可以满足个人使用但很多场景需要把识别能力提供给其它系统调用。最简单的方式是用 FastAPI 写一个推理服务。服务端接收上传图片内部完成预处理和模型推理然后返回 JSON 结果。启动服务后前端、Python 脚本、其它后端服务都可以通过 HTTP 请求调用。先实现推理服务的核心部分。这里为了示例清晰直接在脚本里加载模型、定义处理函数并把图片字节流转为 PIL Image。import io import uvicorn import torch from fastapi import FastAPI, UploadFile, File from PIL import Image import numpy as np app FastAPI() # 加载模型 device torch.device(cuda if torch.cuda.is_available() else cpu) model CNN() model.load_state_dict(torch.load(mnist_cnn.pth, map_locationdevice)) model.to(device).eval() def predict_bytes(image_bytes): # 字节流转为 PIL 图片 img Image.open(io.BytesIO(image_bytes)).convert(L) img img.resize((28, 28), Image.BILINEAR) arr np.array(img, dtypenp.float32) / 255.0 tensor torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): output model(tensor) pred output.argmax(dim1).item() prob torch.softmax(output, dim1).max().item() return {prediction: pred, confidence: round(prob, 4)} app.get(/health) def health(): return {status: ok} app.post(/predict) async def predict_image(file: UploadFile File(...)): content await file.read() return predict_bytes(content) if __name__ __main__: uvicorn.run(app, host0.0.0.0, port8000)启动接口服务后用 curl 做一次快速验证确认服务能正常返回。curl -X POST http://127.0.0.1:8000/predict \ -F filetest_digit.png正常返回结果类似这样{ prediction: 7, confidence: 0.9912 }如果你在写自己的客户端用 requests 调用也一样简单。import requests url http://127.0.0.1:8000/predict files {file: open(test_digit.png, rb)} response requests.post(url, filesfiles, timeout30) print(response.json())接口跑通之后批量任务就变成了一个工程调度问题。最简单的方式是在本地维护一个待处理图片目录脚本依次请求接口拿到结果后写入表格。如果图片量大建议加入失败重试和日志记录。比如某个图片因为格式异常导致接口返回错误不能直接中断整个任务应该记录错误后继续处理下一条。从模型服务角度看也可以把请求改为同步等待批量时控制并发数避免一次性打满服务导致超时。8. 资源占用与性能观察运行这个项目时资源占用是很多读者关心的问题。先说结论MNIST 是一个很小的任务模型结构也很轻量绝大多数普通电脑都可以顺畅运行。CPU 训练时主要压力在 CPU 计算和内存上内存占用通常在 2GB 以内。训练 5 个 epoch 的时间取决于 CPU 型号速度慢可能在十分钟左右速度快可能只需要几分钟这个数字需要结合自己电脑的配置来观察。如果使用 GPU 训练显存占用会非常低。batch size 为 64、图片分辨率只有 28x28显存占用通常在 1GB 以内低显存显卡也可以轻松运行。如果你要观察显存占用可以在训练循环里临时打印torch.cuda.memory_allocated()和torch.cuda.memory_reserved()也可以直接用 NVIDIA 的nvidia-smi命令查看。nvidia-smi影响资源占用的因素主要有四个batch size 决定了单次送入模型的数据量图片分辨率决定卷积层的计算量epoch 数量决定总训练时间数据加载线程数会影响 CPU 占用。这个项目虽然默认图片是 28x28但如果你在扩展时把输入改成 64x64模型参数量和计算量都会明显增加训练时间也会拉长。降低负载的策略很简单。如果显存或内存吃紧把 batch size 调小比如从 64 调到 32 或 16。如果训练时间太长可以先减少 epoch 数量跑通流程后再增加。推理阶段如果要做高并发接口服务可以考虑把模型加载在内存后常驻避免每次请求都重新加载权重。尽量把主机地址绑定到 127.0.0.1 而不是 0.0.0.0可以减少不必要的网络暴露。9. 常见问题与排查方法实际运行这个项目时大概率会遇到几个典型问题。这里整理一份排查表遇到问题可以直接对照。问题现象可能原因排查方式解决方案MNIST 数据集下载失败网络不稳定或下载源不可访问查看下载时的报错信息手动下载数据集并放到 ./data 目录训练时报 shape 不匹配图片尺寸不是 28x28打印输入 tensor 的 shape检查预处理 resize 是否生效推理准确率很低训练和推理预处理不一致对比归一化参数是否相同统一使用相同的预处理代码torch.load 报错权重文件路径不对或设备不匹配检查文件是否存在使用 map_locationtorch.device(cpu)API 启动后页面打不开端口被占用或服务未启动检查终端日志和端口更换端口如 8001curl 请求报 422上传字段名与接口参数不一致检查 -F 后面的字段名把 file 换成接口定义的名字批量任务卡住单张图片处理时间过长或死循环加日志定位处理到哪张图增加单张超时和失败重试中文路径读取失败PIL 老旧版本无法处理某些中文路径检查报错堆栈使用 OpenCV 读取或复制到纯英文路径数据集下载失败是最常见的问题。MNIST 训练时如果网络不好downloadTrue 会中断。解决办法是手动从公开渠道下载四个 gzip 文件放到项目 ./data/MNIST/raw 目录下再重新运行脚本。如果你改成自己的数据集还要注意标签文件是否齐全。另一个容易被忽略的问题是图片归一化不一致。很多教程在训练时使用 transforms.Normalize((0.1307,), (0.3081,))但推理时只把像素除以 255这会导致特征分布完全不同识别准确率明显下降。本文的推理示例使用了除以 255 的简化方案实际如果发现识别变差建议训练和推理都使用 torchvision 的统一预处理。API 服务的问题也值得提前规避。FastAPI 的 uvicorn 默认单进程启动如果批量请求量较大需要修改启动参数增加 worker。但这里要注意模型加载在 worker 进程里会各自占用一份内存所以并发提升会带来内存上升需要根据服务内存大小调整 worker 数量。10. 最佳实践与使用建议项目跑通之后建议按照工程化方式整理目录结构。把代码、数据、模型、图片素材、输出结果分开管理例如src/放代码data/放数据集models/放权重文件inputs/放测试图片outputs/放推理结果。这样项目变大后不容易混乱回滚排查也更方便。训练阶段建议增加数据增强比如随机旋转、随机平移、随机缩放。MNIST 虽然简单但加入这些变换可以让模型对真实手写图片更鲁棒。不过要注意数据增强不能改变数字语义旋转角度控制在 ±15 度以内比较稳妥。推理阶段要重视置信度阈值。不要无条件信任模型的输出可以把置信度低于 0.8 的结果标记为“存疑”交给人工确认。这一步在批量场景里尤其重要否则错误数字一旦直接写入 Excel 或 WPS后续数据处理全都会跟着错。接口服务如果投入生产要把权重文件和启动代码做版本管理。模型迭代后接口返回结果要能追溯是哪个版本模型生成的最好在返回 JSON 里加上模型版本号。另外服务进程要设置访问限制如果是内网使用尽量通过防火墙只允许指定 IP 访问如果是公网必须加认证以免接口被刷。再往后扩展可以考虑把 PyTorch 模型导出为 ONNX 格式然后用 ONNX Runtime 推理。ONNX Runtime 在 CPU 上通常比原始 PyTorch 更快部署时也不需要完整安装 PyTorch占用更小。导出方式也简单import torch model.eval() dummy_input torch.randn(1, 1, 28, 28) torch.onnx.export( model, dummy_input, mnist_cnn.onnx, input_names[input], output_names[output], dynamic_axes{input: {0: batch_size}, output: {0: batch_size}} ) print(ONNX 模型导出完成)导出后可以使用 onnxruntime 加载模型推理速度对轻量模型而言通常有明显改善这对批量任务更友好。识别结果导出到 Excel 后还可以进一步和 WPS 模板批量填充、Word 报表生成等办公自动化场景结合整个链路的价值会比单独一个识别模型大很多。这个项目最值得验证的三个点一是预处理和训练是否一致二是单张推理置信度是否正常三是批量导出表格能否跑通第一条最容易踩坑先把它跑顺后面都是水到渠成的事。