3.1.6 RNN 的 API 使用 — 详细讲解

一、概述

PyTorch 提供了 torch.nn.RNN 模块用于构建循环神经网络(Recurrent Neural Network)。该模块是 PyTorch 对 RNN 的标准实现,功能强大,具备以下核心能力:

  • 支持单层或多层堆叠(通过 num_layers 控制)
  • 支持双向 RNN(通过 bidirectional 控制),可同时捕捉前向和后向的序列依赖
  • 适用于各类序列建模任务:文本分类、序列标注、命名实体识别、时间序列预测等

该模块是传统 RNN 的基础组件,理解其 API 是后续学习 LSTM、GRU 的前提。


二、构造函数参数详解

torch.nn.RNN(
    input_size,          # 必填:输入特征维度
    hidden_size,         # 必填:隐藏状态维度
    num_layers=1,        # RNN 层数
    nonlinearity="tanh", # 激活函数
    bias=True,           # 是否使用偏置
    batch_first=False,   # 批次维度是否在第一维
    dropout=0.0,         # 层间 dropout 概率
    bidirectional=False, # 是否双向
    device=None,         # 设备
    dtype=None,          # 数据类型
)

各参数详细说明

参数类型默认值详细说明
input_sizeint必填每个时间步输入的特征维度。在 NLP 任务中,通常是词向量的维度。例如:若使用 300 维的 Word2Vec 词向量,则 input_size=300;若使用 PyTorch 的 nn.Embedding(vocab_size, 128),则 input_size=128
hidden_sizeint必填隐藏状态的维度。这个值决定了 RNN 的"记忆容量"——越大模型表达能力越强,但参数量和过拟合风险也随之增加。合理的值通常在 128~1024 之间,需根据任务复杂度调整
num_layersint1RNN 的堆叠层数。当 num_layers ≥ 2 时,第 i 层的输出会作为第 i+1 层的输入,形成深层 RNN。深层 RNN 能学习更抽象的时序特征,但也会加重梯度消失问题
nonlinearitystr"tanh"激活函数,可选 "tanh" 或 "relu"。tanh 是 RNN 的传统选择(输出范围 [-1,1]),relu 可在一定程度上缓解梯度消失,但可能导致神经元"死亡"
biasboolTrue是否在线性变换中使用偏置项 bbb。计算隐藏状态时公式为 ht=tanh⁡(Wihxt+bih+Whhht−1+bhh)h_t = \tanh(W_{ih}x_t + b_{ih} + W_{hh}h_{t-1} + b_{hh})ht​=tanh(Wih​xt​+bih​+Whh​ht−1​+bhh​),关闭后则无偏置项
batch_firstboolFalse控制输入和输出张量的维度排列顺序。默认 (seq_len, batch, feature);设为 True 则变为 (batch, seq_len, feature)。强烈建议设为 True,这样与 DataLoader 的输出格式一致,代码更直观
dropoutfloat0.0Dropout 概率,用于防止过拟合。仅当 num_layers ≥ 2 时才生效,且只作用于除最后一层之外的所有层之间的连接。原理是在训练时随机将一部分神经元的输出置零
bidirectionalboolFalse是否为双向 RNN。设为 True 后,模型会同时学习正向序列和反向序列两个方向的信息,num_directions 从 1 变为 2,隐藏状态维度会翻倍
devicestr 或 torch.deviceNone模块运行的设备。如 "cuda"、"cuda:0"、"cpu"。若不指定,模块将在当前默认设备上创建参数
dtypetorch.dtypeNone模块参数的数据类型。如 torch.float32(默认)、torch.float64。混合精度训练时可设为 torch.float16

三、输入与输出详解

3.1 调用方式

rnn = torch.nn.RNN(input_size, hidden_size, ...)
output, h_n = rnn(input, h_0)   # h_0 是可选的初始隐藏状态

3.2 输入参数

参数形状 batch_first=False形状 batch_first=True说明
input(seq_len, batch_size, input_size)(batch_size, seq_len, input_size)输入序列。seq_len 是序列长度(如句子中的词数),batch_size 是批次大小,input_size 是每个时间步的特征维度
h_0(可选)(num_layers × num_directions, batch_size, hidden_size)同左初始隐藏状态。若不传入,PyTorch 自动初始化为全零张量。当需要状态传递(如生成式解码)时显式传入

注意:h_0 的第一维 = num_layers × num_directions。单向时 num_directions = 1,双向时 num_directions = 2。

3.3 输出参数

参数形状 batch_first=False形状 batch_first=True说明
output(seq_len, batch_size, num_directions × hidden_size)(batch_size, seq_len, num_directions × hidden_size)每个时间步的隐藏状态,仅包含最后一层的所有时间步输出。如果 bidirectional=True,每个时间步的最后一维是正向和反向隐藏状态的拼接
h_n(num_layers × num_directions, batch_size, hidden_size)同左最后一个时间步的隐藏状态。包含所有层、所有方向的最终状态。可用于初始化下一个 batch 的 h_0,实现跨 batch 状态传递

