利用PyTorch实现基于MNIST数据集的手写数字识别

GeekLiHua

发布于 2025-01-21 16:08:51

66100

代码可运行

文章被收录于专栏：JavaJava

运行总次数：0

代码可运行

利用PyTorch实现基于MNIST数据集的手写数字识别

简介：如何使用PyTorch实现基于MNIST数据集的手写数字识别。手写数字识别是计算机视觉领域的经典问题之一，旨在将手写数字图像转换为对应的数字标签。

数据集简介

MNIST数据集是一个经典的手写数字数据集，包含了60000张训练图像和10000张测试图像。每张图像的大小为28x28像素，图像内容为0到9的手写数字。我们将使用这个数据集来训练和测试我们的模型。

代码实现

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

# 设置随机种子，以确保结果的可重复性
torch.manual_seed(0)

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),  # 将图像转换为张量
    transforms.Normalize((0.1307,), (0.3081,))  # 标准化图像
])

# 加载MNIST数据集
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)  # 训练数据集
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)  # 测试数据集

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

# 定义神经网络模型（LeNet）
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5, 1)
        self.conv2 = nn.Conv2d(20, 50, 5, 1)
        self.fc1 = nn.Linear(4*4*50, 500)
        self.fc2 = nn.Linear(500, 10)

    def forward(self, x):
        x = nn.functional.relu(self.conv1(x))
        x = nn.functional.max_pool2d(x, 2, 2)
        x = nn.functional.relu(self.conv2(x))
        x = nn.functional.max_pool2d(x, 2, 2)
        x = x.view(-1, 4*4*50)
        x = nn.functional.relu(self.fc1(x))
        x = self.fc2(x)
        return nn.functional.log_softmax(x, dim=1)

# 创建神经网络模型实例
model = Net()

# 定义优化器和损失函数
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
criterion = nn.CrossEntropyLoss()

