[原创] #AI挑战营第一站#MNIST手写数字识别模型训练

本文使用的环境是conda+pycharm

1.环境安装

下面是我的环境

2.数据集准备

MNIST 数据集可在 http://yann.lecun.com/exdb/mnist/ 获取。代码加载数据集如下所示

# 导包
import torch
import torch.nn as nn  # 神经网络
import torch.optim as optim  # 定义优化器
from torchvision import datasets, transforms  # 数据集    transforms完成对数据的处理

# 定义超参数
input_size = 28 * 28  # 输入大小
hidden_size = 512  # 隐藏层大小
num_classes = 10  # 输出大小（类别数）
batch_size = 100  # 批大小
learning_rate = 0.001  # 学习率
num_epochs = 10  # 训练轮数

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

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  # 一批数据为100个
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

3.定义网络

# 定义 MLP 网络
class MLP(nn.Module):
    # 初始化方法
    # input_size 输入数据的维度
    # hidden_size 隐藏层的大小
    # num_classes 输出分类的数量
    def __init__(self, input_size, hidden_size, num_classes):
        # 调用父类的初始化方法
        super(MLP, self).__init__()
        # 定义第1个全连接层
        self.fc1 = nn.Linear(input_size, hidden_size)
        # 定义ReLu激活函数
        self.relu = nn.ReLU()
        # 定义第2个全连接层
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        # 定义第3个全连接层
        self.fc3 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out

这个MLP网络有两个隐藏层，每个隐藏层有512个神经元。

4.训练网络

# 定义损失函数和优化器，接下来使用 PyTorch 的自动求导功能和优化器进行训练
criterion = nn.CrossEntropyLoss() #损失函数
# CrossEntropyLoss = Softmax + log + nllloss
optimizer = optim.Adam(model.parameters(), lr=learning_rate) #Adam优化器
# optimizer = optim.SGD(model.parameters(),0.2)

# 外层for循环控制训练的次数
# 内层for循环控制从DataLoader中循环读取数据
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, 28 * 28)  # 将images转换成向量

        outputs = model(images)  # 将数据送到网络中

        loss = criterion(outputs, labels)  # 计算损失

        optimizer.zero_grad()  # 首先将梯度清零
        loss.backward()  # 反向传播

        optimizer.step()  # 更新参数

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}],Step[{i + 1}/{len(train_loader)}],Loss:{loss.item():.4f}')

5.评估模型的准确率

# 测试网络
with torch.no_grad():
    correct = 0
    total = 0
    # 从test_loader中循环读取测试数据
    for images, labels in test_loader:
        # 将images转换成向量
        images = images.reshape(-1, 28 * 28)
        # 将数据传送到网络
        outputs = model(images)
        # 取出最大值对应的索引 即预测值
        _, predicted = torch.max(outputs.data, 1)  # 返回一个元组：第一个为最大值，第二个是最大值的下标
        # 累加labels数量   labels为形状为(batch_size,1)的矩阵，取size(0)就是取出batch_size的大小（该批的大小）
        total += labels.size(0)
        # 预测值与labels值对比 获取预测正确的数量
        correct += (predicted == labels).sum().item()
    # 打印最终的准确率
    print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')
# 保存模型的状态字典
torch.save(model.state_dict(), 'mnist.pth')

6.转换成ONNX模型

# 保存模型的状态字典
torch.save(model.state_dict(), 'mnist.pth')
model.load_state_dict(torch.load('mnist.pth'))
# 将模型设置为评估模式（如果需要）
model.eval()
#导出为onnx模型
input = torch.randn(1, 28, 28)
torch.onnx.export(model, input, "mnist.onnx", verbose=True)

7.查看模型结构

使用netron一个轻量级的 ONNX 模型查看器来查看

8.完整代码

# 导包
import torch
import torch.nn as nn  # 神经网络
import torch.optim as optim  # 定义优化器
from torchvision import datasets, transforms  # 数据集    transforms完成对数据的处理

# 定义超参数
input_size = 28 * 28  # 输入大小
hidden_size = 512  # 隐藏层大小
num_classes = 10  # 输出大小（类别数）
batch_size = 100  # 批大小
learning_rate = 0.001  # 学习率
num_epochs = 10  # 训练轮数

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

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)  # 一批数据为100个
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)


# 定义 MLP 网络
class MLP(nn.Module):
    # 初始化方法
    # input_size 输入数据的维度
    # hidden_size 隐藏层的大小
    # num_classes 输出分类的数量
    def __init__(self, input_size, hidden_size, num_classes):
        # 调用父类的初始化方法
        super(MLP, self).__init__()
        # 定义第1个全连接层
        self.fc1 = nn.Linear(input_size, hidden_size)
        # 定义ReLu激活函数
        self.relu = nn.ReLU()
        # 定义第2个全连接层
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        # 定义第3个全连接层
        self.fc3 = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # 将输入张量展平为向量
        x = x.view(x.size(0), -1)
        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out


# 实例化 MLP 网络
model = MLP(input_size, hidden_size, num_classes)

# 现在我们已经定义了 MLP 网络并加载了 MNIST 数据集，接下来使用 PyTorch 的自动求导功能和优化器进行训练。首先，定义损失函数和优化器；然后迭代训练数据并使用优化器更新网络参数。

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
# CrossEntropyLoss = Softmax + log + nllloss
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
# optimizer = optim.SGD(model.parameters(),0.2)

# 训练网络
# 外层for循环控制训练的次数
# 内层for循环控制从DataLoader中循环读取数据
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        images = images.reshape(-1, 28 * 28)  # 将images转换成向量

        outputs = model(images)  # 将数据送到网络中

        loss = criterion(outputs, labels)  # 计算损失

        optimizer.zero_grad()  # 首先将梯度清零
        loss.backward()  # 反向传播

        optimizer.step()  # 更新参数

        if (i + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{num_epochs}],Step[{i + 1}/{len(train_loader)}],Loss:{loss.item():.4f}')

# 最后，我们可以在测试数据上评估模型的准确率：
# 测试网络
with torch.no_grad():
    correct = 0
    total = 0
    # 从test_loader中循环读取测试数据
    for images, labels in test_loader:
        # 将images转换成向量
        images = images.reshape(-1, 28 * 28)
        # 将数据传送到网络
        outputs = model(images)
        # 取出最大值对应的索引 即预测值
        _, predicted = torch.max(outputs.data, 1)  # 返回一个元组：第一个为最大值，第二个是最大值的下标
        # 累加labels数量   labels为形状为(batch_size,1)的矩阵，取size(0)就是取出batch_size的大小（该批的大小）
        total += labels.size(0)
        # 预测值与labels值对比 获取预测正确的数量
        correct += (predicted == labels).sum().item()
    # 打印最终的准确率
    print(f'Accuracy of the network on the 10000 test images: {100 * correct / total}%')

# 保存模型
# 保存模型的状态字典
torch.save(model.state_dict(), 'mnist.pth')
model.load_state_dict(torch.load('mnist.pth'))
# 将模型设置为评估模式（如果需要）
model.eval()
#导出为onnx模型
input = torch.randn(1, 28, 28)
torch.onnx.export(model, input, "mnist.onnx", verbose=True)

main.py (4.34 KB)
(下载次数: 2, 2024-4-14 20:48 上传)

mnist.pth (2.56 MB)
(下载次数: 0, 2024-4-14 20:47 上传)

mnist.onnx (2.57 MB)
(下载次数: 0, 2024-4-14 20:47 上传)