3.4 output 与 h_n 的关系

  • 单层单向:output[:, -1, :] == h_n[0]——output 的最后一个时间步等于 h_n 的最后一层
  • 多层单向:output[:, -1, :] == h_n[-1]——output 的最后一个时间步等于 h_n 的最后一层
  • 单层双向:output 每个时间步 = (正向隐状态, 反向隐状态);
    h_n[0] = 正向最后一个时间步,h_n[1] = 反向第一个时间步(即反向视角的最后一个)
  • 多层双向:h_n 按层和方向交错排列,顺序为 (L1正, L1反, L2正, L2反, ...)

四、四种常见结构配置

4.1 单层单向 — 最基础

import torch.nn as nn

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=1, batch_first=True)

# 输入输出形状:
# input  : (batch, seq_len, 128)
# output : (batch, seq_len, 256)
# h_n    : (1, batch, 256)

特点:结构最简单,参数量最少,训练最快。

适用场景:短文本分类、简单序列标注等不需要深层语义理解的任务。


4.2 多层单向 — 深度堆叠

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=3, batch_first=True)

# 输入输出形状:
# input  : (batch, seq_len, 128)
# output : (batch, seq_len, 256)
# h_n    : (3, batch, 256)

特点:通过堆叠多层 RNN,每一层学习不同粒度的时序特征——低层学局部模式,高层学全局依赖。

注意事项:

  • 层数过多会显著增加训练时间
  • 深度增加会加重梯度消失问题,通常 LSTM/GRU 更适合深层堆叠
  • 建议配合 dropout 使用(如 dropout=0.3),对层间连接进行正则化

适用场景:较长序列的建模任务,需要多层次时序特征提取。


4.3 单层双向 — 双向上下文

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=1,
             bidirectional=True, batch_first=True)

# 输入输出形状:
# input  : (batch, seq_len, 128)
# output : (batch, seq_len, 512)   ← 256 × 2(正向+反向拼接)
# h_n    : (2, batch, 256)         ← 1层 × 2方向

特点:模型同时从正向(从左到右)和反向(从右到左)处理序列,每个时间步都能感知完整上下文。

工作原理:

  • 正向 RNN 按 x1→x2→⋯→xTx_1 \to x_2 \to \dots \to x_Tx1​→x2​→⋯→xT​ 处理
  • 反向 RNN 按 xT→xT−1→⋯→x1x_T \to x_{T-1} \to \dots \to x_1xT​→xT−1​→⋯→x1​ 处理
  • 每个时间步 ttt 的输出 = [正向隐状态, 反向隐状态](拼接后维度翻倍)

适用场景:

  • 命名实体识别(NER):判断一个词是否是人名/地名时,需要看上下文
  • 文本分类:需要整句的全面语义
  • 机器翻译的编码器

4.4 多层双向 — 最强配置

rnn = nn.RNN(input_size=128, hidden_size=256, num_layers=3,
             bidirectional=True, batch_first=True)

# 输入输出形状:
# input  : (batch, seq_len, 128)
# output : (batch, seq_len, 512)   ← 256 × 2
# h_n    : (6, batch, 256)         ← 3层 × 2方向

特点:多层 + 双向,代表 RNN 家族中最强的表示能力。低层的双向 RNN 捕获局部上下文,高层的双向 RNN 在此基础上学习更全局的语义。

注意事项:

  • 参数量大,需要在较大的数据集上训练
  • 通常需要配合 Dropout 和 Early Stopping 防止过拟合
  • 在实际工程中,多层双向 LSTM/GRU 比多层双向 RNN 更常用(因为 RNN 的梯度消失问题)

五、完整代码示例

5.1 单层单向 RNN — 基础用法

import torch
import torch.nn as nn

# ===== 超参数 =====
batch_size = 16
seq_len = 20          # 序列长度(如一个句子有 20 个词)
input_size = 128      # 词向量维度
hidden_size = 256     # 隐藏层维度
num_layers = 2        # 堆叠 2 层

# ===== 创建 RNN 层 =====
rnn = nn.RNN(
    input_size=input_size,
    hidden_size=hidden_size,
    num_layers=num_layers,
    nonlinearity="tanh",
    bias=True,
    batch_first=True,        # 推荐设为 True
    dropout=0.3,             # 层间 Dropout(仅 2 层及以上生效)
    bidirectional=False,
)

# ===== 构造输入数据 =====
x = torch.randn(batch_size, seq_len, input_size)   # (16, 20, 128)

# ===== 可选:初始隐藏状态 =====
h0 = torch.zeros(num_layers, batch_size, hidden_size)  # (2, 16, 256)

# ===== 前向传播 =====
output, hn = rnn(x, h0)

print(f"output shape: {output.shape}")   # torch.Size([16, 20, 256])
print(f"hn shape:     {hn.shape}")        # torch.Size([2, 16, 256])

# ===== 取最后一个时间步 =====
# 常用于文本分类任务中,取 "句子级" 表示
last_output = output[:, -1, :]            # (16, 256)
# ↓ 可接入全连接层进行分类
# logits = nn.Linear(256, num_classes)(last_output)

5.2 双向 RNN — 完整上下文感知

