使用PyTorch进行手写数字识别通常涉及以下步骤:
-
准备数据集:MNIST是一个常用的手写数字数据集,它包含了0到9的手写数字图片。PyTorch提供了直接下载MNIST数据集的工具。
-
定义模型:定义一个神经网络模型来对数字图片进行分类。
-
训练模型:使用MNIST训练数据来训练模型。
-
评估模型:使用MNIST测试数据来评估模型的性能。
下面是一个简单的例子,展示如何使用PyTorch进行手写数字识别:
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
# 1. 准备数据集
batch_size = 64
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
# 2. 定义模型
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.fc1 = nn.Linear(28*28, 500)
self.fc2 = nn.Linear(500, 10)
def forward(self, x):
x = x.view(-1, 28*28)
x = torch.relu(self.fc1(x))
x = self.fc2(x)
return x
model = Net()
# 3. 训练模型
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
num_epochs = 5
for epoch in range(num_epochs):
model.train()
running_loss = 0.0
for batch_idx, (inputs, targets) in enumerate(train_loader):
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, targets)
loss.backward()
optimizer.step()
running_loss += loss.item()
print(f'Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}')
# 4. 评估模型
model.eval()
total_correct = 0
with torch.no_grad():
for inputs, targets in test_loader:
outputs = model(inputs)
_, predicted = torch.max(outputs, 1)
total_correct += (predicted == targets).sum().item()
print(f'Test accuracy: {total_correct/len(test_dataset)*100:.2f}%')
在这个例子中,我们首先定义了数据转换和加载器。然后,我们定义了一个简单的全连接神经网络模型。接着,我们使用交叉熵损失函数和随机梯度下降优化器来训练模型。最后,我们在测试集上评估了模型的性能。
请注意,这只是一个简单的例子。在实际应用中,你可能需要更复杂的网络结构、数据增强、正则化等技术来提高模型的性能。