基于Pytorch的CNN网络识别MNIST手写数字识别

1. 导入必要的库

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from tqdm import tqdm

这部分导入了 PyTorch、Torchvision、Matplotlib、以及进度条工具 tqdm。

2. 数据预处理与加载

# 数据预处理：将图像转换为张量，并进行标准化
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

# 下载MNIST数据集（训练集和测试集）
trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)

# 创建数据加载器
trainloader = DataLoader(trainset, batch_size=64, shuffle=True)
testloader = DataLoader(testset, batch_size=64, shuffle=False)

这部分用于加载 MNIST 数据集，并应用预处理（转换为张量并进行标准化）。然后使用 DataLoader 来分批加载数据。

3. 展示数据样本

# 获取9张图片
data_iter = iter(trainloader)
images, labels = next(data_iter)

# 创建一个图像窗口
fig, axes = plt.subplots(3, 3, figsize=(8, 8))

# 展示9张图片
for i, ax in enumerate(axes.flat):
    ax.imshow(images[i].numpy().squeeze(), cmap='gray')  # 将图片从Tensor转为numpy并去掉通道维度
    ax.set_title(f"Label: {labels[i].item()}")
    ax.axis('off')  # 关闭坐标轴

plt.show()

此部分展示了从训练集中获取的 9 张图像，并且为每张图像显示其标签。

4. 定义神经网络结构

class SimpleNN(nn.Module):
    def __init__(self):
        super(SimpleNN, self).__init__()
        # 定义网络结构
        self.fc1 = nn.Linear(28*28, 128)  # 第一层，输入层到隐藏层
        self.fc2 = nn.Linear(128, 64)     # 第二层，隐藏层到输出层
        self.fc3 = nn.Linear(64, 10)      # 输出层，10个分类

    def forward(self, x):
        # 前向传播：将输入展平为一维向量，经过隐藏层和输出层
        x = x.view(-1, 28*28)  # 将输入图片展平为1维向量
        x = torch.relu(self.fc1(x))  # 第一层 + ReLU激活
        x = torch.relu(self.fc2(x))  # 第二层 + ReLU激活
        x = self.fc3(x)  # 输出层（没有激活函数）
        return x

这部分定义了一个简单的全连接神经网络，包含输入层、两个隐藏层和一个输出层。

5. 初始化模型、损失函数和优化器

# 实例化模型
model = SimpleNN()

# 损失函数：交叉熵损失
criterion = nn.CrossEntropyLoss()

# 优化器：Adam
optimizer = optim.Adam(model.parameters(), lr=0.001)

这里创建了神经网络模型，选择了交叉熵作为损失函数，并使用 Adam 优化器。

6. 训练模型

num_epochs = 5  # 训练的轮数

for epoch in range(num_epochs):
    model.train()  # 设置模型为训练模式
    running_loss = 0.0
    correct = 0
    total = 0
    
    # 使用tqdm为训练集的迭代器添加进度条
    progress_bar = tqdm(trainloader, desc=f"Epoch [{epoch + 1}/{num_epochs}]", unit="batch")
    
    for inputs, labels in progress_bar:
        # 清空梯度
        optimizer.zero_grad()

        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 反向传播
        loss.backward()

        # 更新参数
        optimizer.step()

        # 统计损失和准确率
        running_loss += loss.item()
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        # 更新进度条的后缀信息，显示当前的损失和准确率
        progress_bar.set_postfix(loss=running_loss / (progress_bar.n + 1), accuracy=100 * correct / total)

    # 打印每个epoch的平均损失和准确率
    print(f"Epoch [{epoch + 1}/{num_epochs}], Loss: {running_loss/len(trainloader):.4f}, Accuracy: {100 * correct / total:.2f}%")

这部分包含了训练过程，使用了 tqdm 进度条来显示每个 epoch 的损失和准确率。

7. 评估模型

# 评估模型
model.eval()  # 设置模型为评估模式
correct = 0
total = 0

# 使用tqdm为测试集的迭代器添加进度条
progress_bar = tqdm(testloader, desc="Evaluating", unit="batch")

with torch.no_grad():  # 不计算梯度
    for inputs, labels in progress_bar:
        outputs = model(inputs)
        _, predicted = torch.max(outputs, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

        # 更新进度条的后缀信息，显示当前的准确率
        progress_bar.set_postfix(accuracy=100 * correct / total)

# 打印最终的测试集准确率
print(f"Test Accuracy: {100 * correct / total:.2f}%")

这部分对模型在测试集上的表现进行评估，计算并输出最终的测试集准确率。

8. 获取模型的预测结果并展示

# 获取模型的预测
model.eval()  # 设置模型为评估模式
with torch.no_grad():  # 不需要计算梯度
    outputs = model(images)
    _, predicted = torch.max(outputs, 1)  # 获取每张图片的预测标签

# 创建一个图像窗口
fig, axes = plt.subplots(3, 3, figsize=(8, 8))

# 展示9张图片及其预测标签
for i, ax in enumerate(axes.flat):
    ax.imshow(images[i].numpy().squeeze(), cmap='gray')
    ax.set_title(f"Pred: {predicted[i].item()} ")
    ax.axis('off')

plt.show()

这部分展示了模型对前面展示的 9 张图片的预测结果，并显示预测标签。