怎样用循环神经网络（RNN）实现下一个单词的预测

访客自然语言处理 2026-06-05 05:45:52 1

本文目录导读：

核心原理简述
数据准备
模型搭建
训练模型
进行预测
改进与进阶技巧
完整代码汇总

使用循环神经网络（RNN）实现“下一个单词预测”（Next Word Prediction）是一个非常经典的自然语言处理任务，下面我将从数据准备、模型搭建、训练、预测四个核心步骤，结合代码示例（使用 PyTorch）为你详细讲解。

核心原理简述

输入：一个单词序列（["我", "喜欢", "吃"]）
输出：下一个最可能的单词（"苹果"）
RNN 的作用：通过隐藏状态（hidden state）记住前面的单词信息，从而利用上下文来预测下一个词。

数据准备

1 构建词汇表

把每个单词映射为一个唯一的整数 ID。

from collections import Counter
text = "我 喜欢 吃 苹果 我 喜欢 喝 牛奶 他 喜欢 吃 香蕉"
words = text.split()
vocab = list(set(words))
word2idx = {w: i for i, w in enumerate(vocab)}  # 词 → ID
idx2word = {i: w for i, w in enumerate(vocab)}  # ID → 词
vocab_size = len(vocab)
print("词汇表大小:", vocab_size)

2 生成训练样本

使用滑动窗口构造 (输入序列, 目标单词) 对。
窗口大小为 3 → 前 2 个词预测第 3 个词。

import torch
seq_len = 2  # 用前2个词预测第3个
data = []
for i in range(len(words) - seq_len):
    input_seq = words[i:i+seq_len]
    target_word = words[i+seq_len]
    input_ids = [word2idx[w] for w in input_seq]
    target_id = word2idx[target_word]
    data.append((input_ids, target_id))
print("样本示例:", data[:3])
# 输出: ([0, 1], 2) 之类的 (实际ID取决于词汇表顺序)

模型搭建

使用 PyTorch 的 nn.RNN 或更常用的 nn.LSTM（可以缓解长序列梯度消失问题）。

import torch.nn as nn
import torch.nn.functional as F
class NextWordRNN(nn.Module):
    def __init__(self, vocab_size, embedding_dim, hidden_dim):
        super(NextWordRNN, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.rnn = nn.RNN(embedding_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)  # 输出层：词概率分布
    def forward(self, x, hidden=None):
        # x: (batch_size, seq_len)
        embeds = self.embedding(x)           # (batch, seq_len, embed_dim)
        out, hidden = self.rnn(embeds, hidden) # out: (batch, seq_len, hidden_dim)
        # 取最后一个时间步的输出 → 预测下一个词
        out = out[:, -1, :]                  # (batch, hidden_dim)
        logits = self.fc(out)                # (batch, vocab_size)
        return logits, hidden

参数说明：

embedding_dim：词向量维度（如 100）
hidden_dim：RNN 隐层维度（如 128）
vocab_size：词汇表大小

训练模型

# 超参数
embedding_dim = 50
hidden_dim = 100
learning_rate = 0.01
epochs = 100
model = NextWordRNN(vocab_size, embedding_dim, hidden_dim)
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
# 将数据转为张量
inputs = torch.tensor([d[0] for d in data])   # (样本数, seq_len)
targets = torch.tensor([d[1] for d in data])  # (样本数,)
# 训练循环
for epoch in range(epochs):
    model.train()
    optimizer.zero_grad()
    logits, _ = model(inputs)          # 前向传播
    loss = criterion(logits, targets)  # 计算损失
    loss.backward()                    # 反向传播
    optimizer.step()                   # 更新参数
    if (epoch+1) % 20 == 0:
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')

进行预测

给定一个开头序列,预测下一个最可能的单词。

def predict_next_word(model, input_words, word2idx, idx2word):
    model.eval()
    # 将输入词转为 ID
    input_ids = [word2idx[w] for w in input_words]
    input_tensor = torch.tensor([input_ids])  # (1, seq_len)
    with torch.no_grad():
        logits, _ = model(input_tensor)
        probs = F.softmax(logits, dim=1)
        # 取概率最大的词
        predicted_idx = torch.argmax(probs, dim=1).item()
    return idx2word[predicted_idx]
# 测试
test_input = ["我", "喜欢"]
predicted = predict_next_word(model, test_input, word2idx, idx2word)
print(f"{' '.join(test_input)} → {predicted}")

输出示例（取决于数据和训练程度）：

我 喜欢 → 吃

改进与进阶技巧

问题	解决方法
序列太长导致梯度消失	使用 LSTM 或 GRU 代替普通 RNN
生成结果不够多样	用温度采样（Temperature Sampling）代替 argmax
模型容量不足	增加 embedding_dim / hidden_dim / 层数
训练数据太少	使用预训练词向量（如 Word2Vec / GloVe）
需要生成多个词	用自回归方式：每次预测一个词，将其加入输入序列继续预测

温度采样代码片段：

def sample_with_temperature(logits, temperature=0.8):
    probs = F.softmax(logits / temperature, dim=1)
    return torch.multinomial(probs, num_samples=1).item()