使用PyTorch实现MNIST手写数字识别
我们将定义一个简单的卷积神经网络模型。该模型包含两个卷积层、两个最大池化层和两个全连接层。我们使用交叉熵损失函数和Adam优化器来训练模型。通过本教程,我们学习了如何使用PyTorch构建一个简单的卷积神经网络来识别MNIST手写数字。我们涵盖了数据加载、模型定义、训练和测试等步骤。希望本教程对你理解如何使用PyTorch进行深度学习有所帮助。如果你有任何问题或建议,欢迎在评论区留言!
在本教程中,我们将使用PyTorch框架来实现一个简单的卷积神经网络(CNN),用于识别MNIST数据集中的手写数字。MNIST数据集是一个广泛使用的数据集,包含60,000个训练样本和10,000个测试样本,每个样本是一个28x28的灰度图像,表示0到9之间的手写数字。
1. 导入必要的库
首先,我们需要导入一些必要的库,包括PyTorch、torchvision、matplotlib等。torchvision提供了对MNIST数据集的便捷访问,而matplotlib用于可视化数据。
import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import matplotlib.pyplot as plt
2. 设置设备
我们将使用GPU来加速训练过程(如果可用)。通过以下代码,我们可以检查是否有可用的GPU,并将其设置为默认设备。
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
3. 数据预处理与加载
接下来,我们需要对数据进行预处理,并将其加载到DataLoader中。我们将使用transforms.Compose来定义一系列的数据转换操作,包括将图像转换为张量并进行归一化。
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.5,), (0.5,))
])
train_set = datasets.MNIST(root='D:\\python_text\\python\\pytorch深度学习实战\\dataset', transform=transform, train=True, download=True)
test_set = datasets.MNIST(root='D:\\python_text\\python\\pytorch深度学习实战\\dataset', transform=transform, train=False, download=True)
batch_size = 64
train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=4)
test_loader = DataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=4)
4. 可视化数据
在训练模型之前,我们可以先可视化一些训练数据,以确保数据加载正确。
plt.figure(figsize=(20,10))
for i in range(20):
plt.subplot(5,10,i+1)
plt.xticks([])
plt.yticks([])
plt.grid(False)
plt.imshow(train_set.data[i], cmap=plt.cm.binary)
plt.xlabel(train_set.targets[i])
plt.show()
5. 定义模型
我们将定义一个简单的卷积神经网络模型。该模型包含两个卷积层、两个最大池化层和两个全连接层。
from torch import nn, optim
class Model(nn.Module):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
self.fc1 = nn.Linear(64*7*7, 64)
self.fc2 = nn.Linear(64, 10)
self.relu = nn.ReLU()
self.flatten = nn.Flatten()
def forward(self, x):
x = self.pool1(self.relu(self.conv1(x)))
x = self.pool2(self.relu(self.conv2(x)))
x = self.flatten(x)
x = self.relu(self.fc1(x))
return self.fc2(x)
model = Model()
model.to(device)
6. 定义损失函数和优化器
我们使用交叉熵损失函数和Adam优化器来训练模型。
criterion = nn.CrossEntropyLoss()
criterion.to(device)
optimizer = optim.Adam(model.parameters(), lr=0.001)
7. 训练模型
我们定义了一个训练函数,用于训练模型并在每个epoch结束后测试模型的准确率。
def train(model, train_loader, test_loader, optimizer, criterion, num_epochs, device):
for epoch in range(num_epochs):
# 训练模型
model.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
output = model(x)
loss = criterion(output, y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print('epoch:', epoch+1, 'loss:', loss.item())
# 测试模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
for x, y in test_loader:
x, y = x.to(device), y.to(device)
output = model(x)
_, predicted = torch.max(output.data, dim=1)
correct += (predicted == y).sum().item()
total += y.size(0)
print('acc:', correct/total)
train(model, train_loader, test_loader, optimizer, criterion, 10, device)
8. 结果
在训练过程中,我们可以看到每个epoch的损失值和测试集上的准确率。经过10个epoch的训练,模型的准确率可以达到98%以上。
9. 总结
通过本教程,我们学习了如何使用PyTorch构建一个简单的卷积神经网络来识别MNIST手写数字。我们涵盖了数据加载、模型定义、训练和测试等步骤。希望本教程对你理解如何使用PyTorch进行深度学习有所帮助。
如果你有任何问题或建议,欢迎在评论区留言!
更多推荐

所有评论(0)