利用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')
