从零手写一个 RNN:80 行代码搞懂循环神经网络

2 阅读7分钟

很多 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 的三个核心能力:

  1. 记住历史信息:模型需要依次读入 a、b、c,并把它们的信息保存在隐藏状态中
  2. 融合多步信息:最终需要把三步的信息综合起来做决策
  3. 预测未知值:输出一个训练时没见过的、全新的数字

这和"根据前三天股价预测第四天"、"根据前几个词预测下一个词"的本质是一样的。

RNN 的核心公式

标准 RNN 的更新公式只有一行:

h_t = tanh(W_xh · x_t + W_hh · h_{t-1} + b)
  • x_t:当前时刻的输入
  • h_{t-1}:上一时刻的隐藏状态(记忆)
  • h_t:当前时刻的新隐藏状态
  • W_xhW_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=1hidden_size=32

  • i2h:输入维度 1 + 32 = 33,输出维度 32
  • h2o:输入维度 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()

训练的逻辑很简单:

  1. 随机生成三个数,算出它们的和作为正确答案
  2. 让模型预测,计算预测值和正确答案的差距(loss)
  3. 根据差距调整 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 都初始化为零。这意味着模型不会在不同训练样本之间"记忆"。真正跨样本保持的是权重参数,不是隐藏状态。权重参数在训练过程中不断被更新,这才是模型"学会"东西的本质。

这个模型的局限

  1. 序列长度固定:当前只处理 3 个时间步的序列
  2. 简单 RNN 的梯度问题:如果序列很长(比如 100 步),会出现梯度消失或梯度爆炸
  3. 表达能力有限:只有两层线性变换,复杂任务需要更大的模型

解决这些问题的方法就是 LSTM、GRU 等更复杂的 RNN 变体,但它们的核心思想和这个 TinyRNN 完全一样——都是通过循环连接让信息在时间步之间流动。

总结

一个完整的 RNN 只需要三个东西:

  1. 循环单元h = activation(W · [x, h] + b)
  2. 前向循环:逐个时间步处理序列
  3. 反向传播:用损失函数指导参数更新

理解了这三点,你就理解了 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)