手把手教你实现PyTorch的MNIST数据集

在本文中，我们将手把手教你如何使用PyTorch实现MNIST数据集的分类任务。我们将使用两个示例来说明如何完成这些步骤。

示例1：使用全连接神经网络实现MNIST分类

以下是使用全连接神经网络实现MNIST分类的步骤：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)

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

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

model = Net()

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

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_dataset, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_dataset)}')

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for data in test_dataset:
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += 1
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {correct / len(test_dataset)}')

在上述代码中，我们首先定义了数据预处理，然后加载了MNIST数据集。接着，我们定义了一个简单的全连接神经网络Net，它含有一个输入层、一个隐藏层和一个输出层。然后，我们创建了一个模型实例model。我们还定义了损失函数criterion和优化器optimizer。在训练模型的过程中，我们使用交叉熵损失函数nn.CrossEntropyLoss()来计算损失。在测试模型的过程中，我们使用torch.max()函数来获取输出中的最大值，并将其与标签进行比较以计算准确度。

示例2：使用卷积神经网络实现MNIST分类

以下是使用卷积神经网络实现MNIST分类的步骤：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = datasets.MNIST('data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST('data', train=False, download=True, transform=transform)

# 定义模型
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

model = Net()

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

# 训练模型
for epoch in range(10):
    running_loss = 0.0
    for i, data in enumerate(train_dataset, 0):
        inputs, labels = data
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item()

    print(f'Epoch {epoch + 1}, Loss: {running_loss / len(train_dataset)}')

# 测试模型
correct = 0
total = 0
with torch.no_grad():
    for data in test_dataset:
        inputs, labels = data
        outputs = model(inputs)
        _, predicted = torch.max(outputs.data, 1)
        total += 1
        correct += (predicted == labels).sum().item()

print(f'Accuracy: {correct / len(test_dataset)}')

在上述代码中，我们首先定义了数据预处理，然后加载了MNIST数据集。接着，我们定义了一个简单的卷积神经网络Net，它含有两个卷积层、两个池化层和两个全连接层。然后，我们创建了一个模型实例model。我们还定义了损失函数criterion和优化器optimizer。在训练模型的过程中，我们使用交叉熵损失函数nn.CrossEntropyLoss()来计算损失。在测试模型的过程中，我们使用torch.max()函数来获取输出中的最大值，并将其与标签进行比较以计算准确度。