利用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')
15万+
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
