pytorch nn.RNN demo

之前已经讲过关于RNNCell的实现了.

这里用LLM写了一个简单的nn.RNN demo：

python 复制代码

import torch
import torch.nn as nn

# 设置随机种子以便结果可复现
torch.manual_seed(42)

# 定义模型参数
input_size = 4      # 输入特征维度
hidden_size = 8     # 隐藏层维度
num_layers = 2      # RNN 层数（修改为2层）
seq_len = 10        # 序列长度
batch_size = 3      # 批量大小

# 创建2层RNN模型
model = nn.RNN(
    input_size=input_size,
    hidden_size=hidden_size,
    num_layers=num_layers,
    batch_first=False  # 输入输出格式: [seq_len, batch_size, feature_size]
)

# 生成随机输入数据 [seq_len, batch_size, input_size]
x = torch.randn(seq_len, batch_size, input_size)
print(f"输入 x 的形状: {x.shape}  # [seq_len, batch_size, input_size]")

# 初始化隐藏状态 (可选)
h0 = torch.zeros(num_layers, batch_size, hidden_size)
print(f"初始隐藏状态 h0 的形状: {h0.shape}  # [num_layers, batch_size, hidden_size]")

# 前向传播
output, h_n = model(x, h0)
# output: 所有时间步的最后一层隐藏状态
# h_n: 所有层的最后一个时间步的隐藏状态

print(f"\n输出结果:")
print(f"output (所有时间步的最后一层隐藏状态) 的形状: {output.shape}  # [seq_len, batch_size, hidden_size]")
print(f"h_n (所有层的最后时间步隐藏状态) 的形状: {h_n.shape}  # [num_layers, batch_size, hidden_size]")

# 验证 h_n 与 output 的关系（修正后的逻辑）
print(f"\n验证 h_n 与 output 的关系:")
# 最后一层的最后状态应等于 output 的最后时间步
assert torch.allclose(h_n[-1], output[-1]), "最后一层的最后状态应等于output的最后时间步"
print(" 最后一层的最后状态与 output 的最后时间步相等")

# 打印第一层和第二层的最后隐藏状态
print(f"\n第一层的最后隐藏状态:")
print(h_n[0, 0, :5])  # 打印第一个样本的前5个元素
print(f"\n第二层的最后隐藏状态:")
print(h_n[1, 0, :5])  # 打印第一个样本的前5个元素

可以看到，nn.RNN默认会输出两个张量：一个是最后一个时间步的所有层，一个是最后一层的所有时间步。它是不会输出"所有时间步的所有层"的。

最后再给出与RNNCell部分类似的，一个完整的训练+测试的demo：

python 复制代码

import torch
import torch.nn as nn
import torch.optim as optim

# 配置
input_size = 4
hidden_size = 16
seq_len = 6
batch_size = 8
num_classes = 2
epochs = 30

# 模型定义
class RNNClassifier(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super().__init__()
        self.rnn = nn.RNN(input_size, hidden_size, batch_first=False)
        self.fc = nn.Linear(hidden_size, num_classes)

    def forward(self, x):
        # x: [seq_len, batch_size, input_size]
        output, h_n = self.rnn(x)  # h_n: [num_layers=1, batch_size, hidden_size]
        out = self.fc(h_n.squeeze(0))  # 使用最后一层的隐藏状态
        return out

# 数据生成逻辑不变
def generate_batch(batch_size, seq_len, input_size):
    x = torch.randn(seq_len, batch_size, input_size)
    last_step = x[-1]
    labels = (last_step[:, 0] > 0).long()
    return x, labels

# 初始化模型与训练配置
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = RNNClassifier(input_size, hidden_size, num_classes).to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# 训练
for epoch in range(epochs):
    model.train()
    x_batch, y_batch = generate_batch(batch_size, seq_len, input_size)
    x_batch, y_batch = x_batch.to(device), y_batch.to(device)

    logits = model(x_batch)
    loss = criterion(logits, y_batch)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 5 == 0 or epoch == 0:
        pred = logits.argmax(dim=1)
        acc = (pred == y_batch).float().mean().item()
        print(f"[Epoch {epoch+1}] Loss: {loss.item():.4f}, Acc: {acc:.2f}")

# 测试
model.eval()
with torch.no_grad():
    x_test, y_test = generate_batch(1, seq_len, input_size)
    x_test, y_test = x_test.to(device), y_test.to(device)
    pred = model(x_test).argmax(dim=1)
    print("\nTest sample:")
    print("Target label:", y_test.item())
    print("Predicted   :", pred.item())