mnist 源码地址
A mnist demo app.
入门搭建开发环境
- 下载 python
- 安装 选择添加环境变量，直接下一步到完成
- 安装 pipenv
pip install pipenv，使用pipenv管理环境非常方便 - 有条件使用vpn，没有可以配置阿里的源。
运行项目
- 创建虚拟环境
pipenv shell - 安装依赖
pipenv install - 训练模型
python train.py - 测试训练的模型
python test.py - 集成显示
python serve.py
验证识别
- 在集成显示的浏览器中输入
http://127.0.0.1:5000/
训练代码
model.py
import torch.nn as nn
# 导入 PyTorch 的神经网络模块并简写为 nn，方便后续使用 nn.Conv2d、nn.Linear 等
class NeuralNetwork(nn.Module):
# 定义一个名为 NeuralNetwork 的神经网络类，继承自 nn.Module
def __init__(self):
# 构造函数：初始化网络层
super().__init__()
# 调用父类构造函数以正确注册子模块
self.conv1 = nn.Conv2d(1, 32, 3, 1)
# 第一层卷积：输入通道 1（灰度图），输出通道 32，卷积核大小 3，步幅 1
self.conv2 = nn.Conv2d(32, 64, 3, 1)
# 第二层卷积：输入通道 32，输出通道 64，卷积核大小 3，步幅 1
self.fc1 = nn.Linear(in_features=64 * 5 * 5, out_features=128)
# 第一个全连接层：输入特征数 64*5*5（前面两次池化后特征图展平），输出 128 个神经元
self.fc2 = nn.Linear(in_features=128, out_features=10)
# 输出层：输入 128，输出 10（对应 MNIST 的 10 个数字类别）
def forward(self, x):
# 定义前向传播函数，输入 x（形状通常为 [batch, 1, 28, 28]）
x = nn.functional.relu(self.conv1(x))
# 先通过 conv1 卷积，然后使用 ReLU 激活函数增加非线性
x = nn.functional.max_pool2d(x, kernel_size=2)
# 对激活后的特征图进行 2x2 的最大池化，减小空间维度
x = nn.functional.relu(self.conv2(x))
# 再通过 conv2 卷积并使用 ReLU 激活
x = nn.functional.max_pool2d(x, kernel_size=2)
# 再次进行 2x2 最大池化，通常从 28x28 经两次 2x2 池化变为 5x5（取整/边界处理依赖于卷积设置）
x = x.view(-1, 64 * 5 * 5)
# 将张量展平以输入全连接层：-1 表示自动计算 batch 大小，其他维度为 64*5*5
x = nn.functional.relu(self.fc1(x))
# 将展平后的张量输入第一个全连接层并接 ReLU 激活
x = self.fc2(x)
# 最后通过输出层得到 logits（未归一化的类别分数）
return x
train.py
# 导入时间模块，用于计算训练耗时
from time import time
# 导入PyTorch库及其核心模块
import torch
import torch.nn as nn # 神经网络模块
import torch.optim as optim # 优化器模块
# 从torchvision导入数据集和数据处理工具
from torchvision import datasets # 预定义数据集
from torch.utils.data import DataLoader # 数据加载器
from torchvision.transforms import ToTensor # 图像转张量的转换器
# 导入自定义的神经网络模型
from model import NeuralNetwork
# 训练函数：执行模型训练的一个epoch
# 参数:
# dataloader: 数据加载器，提供批量数据
# device: 训练设备(cpu或cuda)
# model: 要训练的神经网络模型
# loss_fn: 损失函数
# optimizer: 优化器
def train(dataloader, device, model, loss_fn, optimizer):
model.train() # 将模型设置为训练模式
running_loss = 0.0 # 累计损失值
# 遍历数据加载器中的每个批次
for batch, (inputs, labels) in enumerate(dataloader):
inputs = inputs.to(device) # 将输入数据移至指定设备
labels = labels.to(device) # 将标签移至指定设备
optimizer.zero_grad() # 清除优化器中的梯度
outputs = model(inputs) # 前向传播，获取模型输出
loss = loss_fn(outputs, labels) # 计算损失值
loss.backward() # 反向传播，计算梯度
optimizer.step() # 更新模型参数
running_loss += loss.item() # 累加批次损失
# 打印该epoch的平均损失值
print(f'loss: {running_loss/len(dataloader):>0.3f}')
# 测试函数：评估模型性能
# 参数:
# dataloader: 测试数据加载器
# device: 测试设备
# model: 要测试的模型
def test(dataloader, device, model):
model.eval() # 将模型设置为评估模式
correct = 0 # 正确预测的样本数
total = 0 # 总样本数
with torch.no_grad(): # 禁用梯度计算，提高推理速度并减少内存使用
# 遍历测试数据
for inputs, labels in dataloader:
inputs = inputs.to(device) # 将输入数据移至指定设备
labels = labels.to(device) # 将标签移至指定设备
outputs = model(inputs) # 前向传播，获取模型输出
_, predicted = torch.max(outputs.data, 1) # 获取预测类别
total += labels.size(0) # 更新总样本数
correct += (predicted == labels).sum().item() # 更新正确预测数
# 打印模型准确率
print(f'accuracy: {100.0*correct/total:>0.2f} %')
# 主函数：执行完整的训练流程
def main():
print('loading training data...') # 加载训练数据
# 加载MNIST训练数据集
train_data = datasets.MNIST(
root='./data', # 数据保存路径
train=True, # 加载训练集
download=True, # 如果数据不存在则下载
transform=ToTensor()) # 数据转换为张量
print('loading test data...') # 加载测试数据
# 加载MNIST测试数据集
test_data = datasets.MNIST(
root='./data', # 数据保存路径
train=False, # 加载测试集
download=True, # 如果数据不存在则下载
transform=ToTensor()) # 数据转换为张量
# 创建训练数据加载器，批量大小为64
train_dataloader = DataLoader(train_data, batch_size=64)
# 创建测试数据加载器，批量大小为64
test_dataloader = DataLoader(test_data, batch_size=64)
# 选择训练设备：优先使用GPU(cuda)，否则使用CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device}') # 打印使用的设备
# 实例化模型并移至指定设备
model = NeuralNetwork().to(device)
print(model) # 打印模型结构
# 定义交叉熵损失函数
loss_fn = nn.CrossEntropyLoss()
# 定义Adam优化器，学习率为0.001
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 定义训练轮数
epochs = 5
# 训练循环
for t in range(epochs):
start_time = time() # 记录epoch开始时间
print(f'epoch {t+1} / {epochs}\n--------------------') # 打印当前epoch信息
# 执行训练
train(train_dataloader, device, model, loss_fn, optimizer)
# 执行测试
test(test_dataloader, device, model)
end_time = time() # 记录epoch结束时间
# 打印当前epoch耗时
print(f'time: {end_time-start_time:>0.2f} seconds')
print('done!') # 训练完成
path = 'mnist.pth' # 模型保存路径
torch.save(model.state_dict(), path) # 保存模型参数
print(f'model saved: {path}') # 打印保存信息
# 当作为主程序运行时，执行main函数
if __name__ == '__main__':
main()
测试训练结果
serve.py
# 导入标准库 base64，用于对前端上传的图像进行 base64 解码
import base64
# 导入 PyTorch
import torch
# BytesIO 用于把解码后的字节数据包装成类文件对象，PIL 可以直接 open
from io import BytesIO
# PIL 用于打开和处理图像
from PIL import Image
# Flask 相关：创建应用、处理请求、重定向和返回 JSON
from flask import Flask, request, redirect, jsonify
# torchvision 的 transforms 用于把 PIL 图像转换为 tensor 并做灰度化
from torchvision import transforms
# 从本地文件导入定义的模型类 NeuralNetwork
from model import NeuralNetwork
# 选择设备：如果有可用的 CUDA 则使用 GPU，否则回退到 CPU
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f'using {device}') # 打印当前使用的设备，便于调试
# 实例化模型并移动到选定设备
model = NeuralNetwork().to(device)
# 模型参数文件路径（相对于当前工作目录）
path = './mnist.pth'
# 加载训练好的模型权重到模型中
model.load_state_dict(torch.load(path))
print(f'loaded model from {path}') # 打印已加载模型的路径
print(model) # 打印模型结构，方便检查
# 获取并打印模型参数（state_dict）以便调试（在生产中可去掉）
params = model.state_dict()
print(params)
# 创建 Flask 应用实例
app = Flask(__name__)
@app.route('/')
def index():
"""根路径：重定向到静态前端页面 index.html"""
return redirect('/static/index.html')
@app.route('/api', methods=['POST'])
def api():
"""API 端点：接收前端发送的 base64 编码图像，返回预测结果和概率。
请求 JSON 格式示例：{"image": "<base64字符串>"}
返回 JSON 格式示例：{"result": 7, "probability": 0.98}
"""
# 从请求中解析 JSON 数据（Flask 会将请求体解析为 Python dict）
data = request.get_json()
# 从 data 中读取字段 'image'，并对 base64 文本进行解码得到二进制图像数据
image_data = base64.b64decode(data['image'])
# 使用 BytesIO 把二进制数据包装成文件对象，再交给 PIL 打开为图像
image = Image.open(BytesIO(image_data))
# 定义图像预处理管道：先转为单通道灰度图，再转换为 PyTorch tensor
trans = transforms.Compose([
transforms.Grayscale(1), # 将图像转换为单通道（1 通道）的灰度图
transforms.ToTensor() # 将 PIL Image 转为 [C,H,W] 范围 [0,1] 的 tensor
])
# 应用预处理并在最前面添加 batch 维度（unsqueeze(0)），最后移动到 device
image_tensor = trans(image).unsqueeze(0).to(device)
# 切换到评估模式，禁用训练时的行为（如 Dropout）
model.eval()
# 在 no_grad 上下文中计算，避免构建计算图，节省内存
with torch.no_grad():
# 前向推理，得到原始输出 logits 或者未归一化的分数
output = model(image_tensor)
# 对 logits 做 softmax 得到每个类别的概率（按 0 维度，即类别维度）
probs = torch.nn.functional.softmax(output[0], 0)
# 取概率最大的索引作为预测结果
predict = torch.argmax(probs).item()
# 取对应的概率值（Tensor）
prob = probs[predict]
# 打印预测信息用于调试：预测类别、该类别概率，以及完整概率向量
print(f'predict: {predict}, prob: {prob}, probs: {probs}')
# 返回 JSON 格式的预测结果，注意要把 Tensor 转为 Python 原生类型（prob.item()）
return jsonify({
'result': predict,
'probability': prob.item()
})
if __name__ == '__main__':
# 启动 Flask 开发服务器，监听本地 5000 端口
app.run(port=5000)
