使用PyTorch构建神经网络：手写数字识别实战

文将通过构建一个简单的全连接神经网络，实现对MNIST手写数字数据集的分类任务。我们将从数据准备、网络结构设计、模型训练、可视化到性能分析的完整流程进行详细讲解，并探讨如何在PyTorch中使用nn.Module与nn.functional构建神经网络，帮助读者理解PyTorch构建神经网络的工具及其实现逻辑。

一、PyTorch构建神经网络的核心工具

在PyTorch中，构建神经网络主要依赖于两个核心模块：

• torch.nn.Module：是所有神经网络模块的基类，提供了模块化构建网络的能力。
• torch.nn.functional（简称F）：提供了大量神经网络操作函数，如激活函数、池化等，通常用于函数式编程风格。

1. nn.Module 与 nn.functional 的区别

✅ 使用建议：

• 含参数的层（如全连接层、卷积层、Dropout等）应使用nn.Module子类。
• 无参数操作（如ReLU、Sigmoid、池化等）推荐使用F函数。

二、实战：构建手写数字识别模型

我们将使用一个简单的前馈神经网络（Feedforward Neural Network），结构如下：

输入层（784） -> 隐藏层1（300） -> 隐藏层2（100） -> 输出层（10）
激活函数：ReLU
损失函数：交叉熵损失（CrossEntropyLoss）
优化器：随机梯度下降（SGD + Momentum）

三、数据准备与预处理

1. 导入必要模块

import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import matplotlib.pyplot as plt
import numpy as np

2. 下载并预处理MNIST数据集

定义预处理管道
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, transform=transform)
 
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False)

3. 数据可视化

examples = enumerate(test_loader)
batch_idx, (example_data, example_targets) = next(examples)
 
fig = plt.figure(figsize=(10, 6))
for i in range(6):
    plt.subplot(2, 3, i+1)
    plt.tight_layout()
    plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
    plt.title("Ground Truth: {}".format(example_targets[i]))
    plt.xticks([])
    plt.yticks([])
plt.show()

四、构建神经网络模型

我们将使用nn.Sequential将网络层组合起来，并在__init__中定义各层结构。

class Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super(Net, self).__init__()
        self.layer1 = nn.Sequential(
            nn.Linear(in_dim, n_hidden_1),
            nn.BatchNorm1d(n_hidden_1),
            nn.ReLU()
        )
        self.layer2 = nn.Sequential(
            nn.Linear(n_hidden_1, n_hidden_2),
            nn.BatchNorm1d(n_hidden_2),
            nn.ReLU()
        )
        self.layer3 = nn.Linear(n_hidden_2, out_dim)
        self.dropout = nn.Dropout(0.2)
 
    def forward(self, x):
        x = self.layer1(x)
        x = self.dropout(x)
        x = self.layer2(x)
        x = self.dropout(x)
        x = self.layer3(x)
        return x

🧠 说明：

• nn.BatchNorm1d用于加速训练并提升泛化性能。
• nn.Dropout用于防止过拟合。
• 激活函数ReLU使用nn.ReLU()类，也可以使用F.relu()函数。

五、实例化模型与设置训练参数

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Net(28 * 28, 300, 100, 10).to(device)
 
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.5)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)

六、训练模型

losses, acces = [], []
eval_losses, eval_acces = [], []
 
for epoch in range(20):
    model.train()
    train_loss, train_acc = 0, 0
    
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
        images = images.view(images.size(0), -1)  # 展平
 
        outputs = model(images)
        loss = criterion(outputs, labels)
 
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
 
        train_loss += loss.item()
        _, pred = outputs.max(1)
        train_acc += (pred == labels).sum().item() / images.shape[0]
 
    losses.append(train_loss / len(train_loader))
    acces.append(train_acc / len(train_loader))
 
    # 测试阶段 
    model.eval()
    eval_loss, eval_acc = 0, 0 
    with torch.no_grad():
        for images, labels in test_loader:
            images, labels = images.to(device), labels.to(device)
            images = images.view(images.size(0), -1)
            outputs = model(images)
            loss = criterion(outputs, labels)
 
            eval_loss += loss.item()
            _, pred = outputs.max(1)
            eval_acc += (pred == labels).sum().item() / images.shape[0]
 
    eval_losses.append(eval_loss / len(test_loader))
    eval_acces.append(eval_acc / len(test_loader))
 
    print(f'Epoch [{epoch+1}/20], Train Loss: {losses[-1]:.4f}, Test Acc: {eval_acces[-1]:.4f}')

七、可视化训练过程

plt.figure(figsize=(12, 4))
 
plt.subplot(1, 2, 1)
plt.plot(losses, label='Train Loss')
plt.plot(eval_losses, label='Test Loss')
plt.legend()
plt.title('Loss Curve')
plt.xlabel('Epoch')
plt.ylabel('Loss')
 
plt.subplot(1, 2, 2)
plt.plot(acces, label='Train Accuracy')
plt.plot(eval_acces, label='Test Accuracy')
plt.legend()
plt.title('Accuracy Curve')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
 
plt.tight_layout()
plt.show()

八、保存模型

torch.save(model.state_dict(), 'mnist_mlp.pth')
print("模型已保存至 mnist_mlp.pth")

九、结果分析与优化建议

1. 模型表现

• 训练准确率接近 100%，测试准确率最高可达 98.4%。
• 损失值稳定下降，未出现明显过拟合。

2. 优化建议

十、总结

本文从PyTorch构建神经网络的核心工具入手，逐步实现了MNIST手写数字识别的完整流程，包括数据加载、预处理、网络构建、训练、评估和可视化。通过本例，我们可以掌握PyTorch的基本使用方法，并理解神经网络的运行机制。

🧩 未来方向：

• 构建更复杂的模型（如CNN、RNN、Transformer）
• 引入迁移学习、模型微调等高级技巧
• 使用PyTorch Lightning提升代码结构和训练效率