这次我们来看一个很经典的深度学习实战项目:手写数字识别。别急着觉得简单,很多教程只讲到“训练一个 MNIST 模型”就停了,真正上手时卡在数据预处理、图片尺寸不一致、推理接口封装、批量任务落地的人不在少数。这篇文章不绕弯,直接按照“数据预处理 -> 模型训练 -> 状态验证 -> API 封装 -> 批量导出”的顺序,把一个完整的手写数字识别项目从头到尾串起来,并给出可以直接运行的代码。
手写数字识别本身是计算机视觉入门最标准的任务,但它足够承载深度学习里不少核心问题:图片如何统一尺寸、上采样和下采样怎么选、归一化怎么做、卷积网络如何设计、训练好的模型如何部署成接口、批量图片怎么处理。项目基于 PyTorch 和 MNIST 数据集,核心思路可以迁移到身份证号码识别、票据数字识别、验证码识别等场景。更关键的是,这个项目门槛不高,手头没有独立显卡也能在 CPU 上完成训练和推理,集成 FastAPI 之后还能对外提供服务。
1. 核心能力速览
| 能力项 | 说明 |
|---|---|
| 项目类型 | 基于 PyTorch 的图片分类实战项目 |
| 主要功能 | 手写数字 0-9 识别、图片批量识别、API 调用、结果导出 |
| 数据集 | MNIST,每张图片 28x28 灰度图 |
| 训练环境 | Python 3.8 以上,PyTorch 2.x,CPU 可完成训练,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 python=3.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, dtype=np.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, interpolation=cv2.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', train=True, download=True, transform=transform) test_dataset = datasets.MNIST('./data', train=False, download=True, transform=transform) train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True) test_loader = DataLoader(test_dataset, batch_size=256, shuffle=False) # 定义卷积网络 class CNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Sequential( nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.ReLU(), nn.MaxPool2d(2) ) self.conv2 = nn.Sequential( nn.Conv2d(32, 64, kernel_size=3, padding=1), 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(), lr=1e-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(f"Epoch {epoch+1}/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(dim=1) correct += (preds == labels).sum().item() total += labels.size(0) return correct / total acc = evaluate() print(f"Test 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, dtype=np.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(dim=1).item() prob = torch.softmax(output, dim=1).max().item() return pred, prob # 模型加载 model = CNN() model.load_state_dict(torch.load("mnist_cnn.pth", map_location=device)) 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", index=False) print(df) return df batch_predict("test_images", model, device)输出文件用 pandas 写入 xlsx,WPS 表格可以直接打开。这样批量识别结果就和办公场景串起来了:图片数字识别、结果入表、后续报表生成或模板填充,整个过程完全可以自动化。
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_location=device)) 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, dtype=np.float32) / 255.0 tensor = torch.from_numpy(arr).unsqueeze(0).unsqueeze(0).to(device) with torch.no_grad(): output = model(tensor) pred = output.argmax(dim=1).item() prob = torch.softmax(output, dim=1).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, host="0.0.0.0", port=8000)启动接口服务后,用 curl 做一次快速验证,确认服务能正常返回。
curl -X POST http://127.0.0.1:8000/predict \ -F "file=@test_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, files=files, timeout=30) 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_location=torch.device('cpu') |
| API 启动后页面打不开 | 端口被占用或服务未启动 | 检查终端日志和端口 | 更换端口,如 8001 |
| curl 请求报 422 | 上传字段名与接口参数不一致 | 检查 -F 后面的字段名 | 把 file 换成接口定义的名字 |
| 批量任务卡住 | 单张图片处理时间过长或死循环 | 加日志定位处理到哪张图 | 增加单张超时和失败重试 |
| 中文路径读取失败 | PIL 老旧版本无法处理某些中文路径 | 检查报错堆栈 | 使用 OpenCV 读取或复制到纯英文路径 |
数据集下载失败是最常见的问题。MNIST 训练时如果网络不好,download=True 会中断。解决办法是手动从公开渠道下载四个 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 报表生成等办公自动化场景结合,整个链路的价值会比单独一个识别模型大很多。这个项目最值得验证的三个点,一是预处理和训练是否一致,二是单张推理置信度是否正常,三是批量导出表格能否跑通,第一条最容易踩坑,先把它跑顺,后面都是水到渠成的事。