Pytorch搭建MyNet实现MNIST手写数字识别

视频：https://www.bilibili.com/video/BV1Wf421B74f/?spm_id_from=333.880.my_history.page.click

1.1 Model类

import torch
import torch.nn as nn


# 改进的三层神经网络
class MyNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 定义全连接层
        self.fc1 = nn.Linear(28 * 28, 256)  # 输入层 输入是28*28的灰度图像，输出是256个神经元
        self.fc2 = nn.Linear(256, 128)  # 第二层，全连接层，输入256个神经元，输出128个神经元
        self.fc3 = nn.Linear(128, 64)  # 第三层，全连接层，输入128个神经元，输出64个神经元
        self.fc4 = nn.Linear(64, 10)  # 第四层，全连接层，输入64个神经元，输出10个类别

    def forward(self, x):
        x = torch.flatten(x, start_dim=1)  # 展平数据，方便进行全连接
        x = torch.relu(self.fc1(x))  # 第一层+ReLU激活
        x = torch.relu(self.fc2(x))  # 第二层+ReLU激活
        x = torch.relu(self.fc3(x))  # 第三层+ReLU激活
        x = self.fc4(x)  # 第四层 输出层 不需要激活函数
        return x

2 train.py创建

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
from .model import MyNet

# 设置随机种子
torch.manual_seed(21)

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device:{device}")

transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转为张量
    transforms.Normalize((0.5,), (0.5,))  # 标准化图像数据，对于灰度图像，只需要一个通道的标准化
])

# 加载MNIST数据集
train_dataset = datasets.MNIST(root='../../datasets', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='../../datasets', train=False, download=True, transform=transform)

# 创建数据加载器
batch_size = 64
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型并将模型移到GPU上
model = MyNet().to(device)

# 定义损失函数和优化器
lr = 0.001  # 学习率为0.001
criterion = nn.CrossEntropyLoss()  # 交叉熵
optimizer = optim.Adam(model.parameters(), lr=lr)  # 使用Adam优化器

# 保存训练过程中的损失和准确率
train_losses = []
train_accuracies = []
test_accuracies = []

epochs = 10
best_accuracy = 0.0  # 记录最佳验证集准确率
best_model_path = 'best_mnist_model.pth'

# 训练10个epoch
for epoch in range(epochs):
    running_loss = 0.0
    corrct_train = 0  # 正确预测的数量
    total_train = 0  # 样本总数

    # 训练过程
    model.train()  # 设定模型为训练模式
    for inputs, labels in train_loader:
        inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到GPU上
        optimizer.zero_grad()  # 梯度清零
        outputs = model(inputs)  # 前向传播
        loss = criterion(outputs, labels)  # 计算损失
        loss.backward()  # 反向传播
        optimizer.step()  # 更新参数
        running_loss += loss.item()  # 累加损失

        # 计算训练集上的准确率
        _, predicted = torch.max(outputs, 1)  # 获取预测结果
        total_train += labels.size(0)  # 累加样本数量
        corrct_train += (predicted == labels).sum().item()  # 累加正确预测的数量

    # 计算训练集上的准确率
    train_accuracy = corrct_train / total_train
    train_losses.append(running_loss / len(train_loader))  # 记录每个 epoch 的平均损失,len(train_loader)为批次数，一个epoch结束后
    train_accuracies.append(train_accuracy)  # 记录每个 epoch 的训练集准确率
    print(
        f"Epoch: {epoch + 1}/{epochs}, Loss: {running_loss / len(train_loader):.4f}, Train Accuracy: {train_accuracy:.2%}")

    # 在测试集上评估模型
    model.eval()  # 设定模型为评估模式
    correct = 0  # 正确的预测数量
    total = 0  # 样本总数

    with torch.no_grad():  # 关闭梯度计算
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)  # 将数据移动到GPU上
            outputs = model(inputs)  # 前向传播
            _, predicted = torch.max(outputs, 1)  # 获取预测结果
            total += labels.size(0)  # 累加样本数量
            correct += (predicted == labels).sum().item()  # 累加正确预测的数量

    # 计算测试集上的准确率
    test_accuracy = correct / total  # 记录每个 epoch 的测试集准确率
    test_accuracies.append(test_accuracy)
    print(f"Epoch: {epoch + 1}/{epochs},Test Accuracy: {test_accuracy:.2%}")

    # 如果测试集准确率提高，保存当前模型的权重
    if test_accuracy > best_accuracy:
        best_accuracy = test_accuracy
        torch.save(model.state_dict(), best_model_path)
        print(f"Best model saved with accuracy: {best_accuracy:.2%}")

