以下是PyTorch实现MNIST数据集手写数字识别的完整攻略。
首先,我们需要导入必要的库,包括PyTorch、torchvision、numpy和matplotlib等。
import torch
import torchvision
import numpy as np
import matplotlib.pyplot as plt
接下来,我们需要加载MNIST数据集。可以使用torchvision中的datasets模块来加载数据集。
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, transform=torchvision.transforms.ToTensor(), download=True)
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, transform=torchvision.transforms.ToTensor(), download=True)
我们使用一个简单的卷积神经网络来实现手写数字识别。定义模型的代码如下:
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = torch.nn.Conv2d(1, 32, kernel_size=5, padding=2)
self.conv2 = torch.nn.Conv2d(32, 64, kernel_size=5, padding=2)
self.fc1 = torch.nn.Linear(7 * 7 * 64, 1024)
self.fc2 = torch.nn.Linear(1024, 10)
def forward(self, x):
x = torch.nn.functional.relu(self.conv1(x))
x = torch.nn.functional.max_pool2d(x, 2)
x = torch.nn.functional.relu(self.conv2(x))
x = torch.nn.functional.max_pool2d(x, 2)
x = x.view(-1, 7 * 7 * 64)
x = torch.nn.functional.relu(self.fc1(x))
x = torch.nn.functional.dropout(x, training=self.training)
x = self.fc2(x)
return torch.nn.functional.log_softmax(x, dim=1)
model = Net()
我们使用交叉熵损失函数和随机梯度下降优化器来训练模型。
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
接下来,我们使用训练集对模型进行训练。
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
for epoch in range(10):
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('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
最后,我们使用测试集对模型进行测试。
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=True)
with torch.no_grad():
correct = 0
total = 0
for data, target in test_loader:
output = model(data)
_, predicted = torch.max(output.data, 1)
total += target.size(0)
correct += (predicted == target).sum().item()
print('Accuracy of the network on the 10000 test images: %d %%' % (
100 * correct / total))
上面的代码实现了PyTorch对MNIST数据集的手写数字识别。
image, label = train_dataset[0]
plt.imshow(image.squeeze().numpy(), cmap='gray')
plt.title('Label: %d' % label)
plt.show()
上面的代码显示了数据集中的一张图片。
image, label = test_dataset[0]
output = model(image.unsqueeze(0))
_, predicted = torch.max(output.data, 1)
plt.imshow(image.squeeze().numpy(), cmap='gray')
plt.title('Predicted: %d, Actual: %d' % (predicted.item(), label))
plt.show()
上面的代码显示了模型的预测结果。
总结:以上就是PyTorch实现MNIST数据集手写数字识别的完整攻略,包括数据集的加载、模型的定义、损失函数和优化器的定义、模型的训练和测试,以及两个示例的展示。
本文链接:http://task.lmcjl.com/news/16548.html