共计 3487 个字符,预计需要花费 9 分钟才能阅读完成。
PyTorch 实战指南:从零搭建前馈 / 卷积 / 循环神经网络实现 MNIST 手写数字识别
为什么选择 MNIST 数据集?
MNIST 手写数字数据集是深度学习领域的经典入门数据集,包含 60,000 张训练图像和 10,000 张测试图像。每张图像都是 28×28 像素的灰度图,表示 0 到 9 的手写数字。这个数据集大小适中,既不会因为太大而难以处理,也不会因为太小而无法展示深度学习模型的优势。

三种神经网络架构对比
1. 前馈神经网络(FNN)
FNN 是最基础的神经网络结构,由全连接层组成。在处理 MNIST 时,我们需要将 28×28 的图像展平为 784 维的向量作为输入。
- 优点:结构简单,易于实现
- 局限:忽略图像的空间局部相关性,参数数量庞大
2. 卷积神经网络(CNN)
CNN 通过局部感知和权重共享机制,特别适合处理图像数据。
- 优点:自动提取空间特征,参数共享减少计算量
- 适用场景:图像分类、目标检测等
3. 循环神经网络(RNN)
RNN 设计用于处理序列数据,虽然不如图像分类的首选,但可以探索将图像作为像素序列处理。
- 优点:能捕捉序列中的时间依赖性
- 局限:对长序列处理能力有限
代码实现详解
数据准备
import torch
from torchvision import datasets, transforms
# 数据预处理
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 = torch.utils.data.DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=1000, shuffle=False)
FNN 模型实现
import torch.nn as nn
import torch.nn.functional as F
class FNN(nn.Module):
def __init__(self):
super(FNN, self).__init__()
self.fc1 = nn.Linear(784, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 10)
self.dropout = nn.Dropout(0.2)
def forward(self, x):
x = x.view(-1, 784) # 展平图像
x = F.relu(self.fc1(x))
x = self.dropout(x)
x = F.relu(self.fc2(x))
x = self.dropout(x)
x = self.fc3(x)
return F.log_softmax(x, dim=1)
CNN 模型实现
class CNN(nn.Module):
def __init__(self):
super(CNN, self).__init__()
self.conv1 = nn.Conv2d(1, 32, 3, 1)
self.conv2 = nn.Conv2d(32, 64, 3, 1)
self.dropout1 = nn.Dropout2d(0.25)
self.dropout2 = nn.Dropout2d(0.5)
self.fc1 = nn.Linear(9216, 128)
self.fc2 = nn.Linear(128, 10)
def forward(self, x):
x = self.conv1(x)
x = F.relu(x)
x = self.conv2(x)
x = F.relu(x)
x = F.max_pool2d(x, 2)
x = self.dropout1(x)
x = torch.flatten(x, 1)
x = self.fc1(x)
x = F.relu(x)
x = self.dropout2(x)
x = self.fc2(x)
return F.log_softmax(x, dim=1)
RNN 模型实现
class RNN(nn.Module):
def __init__(self):
super(RNN, self).__init__()
self.rnn = nn.LSTM(28, 128, batch_first=True)
self.fc = nn.Linear(128, 10)
def forward(self, x):
x = x.squeeze(1) # 移除通道维度
x, _ = self.rnn(x)
x = x[:, -1, :] # 取最后一个时间步的输出
x = self.fc(x)
return F.log_softmax(x, dim=1)
训练循环
def train(model, device, train_loader, optimizer, epoch):
model.train()
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
output = model(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % 100 == 0:
print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}]\tLoss: {loss.item():.6f}')
性能优化技巧
Batch Size 的影响
较大的 batch size 可以:
– 提高 GPU 利用率
– 使梯度估计更稳定
– 但可能降低模型泛化能力
学习率策略
- StepLR: 固定步长衰减
- CosineAnnealing: 余弦退火
# StepLR 示例
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)
# CosineAnnealing 示例
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
GPU 分析工具
with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA],
schedule=torch.profiler.schedule(wait=1, warmup=1, active=3),
on_trace_ready=torch.profiler.tensorboard_trace_handler('./log'),
record_shapes=True
) as prof:
for step, batch_data in enumerate(train_loader):
if step >= (1 + 1 + 3):
break
train_step(batch_data)
prof.step()
常见问题与解决方案
维度错误
- 检查输入张量的 shape
- 使用
print(x.shape)调试 - 确保 view 操作后的维度匹配
梯度爆炸
- 使用梯度裁剪
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
过拟合
- 增加 Dropout 层
- 使用数据增强
- 添加 L2 正则化
延伸思考
- 模型部署:将训练好的模型导出为 ONNX 格式
- 数据不平衡:使用类别权重或过采样技术
- 迁移学习:在 CIFAR-10 数据集上测试模型
# ONNX 导出示例
torch.onnx.export(model, dummy_input, "model.onnx", verbose=True)
总结
通过本教程,我们实现了三种不同类型的神经网络来处理 MNIST 手写数字识别任务。每种网络都有其特点和适用场景。在实践中,CNN 通常是图像分类任务的首选,但了解不同网络结构的特性对于成为一名全面的深度学习工程师至关重要。建议读者尝试调整超参数,观察模型性能的变化,并考虑如何在其他数据集上应用这些知识。
正文完
发表至: 未分类
近两天内
