当前位置：首页 > news >正文

用 PyTorch 实现一个简单的神经网络：从数据到预测

news 2025/8/6 8:05:48

PyTorch 是目前最流行的深度学习框架之一，以其灵活性和易用性受到开发者的喜爱。本文将带你从零开始，用 PyTorch 实现一个简单的神经网络，用于解决经典的 MNIST 手写数字分类问题。我们将涵盖数据准备、模型构建、训练和预测的完整流程，并提供可运行的代码示例。

1. 环境准备

首先，确保你已安装 PyTorch 和相关依赖。本例使用 Python 3.8+ 和 PyTorch。你可以通过以下命令安装：

pip install torch torchvision

我们将使用 MNIST 数据集，它包含 28x28 像素的手写数字图像（0-9），目标是训练一个神经网络来识别这些数字。

2. 数据准备

MNIST 数据集可以通过 PyTorch 的 torchvision 模块直接加载。我们需要将数据加载为张量，并进行归一化处理以加速训练。

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms# 定义数据预处理：将图像转换为张量并归一化
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))  # MNIST 的均值和标准差
])# 加载 MNIST 数据集
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

代码说明：

transforms.ToTensor() 将图像转换为 PyTorch 张量，并将像素值从 [0, 255] 缩放到 [0, 1]。
transforms.Normalize 标准化数据，加速梯度下降收敛。
DataLoader 用于批量加载数据，batch_size=64 表示每次处理 64 张图像。

3. 构建神经网络

我们将定义一个简单的全连接神经网络，包含两个隐藏层，适合处理 MNIST 的分类任务。

import torch.nn as nnclass SimpleNN(nn.Module):def __init__(self):super(SimpleNN, self).__init__()self.flatten = nn.Flatten()  # 将 28x28 图像展平为 784 维向量self.fc1 = nn.Linear(28 * 28, 128)  # 第一个全连接层self.relu = nn.ReLU()  # 激活函数self.fc2 = nn.Linear(128, 64)  # 第二个全连接层self.fc3 = nn.Linear(64, 10)   # 输出层，10 个类别（0-9）def forward(self, x):x = self.flatten(x)x = self.relu(self.fc1(x))x = self.relu(self.fc2(x))x = self.fc3(x)return x# 实例化模型
model = SimpleNN()

代码说明：

nn.Module 是 PyTorch 模型的基类，自定义模型需要继承它。
forward 方法定义了前向传播的计算流程。
网络结构：输入层 (784) → 隐藏层1 (128) → ReLU → 隐藏层2 (64) → ReLU → 输出层 (10)。

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

我们使用交叉熵损失（适合分类任务）和 Adam 优化器来训练模型。

import torch.optim as optim# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

代码说明：

nn.CrossEntropyLoss 结合了 softmax 和负对数似然损失，适合多分类任务。
Adam 优化器以 0.001 的学习率更新模型参数。

5. 训练模型

接下来，我们训练模型 5 个 epoch，观察损失变化。

def train(model, train_loader, criterion, optimizer, epochs=5):model.train()  # 切换到训练模式for epoch in range(epochs):running_loss = 0.0for images, labels in train_loader:optimizer.zero_grad()  # 清零梯度outputs = model(images)  # 前向传播loss = criterion(outputs, labels)  # 计算损失loss.backward()  # 反向传播optimizer.step()  # 更新参数running_loss += loss.item()print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")# 开始训练
train(model, train_loader, criterion, optimizer)

代码说明：

model.train() 启用训练模式（影响 dropout 和 batch norm 等层）。
每次迭代清零梯度、计算损失、反向传播并更新参数。
每轮 epoch 打印平均损失。

6. 测试模型

训练完成后，我们在测试集上评估模型的准确率。

def test(model, test_loader, criterion):model.  # 切换到评估模式correct = 0total = 0test_loss = 0.0with torch.no_grad():  # 禁用梯度计算for images, labels in test_loader:outputs = model(images)loss = criterion(outputs, labels)test_loss += loss.item()_, predicted = torch.max(outputs.data, 1)  # 获取预测类别total += labels.size(0)correct += (predicted == labels).sum().item()accuracy = 100 * correct / totalprint(f"Test Loss: {test_loss/len(test_loader):.4f}, Accuracy: {accuracy:.2f}%")# 测试模型
test(model, test_loader, criterion)

代码说明：

model. 切换到评估模式，禁用 dropout 等。
使用 torch.no_grad() 减少内存消耗。
计算测试集的损失和准确率。

7. 进行预测

最后，我们用训练好的模型对单张图像进行预测。

import matplotlib.pyplot as plt# 获取一张测试图像
dataiter = iter(test_loader)
images, labels = next(dataiter)
image, label = images[0], labels[0]# 预测
model.
with torch.no_grad():output = model(image.unsqueeze(0))  # 增加 batch 维度_, predicted = torch.max(output, 1)# 显示图像和预测结果
plt.imshow(image.squeeze(), cmap='gray')
plt.title(f"Predicted: {predicted.item()}, Actual: {label.item()}")
plt.savefig('prediction.png')  # 保存图像

代码说明：

从测试集取一张图像，调用模型进行预测。
使用 Matplotlib 显示图像及其预测结果，保存为 PNG 文件。

8. 完整代码

以下是完整的可运行代码，整合了上述所有步骤：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt# 数据准备
transform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.1307,), (0.3081,))
])
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)# 定义模型
class SimpleNN(nn.Module):def __init__(self):super(SimpleNN, self).__init__()self.flatten = nn.Flatten()self.fc1 = nn.Linear(28 * 28, 128)self.relu = nn.ReLU()self.fc2 = nn.Linear(128, 64)self.fc3 = nn.Linear(64, 10)def forward(self, x):x = self.flatten(x)x = self.relu(self.fc1(x))x = self.relu(self.fc2(x))x = self.fc3(x)return x# 实例化模型、损失函数和优化器
model = SimpleNN()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)# 训练函数
def train(model, train_loader, criterion, optimizer, epochs=5):model.train()for epoch in range(epochs):running_loss = 0.0for images, labels in train_loader:optimizer.zero_grad()outputs = model(images)loss = criterion(outputs, labels)loss.backward()optimizer.step()running_loss += loss.item()print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader):.4f}")# 测试函数
def test(model, test_loader, criterion):model.correct = 0total = 0test_loss = 0.0with torch.no_grad():for images, labels in test_loader:outputs = model(images)loss = criterion(outputs, labels)test_loss += loss.item()_, predicted = torch.max(outputs.data, 1)total += labels.size(0)correct += (predicted == labels).sum().item()accuracy = 100 * correct / totalprint(f"Test Loss: {test_loss/len(test_loader):.4f}, Accuracy: {accuracy:.2f}%")# 训练和测试
train(model, train_loader, criterion, optimizer)
test(model, test_loader, criterion)# 预测单张图像
dataiter = iter(test_loader)
images, labels = next(dataiter)
image, label = images[0], labels[0]
model.
with torch.no_grad():output = model(image.unsqueeze(0))_, predicted = torch.max(output, 1)
plt.imshow(image.squeeze(), cmap='gray')
plt.title(f"Predicted: {predicted.item()}, Actual: {label.item()}")
plt.savefig('prediction.png')