PyTorch中使用MNIST数据集进行训练和预测

1. 导入必要的库：

   • PyTorch用于构建和训练神经网络。
   • torchvision用于处理图像数据。

2. 加载MNIST数据集：

   • 使用torchvision.datasets模块中的MNIST类加载MNIST数据集。

3. 定义神经网络模型：

   • 构建一个简单的神经网络模型，例如多层感知机（MLP）。

4. 定义损失函数和优化器：

   • 选择适当的损失函数（如交叉熵损失）和优化器（如随机梯度下降SGD）。

5. 训练模型：

   • 使用训练数据集对模型进行训练。

6. 评估模型性能：

   • 使用测试数据集评估模型性能。

以下是一个简单的例子：

import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader

# 设置随机种子以保持结果的一致性
torch.manual_seed(42)

# 定义简单的多层感知机（MLP）模型
class MLP(nn.Module):
    def __init__(self):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(28 * 28, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)  # 将图像展平为一维向量
        x = self.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 设置数据变换
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])

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

# 创建数据加载器
train_loader = DataLoader(dataset=train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(dataset=test_dataset, batch_size=64, shuffle=False)

# 初始化模型、损失函数和优化器
model = MLP()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 训练模型
num_epochs = 5
for epoch in range(num_epochs):
    model.train()
    for images, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

    # 在每个epoch结束后输出训练集上的损失
    print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item():.4f}')

# 在测试集上评估模型
model.eval()
correct = 0
total = 0
with torch.no_grad():
    for images, labels in test_loader:
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

accuracy = correct / total
print(f'Test Accuracy: {accuracy * 100:.2f}%')