🔧 MNIST 手写数字识别

MNIST Handwritten Digit Classification

深度学习 · 计算机视觉 · CNN

---

📖 项目简介

基于 PyTorch 的 MNIST 手写数字识别模型训练与测试框架，采用双层卷积神经网络（CNN）架构，支持 CPU 和 CUDA 计算模式。

📌 功能特性

• ✅ 双层 CNN 架构：16 → 32 通道卷积层 + 全连接层
• ✅ 灵活的训练配置：可调节批次大小、学习率、训练轮数等参数
• ✅ CUDA 加速支持：自动检测并使用 GPU 进行计算
• ✅ 本地数据加载：支持使用本地 MNIST 数据集（无需下载）
• ✅ 模型保存与加载：训练后自动保存模型，测试时加载已有模型

📁 项目结构

查看目录结构

MNIST/
├── main.py              # 🚀 程序入口，解析参数并分发到训练/测试
├── train.py             # 🔧 模型训练逻辑
├── test.py              # 🧪 模型测试/评估逻辑
├── model.py             # 🧠 CNN 模型定义
├── utils/               # 📦 工具模块
│   ├── loadData.py      # 💾 MNIST 数据加载
│   └── paser.py         # ⚙️ 命令行参数解析
└── README.md            # 📄 项目说明文档

🧠 模型架构

查看网络结构详情

Mnist CNN 模型

Input (1, 28, 28)
    ↓
Conv2d(1→16, kernel=5, padding=2) + ReLU + MaxPool(2)
    ↓ (16, 14, 14)
Conv2d(16→32, kernel=5, padding=2) + ReLU + MaxPool(2)
    ↓ (32, 7, 7)
Flatten
    ↓ (1568)
Linear(1568→10)
    ↓
Output (10) - 数字 0-9 的分类 logits

参数量：约 28,938 个可训练参数

🔧 环境配置

前置要求

• Python 3.8 或更高版本
• PyTorch 1.10+
• torchvision

安装依赖

# 使用 pip 安装 PyTorch（CPU 版本）
pip install torch torchvision

# 如需 CUDA 支持，请访问 https://pytorch.org 获取对应版本的安装命令

💾 数据准备

⚠️ 重要说明：本项目使用 download=False 模式，需要手动准备 MNIST 数据集。

获取数据文件

从 MNIST 官网 下载以下四个文件（.gz 格式）：

• train-images-idx3-ubyte.gz
• train-labels-idx1-ubyte.gz
• t10k-images-idx3-ubyte.gz
• t10k-labels-idx1-ubyte.gz

放置数据文件

# 创建数据目录
mkdir -p data/MNIST/raw

# 将下载的文件移动到数据目录
mv /path/to/*.gz data/MNIST/raw/

# 解压文件（必须生成未压缩的 idx 文件）
gunzip -k data/MNIST/raw/*.gz

确认 data/MNIST/raw 目录中存在以下未压缩文件：

• train-images-idx3-ubyte
• train-labels-idx1-ubyte
• t10k-images-idx3-ubyte
• t10k-labels-idx1-ubyte

🚀 快速开始

训练模型

# 使用默认参数训练
python main.py --mode train --data_dir ./data

# 自定义参数训练
python main.py --mode train \
    --data_dir ./data \
    --batch_size 64 \
    --epochs 10 \
    --lr 0.01 \
    --save_model_path ./mnist_model.pth

测试模型

# 使用已保存的模型进行测试
python main.py --mode test \
    --data_dir ./data \
    --save_model_path ./mnist_model.pth

# 禁用 CUDA 测试
python main.py --mode test \
    --data_dir ./data \
    --save_model_path ./mnist_model.pth \
    --no_cuda

⚙️ 命令行参数

查看完整参数列表

参数	类型	默认值	说明
--batch_size	int	64	训练批次大小
--test_batch_size	int	1000	测试批次大小
--epochs	int	10	训练轮数
--lr	float	0.01	学习率
--momentum	float	0.9	SGD 动量因子
--no_cuda	flag	False	禁用 CUDA 训练
--seed	int	42	随机种子
--log_interval	int	10	训练日志打印间隔（批次）
--data_dir	str	./data	数据存储目录
--save_model_path	str	./mnist_model.pth	模型保存/加载路径
--mode	str	train	运行模式：train 或 test

📊 训练输出示例

================= Training the Minist Model =================
 Batch Size: 64
 Epochs: 10
 Learning Rate: 0.01
 Momentum: 0.9
 Data Directory: ./data
 Save Model Path: ./mnist_model.pth
 ------------------------------------------------------------
Train Epoch: 0 [0/60000 (0%)]	Loss: 2.301234
Train Epoch: 0 [6400/60000 (11%)]	Loss: 1.234567
...
Model saved to ./mnist_model.pth

🧪 测试输出示例

================= Testing the Minist Model =================
 Batch Size: 64
 Data Directory: ./data
 Load Model Path: ./mnist_model.pth
 ------------------------------------------------------------
Test Batch: 0	Accuracy: 98.44%
Test Batch: 1	Accuracy: 98.44%
...
Testing completed.

🔍 故障排查

常见问题解决

Dataset not found 错误

错误信息：Dataset not found. You can use download=True to download it

原因：torchvision 在 download=False 模式下未找到原始 MNIST 文件

解决方案：

1. 检查 data/MNIST/raw 目录是否存在
2. 确认目录中包含未压缩的 *-images-idx3-ubyte 和 *-labels-idx1-ubyte 文件
3. 如果只有 .gz 文件，请运行 gunzip -k data/MNIST/raw/*.gz 解压
4. 确保 --data_dir 参数指向正确的路径

CUDA 相关错误

如果遇到 CUDA 相关问题，可以添加 --no_cuda 参数使用 CPU 运行：

python main.py --mode train --no_cuda

📄 许可证

本项目采用 MIT 许可证。

---

Made with ❤️ by airprofly