PyTorch实现手写数字识别的示例代码

以下是“PyTorch实现手写数字识别的示例代码”的完整攻略，包含两个示例说明。

PyTorch实现手写数字识别的示例代码

手写数字识别是计算机视觉中的一个经典问题，它可以用于识别手写数字的图像。在PyTorch中，我们可以使用MNIST数据集来训练一个手写数字识别模型。下面是PyTorch实现手写数字识别的示例代码：

示例1：使用全连接层实现手写数字识别

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 创建模型实例
net = Net()

# 进行模型训练
for epoch in range(10):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

    # 打印训练日志
    print('Epoch: %d, Loss: %.4f' % (epoch+1, loss.item()))

# 进行模型测试
correct = 0
total = 0
with torch.no_grad():
    for inputs, targets in test_loader:
        outputs = net(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += (predicted == targets).sum().item()

# 打印测试结果
print('Accuracy: %.2f%%' % (100 * correct / total))

在这个示例中，我们首先加载了MNIST数据集，并使用DataLoader创建了数据加载器。然后，我们定义了一个包含两个全连接层的模型，并使用交叉熵损失函数和随机梯度下降优化器进行模型训练。最后，我们使用测试集对模型进行测试，并计算了模型的准确率。

示例2：使用卷积神经网络实现手写数字识别

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=128, shuffle=False)

# 定义模型
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=5, padding=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=5, padding=2)
        self.fc1 = nn.Linear(7*7*64, 1024)
        self.fc2 = nn.Linear(1024, 10)

    def forward(self, x):
        x = torch.relu(self.conv1(x))
        x = torch.max_pool2d(x, 2)
        x = torch.relu(self.conv2(x))
        x = torch.max_pool2d(x, 2)
        x = x.view(-1, 7*7*64)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

# 创建模型实例
net = Net()

# 进行模型训练
for epoch in range(10):
    for i, (inputs, targets) in enumerate(train_loader):
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, targets)
        loss.backward()
        optimizer.step()

    # 打印训练日志
    print('Epoch: %d, Loss: %.4f' % (epoch+1, loss.item()))

# 进行模型测试
correct = 0
total = 0
with torch.no_grad():
    for inputs, targets in test_loader:
        outputs = net(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += targets.size(0)
        correct += (predicted == targets).sum().item()

# 打印测试结果
print('Accuracy: %.2f%%' % (100 * correct / total))

在这个示例中，我们首先加载了MNIST数据集，并使用DataLoader创建了数据加载器。然后，我们定义了一个包含两个卷积层和两个全连接层的卷积神经网络，并使用交叉熵损失函数和随机梯度下降优化器进行模型训练。最后，我们使用测试集对模型进行测试，并计算了模型的准确率。

总结

本文介绍了PyTorch实现手写数字识别的示例代码，包括使用全连接层和卷积神经网络两种方法，并提供了两个示例说明。在实现过程中，我们使用了MNIST数据集，并使用DataLoader创建了数据加载器。然后，我们定义了模型、损失函数和优化器，并进行了模型训练和测试。