# ===== 双向 RNN 示例 =====
bi_rnn = nn.RNN(
    input_size=128,
    hidden_size=256,
    num_layers=2,
    batch_first=True,
    bidirectional=True,      # 开启双向
    dropout=0.3,
)

x = torch.randn(16, 20, 128)   # (batch, seq_len, input_size)

output, hn = bi_rnn(x)

print(f"output shape: {output.shape}")   # torch.Size([16, 20, 512]) — 256 × 2
print(f"hn shape:     {hn.shape}")        # torch.Size([4, 16, 256])  — 2层 × 2方向

# ===== 提取双向信息 =====
# output[:, :, :256]   → 正向每个时间步的隐藏状态
# output[:, :, 256:]   → 反向每个时间步的隐藏状态
#
# 注意:反向的第一个位置对应原始序列的最后一个词!
# output[:, -1, :256]  → 正向最后一个时间步(序列末尾)
# output[:, 0, 256:]   → 反向最后一个时间步(序列开头)

# ===== 文本分类场景:拼接正反向最后一个状态 =====
forward_last = output[:, -1, :256]    # (16, 256) 正向末端
backward_last = output[:, 0, 256:]    # (16, 256) 反向末端(= 反向视角的最后)
combined = torch.cat([forward_last, backward_last], dim=1)  # (16, 512)
# ↓ 接入分类器
# logits = nn.Linear(512, num_classes)(combined)

5.3 不传 h_0 的简化写法

# 大部分场景下,不需要显式传入 h_0,PyTorch 会自动初始化为全零
rnn = nn.RNN(128, 256, num_layers=2, batch_first=True)
x = torch.randn(16, 20, 128)

output, hn = rnn(x)          # 自动使用全零的初始隐藏状态
# output: (16, 20, 256)
# hn:     (2, 16, 256)

六、关键注意事项

6.1 batch_first 建议设为 True

PyTorch RNN 模块的默认输入格式为 (seq_len, batch, feature),这与常见的 DataLoader 输出格式 (batch, seq_len, feature) 不一致,容易造成混淆。建议统一设置 batch_first=True,让数据流转更加直观。

6.2 dropout 仅在多层时生效

# num_layers=1 时,dropout 参数被忽略
rnn = nn.RNN(128, 256, num_layers=1, dropout=0.5)  # ← 这个 dropout 不会生效!

# num_layers≥2 时,dropout 作用于层间连接(除最后一层外)
rnn = nn.RNN(128, 256, num_layers=3, dropout=0.3)  # ← 在第 1→2、第 2→3 之间施加 dropout

6.3 输出 output 只包含最后一层

如果你需要访问所有中间层的输出,output 满足不了这个需求。但通常这不是问题,因为:

  • 分类任务只需最后一层最后一个时间步
  • 序列标注需用最后一层所有时间步
  • 如果需要中间层输出做特征融合,需自行改写模块

6.4 双向 RNN 时输出维度翻倍

# 单向:output 最后一维 = hidden_size
rnn = nn.RNN(128, 256, bidirectional=False)
# output: (..., 256)

# 双向:output 最后一维 = hidden_size × 2
rnn = nn.RNN(128, 256, bidirectional=True)
# output: (..., 512)

务必注意:当 RNN 后接全连接层时,全连接层的 in_features 需要匹配翻倍后的维度!

6.5 h_0 的初始化

大多数任务中可以用全零初始化,省去传 h_0 的麻烦。但在以下场景中建议显式初始化:

  • 状态传递(stateful RNN):将上一个 batch 的 h_n.detach() 作为下一个 batch 的 h_0
  • 特殊初始化策略:如使用 Xavier/Kaiming 初始化

6.6 梯度消失与梯度爆炸

这是 RNN 最核心的缺陷。虽然本节讲的是 API 使用,但必须了解:传统 RNN(torch.nn.RNN)在长序列上容易出现梯度消失或梯度爆炸。这是后续引入 LSTM 和 GRU 的根本原因。在实际项目中,长序列任务推荐优先使用 LSTM 或 GRU。


七、参数维度速查表

配置输入形状 (batch_first=True)输出 output 形状h_n 形状
单层单向(B, S, I)(B, S, H)(1, B, H)
多层单向(B, S, I)(B, S, H)(L, B, H)
单层双向(B, S, I)(B, S, 2H)(2, B, H)
多层双向(B, S, I)(B, S, 2H)(2L, B, H)

B = batch_size, S = seq_len, I = input_size, H = hidden_size, L = num_layers


八、与后续内容的关系

本节讲完 RNN 的 API 后,下一节(3.1.7 存在问题)将深入剖析 RNN 的核心缺陷——长期依赖建模困难,其根源在于:

  • 梯度消失:反向传播中,早期时间步的梯度经过多次 tanh 导数的连乘(每项 ≤1),指数级衰减至 0,导致模型只能学到短期依赖
  • 梯度爆炸:若权重过大,梯度指数级增长,参数更新极不稳定

这两个问题直接引出了 LSTM(第 3.2 节)和 GRU(第 3.3 节),它们通过精巧的门控机制有效缓解了上述缺陷,成为实际应用中的主流选择。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