Python AI 与深度学习 - D2.MNIST 手写数字识别

# 导入所有需要的库
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt

# -------------------------- 1. 配置超参数和设备 --------------------------
batch_size = 64  # 每次加载64个样本
epochs = 5       # 整个数据集训练5遍
lr = 0.001       # 学习率
# 设备配置：优先CPU（入门）
device = torch.device("cpu")

# -------------------------- 2. 数据预处理和加载MNIST数据集 --------------------------
# 数据预处理：将图像转为Tensor+归一化（0-1之间）
transform = transforms.Compose([
    transforms.ToTensor(),  # 转为Tensor，形状从(28,28)→(1,28,28)，值从0-255→0-1
    transforms.Normalize((0.1307,), (0.3081,))  # MNIST数据集的均值和标准差，固定值
])
# 加载训练集和测试集（自动下载，无需手动找数据）
train_dataset = datasets.MNIST(
    root='./data',  # 数据保存路径
    train=True,     # 训练集
    download=True,  # 自动下载
    transform=transform
)
test_dataset = datasets.MNIST(
    root='./data',
    train=False,    # 测试集
    download=True,
    transform=transform
)
# 用DataLoader批量加载数据，打乱顺序
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# -------------------------- 3. 定义简单的卷积神经网络（CNN） --------------------------
class MNIST_CNN(nn.Module):
    def __init__(self):
        super(MNIST_CNN, self).__init__()
        # 卷积层：输入1通道（灰度图），输出16通道，卷积核3x3，步长1，填充1
        self.conv1 = nn.Conv2d(1, 16, 3, 1, 1)
        self.relu = nn.ReLU()  # 激活函数
        self.pool = nn.MaxPool2d(2, 2)  # 池化层：2x2，步长2，缩小图像尺寸
        # 卷积层：输入16通道，输出32通道
        self.conv2 = nn.Conv2d(16, 32, 3, 1, 1)
        # 全连接层：将特征展平后，映射到10个输出（0-9数字）
        self.fc1 = nn.Linear(32 * 7 * 7, 10)  # 28→池化2次→7，32*7*7是展平后的特征数

    def forward(self, x):
        # 前向传播：卷积→激活→池化→卷积→激活→池化→展平→全连接
        x = self.pool(self.relu(self.conv1(x)))
        x = self.pool(self.relu(self.conv2(x)))
        x = x.view(-1, 32 * 7 * 7)  # 展平：(batch_size, 32,7,7)→(batch_size, 32*7*7)
        x = self.fc1(x)
        return x

# 实例化模型，移到指定设备
model = MNIST_CNN().to(device)

# -------------------------- 4. 定义损失函数和优化器 --------------------------
criterion = nn.CrossEntropyLoss()  # 交叉熵损失，适合分类任务
optimizer = optim.Adam(model.parameters(), lr=lr)  # Adam优化器，比SGD更快收敛

# -------------------------- 5. 模型训练 --------------------------
print("开始训练MNIST模型...")
model.train()  # 模型进入训练模式
for epoch in range(epochs):
    running_loss = 0.0
    for i, (images, labels) in enumerate(train_loader):
        # 将数据移到指定设备
        images, labels = images.to(device), labels.to(device)
        # 1. 清空梯度（必须步骤，否则梯度累加）
        optimizer.zero_grad()
        # 2. 前向传播：输入图像，得到预测结果
        outputs = model(images)
        # 3. 计算损失：预测结果与真实标签的差距
        loss = criterion(outputs, labels)
        # 4. 反向传播：计算梯度
        loss.backward()
        # 5. 优化器更新参数
        optimizer.step()
        # 累计损失
        running_loss += loss.item()
    # 打印每个epoch的平均损失
    avg_loss = running_loss / len(train_loader)
    print(f"第{epoch+1}/{epochs}轮训练，平均损失：{avg_loss:.4f}")

print("训练完成！")

# -------------------------- 6. 模型测试（评估准确率） --------------------------
print("开始测试模型...")
model.eval()  # 模型进入测试模式，关闭Dropout/BN等
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)
        # 取预测概率最大的类别作为预测结果
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

# 打印测试准确率
accuracy = 100 * correct / total
print(f"MNIST测试集准确率：{accuracy:.2f}%")  # 正常训练后准确率≥98%

# -------------------------- 7. 可视化预测结果（可选） --------------------------
# 取测试集中的前5张图像，展示预测结果
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs, 1)

# 绘制图像
plt.figure(figsize=(10, 5))
for i in range(5):
    plt.subplot(1, 5, i+1)
    plt.imshow(images[i].squeeze().numpy(), cmap='gray')  # 去掉通道维度，显示灰度图
    plt.title(f"真实：{labels[i].item()}\n预测：{predicted[i].item()}")
    plt.axis('off')
plt.show()

# -------------------------- 8. 保存模型（后续可直接加载使用） --------------------------
torch.save(model.state_dict(), './mnist_cnn_model.pth')
print("模型已保存为：mnist_cnn_model.pth")

# -------------------------- 加载模型（后续使用） --------------------------
# new_model = MNIST_CNN().to(device)
# new_model.load_state_dict(torch.load('./mnist_cnn_model.pth'))
# new_model.eval()  # 加载后必须进入测试模式

 

运行完成后，会得到 3 个关键有效输出，代表项目成功：

当前使用设备：cuda:0
开始训练MNIST模型...
第1/5轮训练，平均损失：0.2xx
第2/5轮训练，平均损失：0.0xx
...
训练完成！
开始测试模型...
MNIST测试集准确率：98.xx%
模型已保存为：mnist_cnn_model.pth