用 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 和深度学习有所帮助！

相关推荐
zhangfeng1133几秒前
把“距离过近”的节点（或端点）合并成一个，避免重复。机器学习 python
人工智能·python·机器学习
AscentStream16 分钟前
技术文档 | 使用 Pulsar Functions 构建实时 AI Pipeline
人工智能
POLOAPI17 分钟前
为什么Claude Code让传统IDE开发者"失业"？深度解析AI编程的技术革命
人工智能·ai编程·claude
沐森18 分钟前
屏幕截图 (OCR今日ai下必备工具)
人工智能
YBCarry_段松啓24 分钟前
DeerFlow：深度研究的多智能体框架
人工智能·开源·llm
玄明Hanko24 分钟前
不想写测试脚本AI帮你搞定自动化测试
人工智能·测试
lingling00942 分钟前
光伏清洗机器人是什么？艾利特协作机器人如何重塑新能源运维效率
大数据·运维·人工智能
无限大.1 小时前
《计算机“十万个为什么”》之 面向对象 vs 面向过程：编程世界的积木与流水线
网络·人工智能·python
wangjun51591 小时前
人工智能、机器学习、深度学习、大模型、智能体知识点汇总
人工智能·深度学习·机器学习
音视频牛哥1 小时前
无人机 × 巡检 × AI识别：一套可复制的超低延迟低空视频感知系统搭建实践
人工智能·音视频·无人机·大牛直播sdk·低空感知·无人机视频回传·ai边缘识别
热门推荐
01Qwen3-Coder 快速上手教程 | Qwen Code + Claude Code02全球最强模型Grok4，国内已可免费使用！（附教程）03UV安装并设置国内源04腾讯还是太全面了，限时免费！超全CodeBuddy IDE保姆级教程！（附案例）05KGG转MP3工具|非KGM文件|解密音频06如何在 Cursor 中继续使用 Claude07Cursor 终端“卡死/无响应”问题的解法08VMware Workstation Pro虚拟机的下载和安装图文保姆级教程(附下载链接)09Claude Code+智谱GLM-4.5接入指南（附实测体验）10NVIDIA显卡驱动、CUDA、cuDNN 和 TensorRT 版本匹配指南