很多 RNN 教程一上来就讲 LSTM、GRU、注意力机制,但最核心的 RNN 原理其实非常简单。本文用一个不到 80 行的 PyTorch 示例,带你从零实现一个完整的 RNN,包括模型定义、训练和测试。
我们要做什么?
给定一个数字序列 [a, b, c],让 RNN 预测下一个数字 d,其中 d = a + b + c。
比如:
- 输入
[1, 2, 3],预测6 - 输入
[2, -1, 4],预测5 - 输入
[-3, 1, 2],预测0
模型事先不知道求和规则,它需要通过大量训练数据自己发现这个规律。
为什么选这个任务?
这个任务足够简单,但涵盖了 RNN 的三个核心能力:
- 记住历史信息:模型需要依次读入 a、b、c,并把它们的信息保存在隐藏状态中
- 融合多步信息:最终需要把三步的信息综合起来做决策
- 预测未知值:输出一个训练时没见过的、全新的数字
这和"根据前三天股价预测第四天"、"根据前几个词预测下一个词"的本质是一样的。
RNN 的核心公式
标准 RNN 的更新公式只有一行:
h_t = tanh(W_xh · x_t + W_hh · h_{t-1} + b)
x_t:当前时刻的输入h_{t-1}:上一时刻的隐藏状态(记忆)h_t:当前时刻的新隐藏状态W_xh和W_hh:可学习的权重矩阵
这个公式的含义很直观:把"当前看到的东西"和"之前记住的东西"融合起来,形成新的记忆。
完整代码
import torch
import torch.nn as nn
import torch.optim as optim
class TinyRNN(nn.Module):
def __init__(self, input_size, hidden_size):
super().__init__()
self.hidden_size = hidden_size
self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
self.h2o = nn.Linear(hidden_size, 1)
def forward(self, input_seq):
h = torch.zeros(self.hidden_size)
for x in input_seq:
combined = torch.cat((x, h), dim=0)
h = torch.relu(self.i2h(combined))
return self.h2o(h)
这就是整个模型,只有两个线性层。下面逐行拆解。
模型结构拆解
两个线性层
self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
self.h2o = nn.Linear(hidden_size, 1)
- i2h(input-to-hidden):负责更新隐藏状态。输入是"当前输入拼接上一时刻隐藏状态",输出是新的隐藏状态
- h2o(hidden-to-output):负责从最终隐藏状态生成预测结果
用我们示例的具体数字来说:input_size=1,hidden_size=32
i2h:输入维度1 + 32 = 33,输出维度32h2o:输入维度32,输出维度1
总参数量 = 33 × 32 + 32 + 32 × 1 + 1 = 1121 个可训练参数。
前向传播
def forward(self, input_seq):
h = torch.zeros(self.hidden_size) # 初始记忆为空
for x in input_seq: # 逐个时间步处理
combined = torch.cat((x, h), dim=0) # 拼接输入和记忆
h = torch.relu(self.i2h(combined)) # 更新记忆
return self.h2o(h) # 用最终记忆做预测
这个循环就是 RNN 的灵魂。我们用具体数据走一遍:
假设输入 seq = [[1], [2], [3]],形状为 (3, 1)。
第 1 步(t=0):读入 a=1
x = [1] # 当前输入,shape: (1,)
h = [0, 0, ..., 0] # 初始记忆,shape: (32,),全零
combined = [1, 0, 0, ..., 0] # 拼接后 shape: (33,)
h = relu(i2h(combined)) # 新记忆 shape: (32,)
此时 h 里包含了 a=1 的信息。
第 2 步(t=1):读入 b=2
x = [2] # 当前输入
h = 上一步的结果 # 已经包含了 a 的信息
combined = [2, h_0, h_1, ..., h_31] # 拼接
h = relu(i2h(combined)) # 新记忆
此时 h 里同时包含了 a 和 b 的信息。
第 3 步(t=2):读入 c=3
x = [3] # 当前输入
h = 上一步的结果 # 已经包含了 a 和 b 的信息
combined = [3, h_0, h_1, ..., h_31] # 拼接
h = relu(i2h(combined)) # 新记忆
此时 h 里包含了 a、b、c 三步的全部信息。
输出
pred = h2o(h) # 从 32 维记忆映射到 1 维预测值
数据流图示
a=1 b=2 c=3
↓ ↓ ↓
┌────────┐ ┌────────┐ ┌────────┐
│ │ │ │ │ │
│ RNN │──→│ RNN │──→│ RNN │──→ h_final ──→ h2o ──→ pred=6
│ │ │ │ │ │
└────────┘ └────────┘ └────────┘
↑ ↑ ↑
h=0(初始) h(含a信息) h(含a+b信息)
这就是 RNN "按时间展开"的样子。同一个 RNN 单元在每个时间步复用同一组权重参数。
训练过程
def train():
model = TinyRNN(input_size=1, hidden_size=32)
optimizer = optim.Adam(model.parameters(), lr=0.005)
criterion = nn.MSELoss()
for epoch in range(300):
for _ in range(256):
# 1. 生成随机训练数据
a, b, c = torch.rand(3) # 三个 [0, 1] 的随机数
seq = torch.tensor([[a], [b], [c]])
target = torch.tensor([a + b + c]) # 正确答案:三数之和
# 2. 前向传播
pred = model(seq)
loss = criterion(pred, target)
# 3. 反向传播 + 更新参数
optimizer.zero_grad()
loss.backward()
optimizer.step()
训练的逻辑很简单:
- 随机生成三个数,算出它们的和作为正确答案
- 让模型预测,计算预测值和正确答案的差距(loss)
- 根据差距调整 1121 个权重参数
重复 300 × 256 = 76800 次后,模型就逐渐学会了"把三个数加起来"这个规律。
测试效果
输入: [1, 2, 3] -> 真实: 6.0 | 预测: 5.96 | 误差: 0.04
输入: [2, -1, 4] -> 真实: 5.0 | 预测: 5.26 | 误差: 0.26
输入: [-3, 1, 2] -> 真实: 0.0 | 预测: 0.12 | 误差: 0.12
输入: [0.5, 0.5, 0.5]-> 真实: 1.5 | 预测: 2.08 | 误差: 0.58
输入: [-1.5, 2.0, -0.5] -> 真实: 0.0 | 预测: 0.12 | 误差: 0.12
误差基本在可接受范围内。模型确实学会了求和的近似规律。
关键细节
torch.cat 的作用
combined = torch.cat((x, h), dim=0)
把当前输入(1 维)和上一时刻的隐藏状态(32 维)拼成一个 33 维的向量。这样 i2h 这个线性层就能同时看到"新信息"和"旧记忆",决定如何融合它们。
为什么用 ReLU 而不是 tanh?
传统的 RNN 用 tanh 做激活函数,但这里用了 ReLU。原因是:
- tanh 把输出限制在
(-1, 1)区间,当输入值较大时信息会被压缩丢失 - ReLU 没有上界限制,能更好地保留数值的大小信息
对于"求和"这种需要保持数值幅度的任务,ReLU 更合适。
隐藏状态每次从零开始
注意 forward 里每次处理新序列时 h 都初始化为零。这意味着模型不会在不同训练样本之间"记忆"。真正跨样本保持的是权重参数,不是隐藏状态。权重参数在训练过程中不断被更新,这才是模型"学会"东西的本质。
这个模型的局限
- 序列长度固定:当前只处理 3 个时间步的序列
- 简单 RNN 的梯度问题:如果序列很长(比如 100 步),会出现梯度消失或梯度爆炸
- 表达能力有限:只有两层线性变换,复杂任务需要更大的模型
解决这些问题的方法就是 LSTM、GRU 等更复杂的 RNN 变体,但它们的核心思想和这个 TinyRNN 完全一样——都是通过循环连接让信息在时间步之间流动。
总结
一个完整的 RNN 只需要三个东西:
- 循环单元:
h = activation(W · [x, h] + b) - 前向循环:逐个时间步处理序列
- 反向传播:用损失函数指导参数更新
理解了这三点,你就理解了 RNN 的本质。
# -*- coding: utf-8 -*-
"""
RNN 教学示例:预测下一个数字(优化版)
========================================
任务:给定序列 [a, b, c],预测 d = a + b + c
"""
import torch
import torch.nn as nn
import torch.optim as optim
class TinyRNN(nn.Module):
def __init__(self, input_size, hidden_size):
super().__init__()
self.hidden_size = hidden_size
self.i2h = nn.Linear(input_size + hidden_size, hidden_size)
self.h2o = nn.Linear(hidden_size, 1)
def forward(self, input_seq):
h = torch.zeros(self.hidden_size)
for x in input_seq:
combined = torch.cat((x, h), dim=0)
h = torch.relu(self.i2h(combined))
return self.h2o(h)
def train():
INPUT_SIZE = 1
HIDDEN_SIZE = 32
LR = 0.005
EPOCHS = 300
model = TinyRNN(INPUT_SIZE, HIDDEN_SIZE)
optimizer = optim.Adam(model.parameters(), lr=LR)
criterion = nn.MSELoss()
print("训练 RNN:给定 [a, b, c],预测 d = a+b+c")
for epoch in range(EPOCHS):
total_loss = 0
for _ in range(256):
# 用较小的随机数训练(-2 ~ 2 范围)
a, b, c = torch.rand(3)
seq = torch.tensor([[a], [b], [c]])
target = torch.tensor([a + b + c])
pred = model(seq)
loss = criterion(pred, target)
total_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (epoch + 1) % 50 == 0:
print(f"Epoch {epoch+1}/{EPOCHS}, Loss: {total_loss / 256:.4f}")
return model
def test(model):
print("\n测试结果:")
test_cases = [
[1, 2, 3],
[2, -1, 4],
[-3, 1, 2],
[0.5, 0.5, 0.5],
[-1.5, 2.0, -0.5],
]
model.eval()
with torch.no_grad():
for seq in test_cases:
seq_t = torch.tensor(seq).float().view(-1, 1)
pred = model(seq_t).item()
actual = sum(seq)
print(f"输入: {seq} -> 真实: {actual:5.1f} | 预测: {pred:6.2f} | 误差: {abs(pred - actual):.2f}")
if __name__ == "__main__":
model = train()
test(model)