print(f"Best Accuracy on test set: {best_accuracy:.2%}")

# 绘制并保存损失和准确率曲线
plt.figure(figsize=(12, 5))

# 绘制损失曲线
plt.subplot(1, 2, 1)  # 选择第一个子图
plt.plot(train_losses, label='Training Loss')  # 传入数据、设置标签为Training Loss
plt.xlabel('Epoch')  # x轴数据
plt.ylabel('Loss')
plt.title('Training Loss over Epochs')  # 设置标题
plt.legend()  # 添加图例
plt.grid(True)  # 添加网格

# 绘制训练集和测试集准确率曲线
plt.subplot(1, 2, 2)
plt.plot(train_accuracies, label='Train Accuracy')
plt.plot(test_accuracies, label='Test Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Train and Test Accuracy over Epochs')
plt.legend()
plt.grid(True)

# 保存图像
plt.tight_layout()
plt.savefig('loss_and_accuracy.png')
plt.show()

3 predict.py 创建

import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
from .model import MyNet

# 设置随机种子
torch.manual_seed(21)

# 检查是否有可用的GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device:{device}")

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转为张量
    transforms.Normalize((0.5,), (0.5,))  # 标准化图像数据，对于灰度图像，只需要一个通道的标准化
])

# 加载MNIST数据集
test_dataset = datasets.MNIST(root='../../datasets', train=False, download=True, transform=transform)

# 创建数据加载器
batch_size = 10
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# 初始化模型并将模型移到GPU上
model = MyNet().to(device)

# 加载保存的权重
best_model_path = 'best_mnist_model.pth'
model.load_state_dict(torch.load(best_model_path))

# 设置模型为评估模式
model.eval()

# 从测试集中取10张图片
examples = enumerate(test_loader)  # 从 test loader 中获取一个带索引的迭代器
# #使用 next()函数从迭代器中获取下一个批次的数据。
# next()返回一个元组，第一个元素是批次的索引(batch idx)，
# #第二个元素是一个包含输入数据和目标标签的元组(example_data，example_targets)
batch_idx, (example_data, example_targets) = next(examples)
# 将数据移动到指定的设备上(例如，CPU或GPU)，
# 以便利用设备的计算能力进行进一步的操作(如模型推理或训练)
example_data, example_targets = example_data.to(device), example_targets.to(device)

# 进行推理
with torch.no_grad():  # 使用 torch.no_grad()上下文管理器，表示在此上下文中不计算梯度。
    outputs = model(example_data)  # 使用模型对 example_data 进行前向传播，获取输出结果。
    # 使用 torch.max()函数，获取 outputs 张量中每行的最大值及其索引。
    # 1 表示在第一个维度上(通常是类别维度)寻找最大值。
    _, predicted = torch.max(outputs, 1)

# 可视化并显示预测结果
fig, axes = plt.subplots(1, 10, figsize=(15, 2))  # 创建一个包含1行10列的子图(subplots)的图形，并指定每个子图的大小
for i in range(10):
    axes[i].imshow(example_data[i].cpu().squeeze(), cmap='gray')  # 在指定的子图(axes[i])中显示图像
    axes[i].set_title(f"Pred:{predicted[i].item()}")  # 在每个子图(axes[i])上设置标题，标题内容显示模型对该样本的预测结果
    axes[i].axis('off')  # 关闭第 i个子图(axes[i])的坐标轴显示

plt.show()
# 输出批次中部分样本的真实标签和模型的预测标签
print(f"True labels: {example_targets.cpu().numpy()}")
print(f"Predicted labels: {predicted.cpu().numpy()}")