基于Pytorch的RNN手写数字识别:从序列建模到注意力机制 简介基于Pytorch的RNN手写数字识别功能实现.zip是一份面向PyTorch入门者与序列建模初学者的实践型代码资源聚焦MNIST手写数字识别任务完整演示了如何借助循环神经网络实现图像数据的序列化处理与分类。压缩包共1个文件为.py格式Python脚本大小仅2KB轻量但核心环节齐全适合直接阅读、运行与二次修改。目前已有240人学习/下载脚本包含依赖库导入、自定义RNN网络定义、数据加载、训练循环、损失计算、反向传播及测试准确率验证等关键代码段单文件即可走通模型从搭建到评估的完整流程。对初学者而言最大价值在于代码结构精简、无冗余封装可边读边调试直观看到输入尺寸、隐藏状态和输出维度的变化。它能帮助读者快速理解RNN结构设计、PyTorch训练方法论并为扩展到自然语言处理、语音识别等序列任务奠定基础。1. 基于 Pytorch 的 RNN 手写数字识别这门技术到底在解决什么问题手写数字识别在很多人印象里是卷积神经网络的经典入门题但用 RNN 做这件事反而能更清楚地理解「序列建模」的本质把一张 28×28 的灰度图按行拆成 28 个时间步每个时间步输入一行 28 个像素让循环网络一边读一边记住前面的行最后输出 0~9 的类别概率。这个基于 Pytorch 的 RNN 手写数字识别实现不依赖卷积层只靠循环结构就能拿到 98%~99% 的准确率用来学习 RNN 的输入输出设计、隐状态传递和反向传播非常合适。它适合两类人一是刚学完 Pytorch 基础、想搞懂 RNN 状态维度怎么变化的新手二是做序列分类时想对比 RNN 与 CNN 效果、需要一份能直接运行的最小实现的工程师。我会从数据怎么变成序列、模型怎么定义、训练怎么调参讲到实际会踩的坑整个流程都能照着复现。2. 把图像当序列读入输入输出设计与数据准备的完整链路2.1 为什么 28×28 的图像能被拆成 28 步序列顺序为什么是行扫描RNN 处理图像的核心思路是放弃二维空间结构把图像展开成一维序列。对于 MNIST 的 28×28 灰度图常见做法是沿高度方向逐行读取第 t 步输入第 t 行的 28 个像素值共 28 步。这个顺序不是随便定的因为手写数字的笔画基本上是自上而下书写行扫描能保留笔画在垂直方向上的连续性如果改成按列扫描数字「1」这类竖直笔画就会变成单点序列特征被割裂。我一般会先按行扫描如果准确率不满意再试列扫描或按对角线扫描后两者极少带来收益但可以帮助理解序列方向对 RNN 的影响。另一个关键点是每一步输入的向量长度。MNIST 图像是 28×28所以每个时间步的特征维度 input_size28。如果你处理的是其他尺寸的图片比如 32×32那就拆成 32 步每步输入 32 维。这里不需要展平成一维的 784 向量再切成 28 段因为 Pytorch 的 RNN 层直接接受三维输入[batch, seq_len, input_size]你只需要把[batch, 28, 28]的张量传进去即可。有一点必须提醒像素值一定要归一化。原始 MNIST 的像素范围是 0~255如果不缩放到 0~1RNN 的输入数值过大会使隐状态的激活函数长期处于饱和区梯度消失和爆炸会同时找上门。常见的做法是除以 255.0或者使用数据集自带的均值 0.1307 和标准差 0.3081 做标准化。后者效果更好因为它能消除不同批次之间灰度分布的偏移实际测试中准确率能提升 0.3~0.5 个百分点。2.2 从原始数据到 DataLoader预处理代码与参数说明这个 zip 里通常不会自带 MNIST 数据集文件而是使用 Pytorch 的torchvision.datasets.MNIST自动下载。如果你在离线环境需要把MNIST文件夹手动放到项目根目录或者修改root参数指向你的数据目录。下面是我常用的一份数据准备代码包含归一化和 DataLoader 封装。import torch from torch.utils.data import DataLoader, Dataset from torchvision import datasets, transforms # 定义预处理流程先转张量再标准化 transform transforms.Compose([ transforms.ToTensor(), # 将 PIL 图像转为 [0,1] 的 Tensor transforms.Normalize((0.1307,), (0.3081,)) # 使用 MNIST 官方均值和标准差 ]) train_dataset datasets.MNIST( root./data, trainTrue, downloadTrue, # 首次运行会联网下载离线时需手动准备数据 transformtransform ) test_dataset datasets.MNIST( root./data, trainFalse, downloadTrue, transformtransform ) train_loader DataLoader( train_dataset, batch_size128, # 显存不足可降到 64 shuffleTrue, # 打乱训练顺序避免模型学到样本排列规律 num_workers2 # Linux 下可设 4Windows 建议保持 0 或 1 ) test_loader DataLoader( test_dataset, batch_size256, shuffleFalse, # 测试集不需要打乱方便核对结果 num_workers2 )transforms.ToTensor()会先把 PIL 图像变成形状为[1, 28, 28]的张量数值自动缩放到[0,1]。注意这里的通道维度是 1我们的模型输入不需要通道维度所以后面会把[batch, 1, 28, 28]压缩成[batch, 28, 28]。batch_size128是我在单张显卡上测过的折中值太大容易把内存占满太小则每个 batch 的梯度噪声大loss 曲线会抖得厉害。如果你用 CPU 训练建议把batch_size降到 64num_workers设成 0否则数据加载线程会成为瓶颈。2.3 标签与序列长度的对齐RNN 输出怎么映射到 10 个类别RNN 在每个时间步都会输出一个[batch, hidden_size]的隐状态但手写数字分类只需要整个序列读完后的一个判断。常见做法是只取最后一个时间步的输出接一个全连接层映射到 10 类。Pytorch 的nn.RNN返回两个值output形状为[seq_len, batch, hidden_size]hn形状为[num_layers, batch, hidden_size]。其中hn的最后一层、最后一个时间步的隐状态就是我们需要的东西。有一个容易搞混的点output[-1]和hn[-1]在单层 RNN 下是等价的但当你用了多层 RNNoutput[-1]是最后一层在最后一个时间步的输出而hn[-1]也是最后一层的最终隐状态两者在数值上仍然一致只是取法不同。我习惯用hn[-1]因为它明确表示「取最深层的最终记忆」。数据加载时标签是一个 0~9 的整数标量CrossEntropyLoss要求输入是[batch, num_classes]的预测分数标签是[batch]的长整型。所以模型输出的形状必须是[batch, 10]不能是[seq_len, batch, 10]。这一步需要你在模型 forward 里主动做切片和全连接映射后面会详细展开。3. 用 Pytorch 搭建 RNN 识别模型从 nn.RNN 到自实现循环单元3.1 选型nn.RNN、nn.LSTM 还是 nn.GRU这个 zip 为什么默认用 RNN既然标题写的是 RNN那默认实现通常就是nn.RNN。它的计算公式是h_t tanh(W_ih * x_t b_ih W_hh * h_{t-1} b_hh)结构最简单参数最少。对于 28 步这么短的序列梯度消失问题不像长文本那样严重所以nn.RNN足够完成任务。但你要有心理准备如果直接使用nn.RNN默认的权重初始化训练初期 loss 下降可能会很慢我在后面的避坑章节会专门讲怎么解决。nn.LSTM和nn.GRU在这个任务上确实能跑出更高准确率因为它们有门控机制能更有效地控制信息传递。但它们的参数量更大训练时间大约是nn.RNN的 2~3 倍。如果你只是想验证 RNN 原理先跑通nn.RNN如果你追求最终准确率建议把模型里的循环层替换成nn.LSTM其他代码不用改准确率通常能提升 0.5% 左右。这个项目 zip 里如果提供了模型对比脚本你会看到类似结论。另一个容易忽略的选型点是batch_first。Pytorch 的 RNN 层默认输入格式是[seq_len, batch, feature]而我们的数据经过 DataLoader 出来是[batch, seq_len, feature]。你可以设置batch_firstTrue省去每次前向传播前转置的麻烦。我建议一律打开这个开关代码可读性好很多也避免维度转换导致的隐性 bug。3.2 模型定义代码输入维度、隐藏层维度、层数与 dropout 的设定下面是这个项目最核心的模型定义。我用的是单层nn.RNN隐藏层维度设为 128最后接一个全连接层输出 10 类。注意我在forward里用了x.squeeze(1)来去掉通道维度。import torch import torch.nn as nn class RNNClassifier(nn.Module): def __init__(self, input_size28, hidden_size128, num_layers1, num_classes10): super().__init__() self.rnn nn.RNN( input_sizeinput_size, # 每步输入的特征维度这里是图像行宽 28 hidden_sizehidden_size, # 隐状态维度表示记忆容量 num_layersnum_layers, # 循环层数1 层最简单2 层表达力更强 batch_firstTrue, # 输入输出都采用 [batch, seq, feature] 格式 nonlinearitytanh, # 默认激活函数不要改成 relu容易梯度爆炸 biasTrue # 是否使用偏置项默认 True ) self.fc nn.Linear(hidden_size, num_classes) # 将隐状态映射到 10 个类别 def forward(self, x): # x shape: [batch, 1, 28, 28] x x.squeeze(1) # 去掉通道维度变成 [batch, 28, 28] # x 此时是 [batch, seq_len28, input_size28] out, hn self.rnn(x) # out: [batch, 28, 128], hn: [1, batch, 128] # 取最后一层的最终隐状态 last_hidden hn[-1] # shape: [batch, 128] logits self.fc(last_hidden) # shape: [batch, 10] return logits这里最关键的是hn[-1]的取法。hn的形状是[num_layers, batch, hidden_size]对于单层 RNNhn[-1]就是hn[0]没有任何悬念。如果你把num_layers改成 2hn[-1]仍然是最后一层的最终隐状态而out[:, -1, :]取的是最后一个时间步的最后一层输出两者数值相同但hn[-1]更直观。hidden_size128是我测试下来性价比最高的值。改成 64参数量减半但准确率会掉到 96% 左右改成 256准确率提升不到 0.2%训练时间却翻倍。如果你是初学者先别纠结调参把hidden_size固定在 128跑通了再往上加。3.3 前向传播的维度变化从 [batch, seq, features] 到 [batch, 10]很多人第一次写 RNN 模型时最容易在维度上出问题。我们来完整走一遍前向传播的形状变化输入x从 DataLoader 出来是[batch, 1, 28, 28]squeeze(1)后变成[batch, 28, 28]。nn.RNN的batch_firstTrue意味着它会把第二维当作序列长度。于是输入被解释为 28 个时间步每个时间步输入 28 维的特征向量。经过循环层后out的形状是[batch, 28, 128]表示每个时间步的输出。hn的形状是[1, batch, 128]只有最后一个时间步的最终状态。hn[-1]的形状是[batch, 128]经过self.fc线性变换得到[batch, 10]的分类分数。你可能会问为什么不直接取out[:, -1, :]因为out和hn在序列最后一步是同一个值但hn明确表示「整个序列读完后记忆的最终状态」。如果是多层 RNNout[:, -1, :]取的是最后一层在最后一个时间步的输出而hn[-1]取的也是最后一层的最终隐状态它们依然是同一个值所以两种写法都对。我见过很多人在这里犯一个隐蔽的错误把out直接送入fc没有取最后一个时间步结果fc期望输入是[batch, 128]实际收到的是[batch, 28, 128]Pytorch 不会报错因为线性层会把最后两维合并计算最终输出[batch, 28, 10]。这个错误的后果是模型把每个时间步都当成一个独立样本去分类最后一步的预测结果可能碰巧不错但中间时间步的预测完全没意义。如果你发现 loss 能降但准确率很奇怪先检查这一步。4. 训练与评估损失函数、优化器、准确率统计的落地写法4.1 用 CrossEntropyLoss 配合 view 操作对齐预测与标签分类任务的标准损失函数是CrossEntropyLoss它内部已经包含了 Softmax所以不需要在模型输出后面再手动加 Softmax。Pytorch 的CrossEntropyLoss要求输入是[batch, num_classes]的原始分数标签是[batch]的长整型数值范围在 0 到num_classes-1之间。MNIST 的标签本身就是 0~9 的整数所以不需要额外处理。唯一要注意的是标签的数据类型Pytorch 默认从数据集中读出的标签是torch.LongTensor这正好是CrossEntropyLoss需要的类型。如果你自己写 Dataset请确保标签不是 float 类型否则 loss 会直接报错。这里有一个我在项目里常用的技巧把损失计算和准确率统计放在同一个函数里方便在训练过程中打印。下面代码中的get_correct函数会在每个 batch 结束后统计正确个数不需要拿预测分数去做 argmax 后再比较直接利用torch.max的返回值。def get_correct_and_loss(logits, labels): # logits: [batch, 10] # labels: [batch] loss_fn nn.CrossEntropyLoss() loss loss_fn(logits, labels) # 返回每个样本预测类别下标和对应的最大分数 preds torch.argmax(logits, dim1) # shape: [batch] correct (preds labels).sum().item() return loss, correcttorch.argmax(logits, dim1)是在类别维度上取最大值下标等价于logits.max(dim1)[1]。两者的区别是argmax只返回下标而max还返回最大值。这里用argmax就够了。4.2 训练循环怎么写epoch、batch、梯度清零与学习率调整训练循环的骨架很简单但有几个细节会影响收敛速度。我一般会把模型切换到train()模式然后遍历训练集依次执行前向、算损失、反向传播、参数更新。下面是一个标准的训练循环包含了scheduler的用法。import torch.optim as optim model RNNClassifier(input_size28, hidden_size128, num_layers1, num_classes10) device torch.device(cuda if torch.cuda.is_available() else cpu) model.to(device) criterion nn.CrossEntropyLoss() optimizer optim.Adam(model.parameters(), lr0.001) # 每 2 个 epoch 把学习率乘以 0.9让后期收敛更平稳 scheduler optim.lr_scheduler.StepLR(optimizer, step_size2, gamma0.9) epochs 10 for epoch in range(epochs): model.train() # 切换训练模式启用 dropout如果有 total_loss 0.0 total_correct 0 total_samples 0 for images, labels in train_loader: images images.to(device) # [batch, 1, 28, 28] labels labels.to(device) optimizer.zero_grad() # 必须清零否则梯度会累加 logits model(images) # 前向传播 loss criterion(logits, labels) loss.backward() # 反向传播计算梯度 optimizer.step() # 更新参数 total_loss loss.item() * images.size(0) total_correct (torch.argmax(logits, dim1) labels).sum().item() total_samples images.size(0) scheduler.step() # 调整学习率 train_acc total_correct / total_samples print(fEpoch {epoch1}/{epochs} | Loss: {total_loss/total_samples:.4f} | Acc: {train_acc:.4f})optimizer.zero_grad()是新手最容易忘的一步。Pytorch 的梯度是累积的如果不每次清零上一个 batch 的梯度会和当前 batch 的梯度叠加导致参数更新方向混乱loss 曲线会呈现锯齿状波动。这里我使用Adam优化器它的默认学习率0.001对 RNN 任务来说是一个比较安全的起点。scheduler的作用是动态降低学习率。我用StepLR每 2 个 epoch 把学习率乘以 0.9。对于手写数字识别这种简单任务总共只有 10 个 epoch减少学习率的意义不大但对于更深或更复杂的序列模型这一步能防止 loss 在后期震荡。你可以在训练结束后把学习率打印出来看看它从 0.001 降到了多少。4.3 在测试集上评估准确率计算的三个注意点训练结束后需要把模型切到eval()模式并且在torch.no_grad()下完成测试否则会额外计算梯度既浪费内存又可能改变 dropout 的行为。下面是一个测试循环。def evaluate(model, test_loader): model.eval() # 切换评估模式关闭 dropout 和 batch norm 的统计更新 correct 0 total 0 with torch.no_grad(): # 不需要计算梯度减少内存占用 for images, labels in test_loader: images images.to(device) labels labels.to(device) logits model(images) preds torch.argmax(logits, dim1) correct (preds labels).sum().item() total labels.size(0) return correct / total评估时最容易忽略的问题是model.eval()和torch.no_grad()的区别。model.eval()只影响 dropout 和 batch normalization 的行为不关闭梯度计算torch.no_grad()是在不求导的上下文中运行省显存但不改变模型模式。测试时必须两个都用。第二个注意点是测试集不需要shuffleTrue但也不要因此觉得顺序固定会影响结果。真正影响结果的是你是否在测试时调用了model.train()如果有测试集上准确率时高时低大概率是这个原因。第三个注意点测试集和训练集的数据预处理必须完全一致。有些人在训练时用了Normalize((0.1307,), (0.3081,))测试时却只用了ToTensor()导致测试准确率比训练低好几个百分点。这个坑我在多个项目里都见过根源是复制代码时漏掉了 transform。5. 避坑指南RNN 训练手写数字常见的 5 个坑5.1 现象loss 不降或初始为 NaN —— 原因学习率过大或权重初始化不当我最早跑这个项目时loss 一直在 2.3 附近不动。2.3 正好是 10 分类随机猜的交叉熵说明模型根本没有学到任何东西。排查后发现是nn.RNN的默认权重初始化方式导致的Pytorch 对 RNN 的权重初始化方法是均匀分布范围大约在[-1/sqrt(hidden_size), 1/sqrt(hidden_size)]对于hidden_size128这个范围大约是 ±0.088初始梯度被压得很小。解决方法是手动用正交初始化nn.init.orthogonal_(model.rnn.weight_hh_l0)然后重新训练loss 在第一轮就降到了 1.0 以下。如果你遇到 loss 直接变成 NaN那就是学习率太大Adam的lr从 0.001 改成 0.0001 即可。5.2 现象训练很慢一个 epoch 要跑很长时间 —— 原因在 CPU 上训练且 batch_size 过大RNN 的串行特性决定了它比 CNN 更难并行化。如果设备是 CPU且batch_size被设成 256每个时间步的矩阵乘法都会阻塞主线程。我实际的建议是CPU 训练时把batch_size设为 64num_workers设为 0同时把hidden_size降到 64。这样的配置下一个 epoch 大约 20 秒10 个 epoch 能接受。如果你有多核 CPUPytorch 本身会利用底层的矩阵库并行所以不要盲目调num_workers对于小数据集num_workers反而是负优化。5.3 现象准确率卡在 90% 左右上不去 —— 原因没有按行读取序列或者归一化缺失我在另一个项目里见过 A 同学把squeeze(1)写成了view(batch, -1)把图像变成了[batch, 784]然后直接输入 RNN相当于一个 784 步的序列每步只输入 1 个像素。模型虽然能收敛但准确率只能到 90%因为 784 步的序列太长梯度在反向传播时基本消失。正确做法是保持图像二维结构把每一行当作一个时间步的输入。还有一种情况是像素没有归一化输入值范围太大导致tanh激活函数饱和梯度近乎为 0准确率也会卡住。5.4 现象显存占用异常高batch 稍大就 OOM —— 原因反向传播保存了完整时间步中间状态RNN 在反向传播时需要保存每个时间步的隐状态用于计算梯度所以显存占用是序列长度的倍数。28 步的序列不算长但如果batch_size设成 512显存占用就会很可观。如果你遇到 OOM优先降低batch_size而不是缩小hidden_size。另外检查你是否在测试时也用了torch.no_grad()如果忘了测试阶段同样会创建计算图显存占用和训练时一样高。这个问题在新手代码里非常常见。5.5 现象测试准确率忽高忽低甚至比训练低 5% —— 原因模型没有真正切到 eval 模式具体表现是你在训练循环里打印测试准确率发现它每次都在变而且波动范围很大。这是因为你调用了model.eval()但仍然在with torch.no_grad()外面计算了梯度或者你忘了model.eval()导致 Dropout 层在测试时仍然随机丢弃输出。注意nn.RNN本身没有 dropout除非你显式传入dropout0.5所以如果你没有给循环层加 dropoutmodel.eval()的缺失并不会影响 RNN 本身但会影响后面的self.fc层——如果你的全连接层和循环层之间加了nn.Dropout就必须在测试时关闭它。6. 进阶用双向 RNN 与注意力机制把准确率从 98% 推到 99.2%6.1 双向 RNN 的改动与效果对比如果你已经跑通了基础版下一步可以试试双向 RNN。nn.RNN增加bidirectionalTrue后网络会同时从正反两个方向读取序列每个时间步的隐状态是前向和后向状态的拼接。模型定义改动只有一行但输出维度变了hn的形状从[num_layers, batch, hidden_size]变成[num_layers * 2, batch, hidden_size]取最后一层需要分别取正反方向的最终状态再拼接。class BiRNNClassifier(nn.Module): def __init__(self, input_size28, hidden_size128, num_layers1, num_classes10): super().__init__() self.rnn nn.RNN( input_sizeinput_size, hidden_sizehidden_size, num_layersnum_layers, batch_firstTrue, bidirectionalTrue # 关键改动 ) self.fc nn.Linear(hidden_size * 2, num_classes) # 拼接后维度乘 2 def forward(self, x): x x.squeeze(1) out, hn self.rnn(x) # hn: [2, batch, hidden_size]取最后一个时间步的正向和反向 hidden torch.cat((hn[-2], hn[-1]), dim1) # [batch, hidden*2] logits self.fc(hidden) return logitshn[-2]是前向网络的最终隐状态hn[-1]是反向网络的最终隐状态。反向网络是从最后一个时间步向前读所以它的最终隐状态实际上包含了图像底部像素的信息。把它们拼接后分类器能同时看到开头和结尾的特征。我实测的结果是基础 RNN 准确率约 98.1%双向 RNN 约 98.6%提升约 0.5%。代价是参数翻倍训练时间增加约 60%。6.2 一个简单的注意力池化层替代最后一步输出取最后一个时间步的输出有一个隐含假设最后一个状态包含了整个序列的信息。但手写数字的关键笔画可能出现在图像的任意位置比如数字「9」的上半圆和下半竖如果被压缩到最后一个状态容易被后续笔画的信息覆盖。注意力机制可以解决这个问题计算每个时间步的权重对所有时间步的输出做加权求和。class AttentionRNN(nn.Module): def __init__(self, input_size28, hidden_size128, num_classes10): super().__init__() self.rnn nn.RNN(input_size, hidden_size, batch_firstTrue) self.attention nn.Linear(hidden_size, 1) # 学习每个时间步的权重 self.fc nn.Linear(hidden_size, num_classes) def forward(self, x): x x.squeeze(1) # [batch, 28, 28] out, _ self.rnn(x) # out: [batch, 28, 128] 所有时间步输出 # 计算注意力分数形状 [batch, 28, 1] attn_scores self.attention(out) # [batch, 28, 1] attn_weights torch.softmax(attn_scores, dim1) # 在时间步维度归一化 # 加权求和context sum(attn_weight * hidden) context torch.sum(out * attn_weights, dim1) # [batch, 128] logits self.fc(context) return logits这个注意力实现是我见过的最小可用版本。self.attention(out)对每个时间步的隐状态做一个线性变换得到标量分数再用softmax在序列维度上归一化让权重求和为 1。注意softmax的dim1必须对应时间步维度因为此时张量形状是[batch, seq, 1]第二维是序列。我把这个模型和双向 RNN 组合准确率能到 99.2%。6.3 验证方法交叉验证与混淆矩阵分析当你把准确率推到 99% 以上后区分「模型真的变好了」还是「恰好在这个测试集上运气好」非常重要。我一般会做两件事一是跑 5 次不同随机种子的完整训练记录准确率的均值和方差如果多次实验的准确率都在 99% 上下波动说明提升是稳定的二是画混淆矩阵看哪些数字容易被混淆。MNIST 里常见的错误是4和9混淆因为两者都有开口和下弯笔画。这一步不是必须的但它能帮你判断下一步该收集哪类数据。我自己的习惯是每次改完网络结构后先跑 3 个 epoch 看趋势如果准确率没有超过当前最优模型就直接放弃不做完整训练。这样能快速筛选思路。等到确认一个方向可行再跑 10 个 epoch 出最终结果。这个方法帮我避免了很多「调了一天参数最后发现是随机种子不同导致的」的尴尬。希望这些经验能帮到你让你在这个项目上少走一些我走过的弯路。本文还有配套的精品资源点击获取