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

unicrom_深圳市由你创科技2025-08-06 16:46

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 nn

class 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.0
        for 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 = 0
    total = 0
    test_loss = 0.0
    with 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 / total
    print(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.0
        for 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 = 0
    total = 0
    test_loss = 0.0
    with 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 / total
    print(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')

9. 总结

通过本文，可以了解到如何用 PyTorch 实现一个简单的神经网络，包括：

  • 加载和预处理 MNIST 数据集。

  • 构建一个全连接神经网络。

  • 使用交叉熵损失和 Adam 优化器进行训练。

  • 在测试集上评估模型性能。

  • 对单张图像进行预测并可视化结果。

这个模型虽然简单，但在 MNIST 数据集上通常能达到 95% 以上的准确率。可以进一步尝试调整网络结构（如增加层数）、优化超参数（如学习率）或使用卷积神经网络（CNN）来提升性能。希望这篇文章对你理解 PyTorch 和深度学习有所帮助！

相关推荐
2501_924534891 小时前
智慧零售商品识别误报率↓74%！陌讯多模态融合算法在自助结算场景的落地优化
大数据·人工智能·算法·计算机视觉·目标跟踪·视觉检测·零售
盖雅工场1 小时前
连锁零售排班难？自动排班系统来解决
大数据·人工智能·物联网·算法·零售
bryant_meng3 小时前
【Apache MXNet】
人工智能·apache·mxnet
UMI赋能企业3 小时前
企业视频库管理高效策略
大数据·人工智能
一念&5 小时前
今日科技热点 | AI加速变革，量子计算商用化，5G应用新机遇
人工智能·科技·量子计算
Ysn07195 小时前
pytorch_grad_cam 库学习笔记—— Ablation-CAM 算法的基类 AblationCAM 和 AblationLayer
pytorch·笔记·学习
严文文-Chris5 小时前
【GPT-5 与 GPT-4 的主要区别？】
人工智能·gpt
过往入尘土6 小时前
计算机视觉：从 “看见” 到 “理解”，解锁机器感知世界的密码
人工智能
飞哥数智坊7 小时前
别再组团队了，AI时代一个人就能创业
人工智能·创业
严文文-Chris7 小时前
GPT5的Test-time compute（测试时计算）是什么？
人工智能
热门推荐
01UV安装并设置国内源02KGG转MP3工具|非KGM文件|解密音频03不再让Windows更新！&Edge游戏助手卸载及关闭自动更新04Qwen3-Coder 快速上手教程 | Qwen Code + Claude Code05【2025.08.06最新版】Android Studio下载、安装及配置记录（自动下载sdk）06Claude Code VSCode集成开发指南：AI编程助手完整配置07突破百度网盘的下载限速，两种方法教会你【超详细】08Linux下V2Ray安装配置指南09jdk21下载、安装（Windows、Linux、macOS）10如何在Windows/Linux系统上安装adb