# 训练模型
def train(epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('训练 Epoch: {} [{}/{} ({:.0f}%)]\t损失: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

# 测试模型
def test():
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            output = model(data)
            test_loss += criterion(output, target).item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\n测试集: 平均损失: {:.4f}, 准确率: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

# 进行模型训练和测试
for epoch in range(1, 11):
    train(epoch)
    test()

结果展示

训练 Epoch: 1 [0/60000 (0%)]	损失: 2.320024
训练 Epoch: 1 [6400/60000 (11%)]	损失: 0.577842
训练 Epoch: 1 [12800/60000 (21%)]	损失: 0.452133
训练 Epoch: 1 [19200/60000 (32%)]	损失: 0.278623
训练 Epoch: 1 [25600/60000 (43%)]	损失: 0.277499
训练 Epoch: 1 [32000/60000 (53%)]	损失: 0.115603
训练 Epoch: 1 [38400/60000 (64%)]	损失: 0.148797
训练 Epoch: 1 [44800/60000 (75%)]	损失: 0.092067
训练 Epoch: 1 [51200/60000 (85%)]	损失: 0.125068
训练 Epoch: 1 [57600/60000 (96%)]	损失: 0.062419

测试集: 平均损失: 0.0014, 准确率: 9732/10000 (97%)

训练 Epoch: 2 [0/60000 (0%)]	损失: 0.078582
训练 Epoch: 2 [6400/60000 (11%)]	损失: 0.172461
训练 Epoch: 2 [12800/60000 (21%)]	损失: 0.121057
训练 Epoch: 2 [19200/60000 (32%)]	损失: 0.037924
训练 Epoch: 2 [25600/60000 (43%)]	损失: 0.292302
训练 Epoch: 2 [32000/60000 (53%)]	损失: 0.022772
训练 Epoch: 2 [38400/60000 (64%)]	损失: 0.028810
训练 Epoch: 2 [44800/60000 (75%)]	损失: 0.080398
训练 Epoch: 2 [51200/60000 (85%)]	损失: 0.032344
训练 Epoch: 2 [57600/60000 (96%)]	损失: 0.082773

测试集: 平均损失: 0.0011, 准确率: 9773/10000 (98%)

训练 Epoch: 3 [0/60000 (0%)]	损失: 0.063188
训练 Epoch: 3 [6400/60000 (11%)]	损失: 0.017849
训练 Epoch: 3 [12800/60000 (21%)]	损失: 0.009525
训练 Epoch: 3 [19200/60000 (32%)]	损失: 0.020973
训练 Epoch: 3 [25600/60000 (43%)]	损失: 0.109975
训练 Epoch: 3 [32000/60000 (53%)]	损失: 0.110332
训练 Epoch: 3 [38400/60000 (64%)]	损失: 0.050331
训练 Epoch: 3 [44800/60000 (75%)]	损失: 0.060158
训练 Epoch: 3 [51200/60000 (85%)]	损失: 0.091192
训练 Epoch: 3 [57600/60000 (96%)]	损失: 0.039621

测试集: 平均损失: 0.0007, 准确率: 9847/10000 (98%)

训练 Epoch: 4 [0/60000 (0%)]	损失: 0.025815
训练 Epoch: 4 [6400/60000 (11%)]	损失: 0.041258
训练 Epoch: 4 [12800/60000 (21%)]	损失: 0.020587
训练 Epoch: 4 [19200/60000 (32%)]	损失: 0.105892
训练 Epoch: 4 [25600/60000 (43%)]	损失: 0.093032
训练 Epoch: 4 [32000/60000 (53%)]	损失: 0.012649
训练 Epoch: 4 [38400/60000 (64%)]	损失: 0.073835
训练 Epoch: 4 [44800/60000 (75%)]	损失: 0.035451
训练 Epoch: 4 [51200/60000 (85%)]	损失: 0.024816
训练 Epoch: 4 [57600/60000 (96%)]	损失: 0.020467

测试集: 平均损失: 0.0007, 准确率: 9843/10000 (98%)

训练 Epoch: 5 [0/60000 (0%)]	损失: 0.024172
训练 Epoch: 5 [6400/60000 (11%)]	损失: 0.095893
训练 Epoch: 5 [12800/60000 (21%)]	损失: 0.075022
训练 Epoch: 5 [19200/60000 (32%)]	损失: 0.189776
训练 Epoch: 5 [25600/60000 (43%)]	损失: 0.007984
训练 Epoch: 5 [32000/60000 (53%)]	损失: 0.013882
训练 Epoch: 5 [38400/60000 (64%)]	损失: 0.059460
训练 Epoch: 5 [44800/60000 (75%)]	损失: 0.151336
训练 Epoch: 5 [51200/60000 (85%)]	损失: 0.077553
训练 Epoch: 5 [57600/60000 (96%)]	损失: 0.039423

测试集: 平均损失: 0.0006, 准确率: 9863/10000 (99%)

训练 Epoch: 6 [0/60000 (0%)]	损失: 0.025369
训练 Epoch: 6 [6400/60000 (11%)]	损失: 0.029651
训练 Epoch: 6 [12800/60000 (21%)]	损失: 0.009427
训练 Epoch: 6 [19200/60000 (32%)]	损失: 0.052345
训练 Epoch: 6 [25600/60000 (43%)]	损失: 0.050803
训练 Epoch: 6 [32000/60000 (53%)]	损失: 0.032373
训练 Epoch: 6 [38400/60000 (64%)]	损失: 0.011127
训练 Epoch: 6 [44800/60000 (75%)]	损失: 0.017743
训练 Epoch: 6 [51200/60000 (85%)]	损失: 0.022161
训练 Epoch: 6 [57600/60000 (96%)]	损失: 0.032988

测试集: 平均损失: 0.0006, 准确率: 9878/10000 (99%)

训练 Epoch: 7 [0/60000 (0%)]	损失: 0.010487
训练 Epoch: 7 [6400/60000 (11%)]	损失: 0.038379
训练 Epoch: 7 [12800/60000 (21%)]	损失: 0.020085
训练 Epoch: 7 [19200/60000 (32%)]	损失: 0.017111
训练 Epoch: 7 [25600/60000 (43%)]	损失: 0.005051
训练 Epoch: 7 [32000/60000 (53%)]	损失: 0.011233
训练 Epoch: 7 [38400/60000 (64%)]	损失: 0.011859
训练 Epoch: 7 [44800/60000 (75%)]	损失: 0.039866
训练 Epoch: 7 [51200/60000 (85%)]	损失: 0.120959
训练 Epoch: 7 [57600/60000 (96%)]	损失: 0.017959

测试集: 平均损失: 0.0005, 准确率: 9884/10000 (99%)

训练 Epoch: 8 [0/60000 (0%)]	损失: 0.005525
训练 Epoch: 8 [6400/60000 (11%)]	损失: 0.006888
训练 Epoch: 8 [12800/60000 (21%)]	损失: 0.032923
训练 Epoch: 8 [19200/60000 (32%)]	损失: 0.016291
训练 Epoch: 8 [25600/60000 (43%)]	损失: 0.004896
训练 Epoch: 8 [32000/60000 (53%)]	损失: 0.091328
训练 Epoch: 8 [38400/60000 (64%)]	损失: 0.038165
训练 Epoch: 8 [44800/60000 (75%)]	损失: 0.005650
训练 Epoch: 8 [51200/60000 (85%)]	损失: 0.050541
训练 Epoch: 8 [57600/60000 (96%)]	损失: 0.017801

测试集: 平均损失: 0.0005, 准确率: 9891/10000 (99%)

训练 Epoch: 9 [0/60000 (0%)]	损失: 0.036731
训练 Epoch: 9 [6400/60000 (11%)]	损失: 0.025160
训练 Epoch: 9 [12800/60000 (21%)]	损失: 0.006145
训练 Epoch: 9 [19200/60000 (32%)]	损失: 0.026249
训练 Epoch: 9 [25600/60000 (43%)]	损失: 0.020545
训练 Epoch: 9 [32000/60000 (53%)]	损失: 0.029771
训练 Epoch: 9 [38400/60000 (64%)]	损失: 0.007581
训练 Epoch: 9 [44800/60000 (75%)]	损失: 0.032595
训练 Epoch: 9 [51200/60000 (85%)]	损失: 0.002854
训练 Epoch: 9 [57600/60000 (96%)]	损失: 0.061837

测试集: 平均损失: 0.0006, 准确率: 9870/10000 (99%)

训练 Epoch: 10 [0/60000 (0%)]	损失: 0.008528
训练 Epoch: 10 [6400/60000 (11%)]	损失: 0.014426
训练 Epoch: 10 [12800/60000 (21%)]	损失: 0.012523
训练 Epoch: 10 [19200/60000 (32%)]	损失: 0.001507
训练 Epoch: 10 [25600/60000 (43%)]	损失: 0.027995
训练 Epoch: 10 [32000/60000 (53%)]	损失: 0.026279
训练 Epoch: 10 [38400/60000 (64%)]	损失: 0.001737
训练 Epoch: 10 [44800/60000 (75%)]	损失: 0.052491
训练 Epoch: 10 [51200/60000 (85%)]	损失: 0.003997
训练 Epoch: 10 [57600/60000 (96%)]	损失: 0.020422

测试集: 平均损失: 0.0005, 准确率: 9885/10000 (99%)

Net(
  (conv1): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
  (conv2): Conv2d(20, 50, kernel_size=(5, 5), stride=(1, 1))
  (fc1): Linear(in_features=800, out_features=500, bias=True)
  (fc2): Linear(in_features=500, out_features=10, bias=True)
)

运行结果的总结如下：

每个训练周期（Epoch）包含多个批次（Batches），每个批次大小为64。
训练周期逐渐增加，损失值逐渐减小，这表明模型在训练过程中逐渐学习到了数据的特征。
测试集的平均损失逐渐减小，准确率逐渐提高，这表明模型在训练后在测试集上表现良好。
训练周期增加到10时，测试集的准确率达到了99%左右，模型已经取得了不错的分类效果。
最后的输出展示了训练后的神经网络模型的结构，包括卷积层和全连接层的参数设置。

这个结果表明，LeNet模型在MNIST数据集上取得了良好的分类效果，并且模型的结构也得到了有效的训练和优化。

结果可视化

保存模型

最后，我们可以将训练好的模型保存起来，以便在需要时重新加载并使用。

# 保存模型
torch.save(model.state_dict(), 'mnist_model.pth')

本文参与腾讯云自媒体同步曝光计划，分享自作者个人站点/博客。

原始发表：2024-03-12，如有侵权请联系 cloudcommunity@tencent.com 删除

mnist

本文分享自作者个人站点/博客前往查看

如有侵权，请联系 cloudcommunity@tencent.com 删除。

本文参与腾讯云自媒体同步曝光计划，欢迎热爱写作的你一起参与！

登录后参与评论

0 条评论

热度

利用PyTorch实现基于MNIST数据集的手写数字识别

利用PyTorch实现基于MNIST数据集的手写数字识别

利用PyTorch实现基于MNIST数据集的手写数字识别

代码实现

结果可视化

保存模型

社区

活动

圈层

关于

腾讯云开发者

热门产品

热门推荐

更多推荐