基于Keras的下一词预测：如何获取每个输入词的预测结果

我来帮你一步步理清这两个问题，解决你的困惑～

问题一：训练数据的划分方式选择 #

首先咱们得明确两种划分方式的适用场景：

你看到的常规划分方式 #

X, y = sequences[:,:-1], sequences[:,-1]
y = to_categorical(y, num_classes=vocab_size)

这种方式是把每个固定长度的序列拆成「前n-1个词作为输入」，「最后1个词作为预测目标」，属于单步下一词预测的典型做法。比如你的代码里input_length=55，那每个X样本就是55个词，y是这55个词之后的第1个词，适合那种给定一段固定长度上下文，预测单个下一个词的场景。

你考虑的滑动序列划分方式 #

X = [[10, 9, 4, 5]]
X_train = [[10, 9], [9, 4], [4, 5]]
y_train = [[9], [4], [5]]

这是滑动窗口式的多样本生成，相当于把长文本切成多个重叠的短上下文-目标对，能让模型接触到更多局部上下文的关联，训练样本量也会更大。

怎么选？

两种方式都可以用于下一词预测，核心是要保证你的输入长度和模型Embedding层的input_length匹配。比如你代码里设置了input_length=55，那不管用哪种方式，每个X_train样本都得是55个词的序列：

• 如果用滑动窗口，就把原始长文本切成多个长度为56的片段（前55个是输入，第56个是目标），然后批量生成X和y；
• 如果是已经整理好的固定长度序列，用第一种常规划分更直接。

问题二：获取每个词的完整概率列表 & TimeDistributed层的维度问题 #

你当前的模型只能输出单个预测结果，是因为LSTM默认只返回最后一个时间步的输出，要实现每个时间步都输出对应下一词的概率，得调整模型结构和训练数据：

1. 调整模型结构（解决TimeDistributed的维度问题） #

关键是让LSTM返回每个时间步的输出，再用TimeDistributed层对每个时间步的输出做分类。修改后的代码如下：

from keras.layers import TimeDistributed

model = Sequential()
# Embedding层输出是3D张量：(batch_size, input_length, embedding_size)，正好符合TimeDistributed的输入要求
model.add(Embedding(vocab_size, embedding_size, input_length=55, weights=[pretrained_weights]))
# 设置return_sequences=True，让LSTM返回每个时间步的输出（3D张量）
model.add(Bidirectional(LSTM(units=embedding_size, return_sequences=True)))
# 用TimeDistributed包裹Dense层，对每个时间步的输出单独做softmax分类
model.add(TimeDistributed(Dense(vocab_size, activation='softmax')))

model.compile(loss='categorical_crossentropy', optimizer=RMSprop(lr=0.0005), metrics=['accuracy'])

2. 调整训练数据格式 #

现在模型的输出是3D张量：(batch_size, input_length, vocab_size)，对应每个输入序列的每个时间步的下一词概率，所以你的y_train也得改成3D格式：

• 假设你的原始sequences是长度为56的序列（前55个是输入，后55个是每个时间步对应的目标词）；
• 划分方式改成：

X = sequences[:, :-1]  # 取每个序列的前55个词作为输入
y = sequences[:, 1:]   # 取每个序列的后55个词作为目标（每个位置对应输入序列的下一个词）
y = to_categorical(y, num_classes=vocab_size)  # 转换成3D的one-hot编码：(batch_size, 55, vocab_size)

3. 获取完整概率列表 #

训练完成后，用model.predict()得到的就是3D数组，比如对于一个输入样本，你可以这样获取每个时间步的概率：

predictions = model.predict(np.array(X_test), batch_size=32)
# predictions的形状是：(样本数, 55, 23690)
# 取第一个样本的第一个时间步的概率列表
first_step_probs = predictions[0][0]
# 取第一个样本的所有时间步的概率列表
all_steps_probs = predictions[0]

每个概率列表的长度就是你的词汇表大小（23690），对应每个词的预测概率。

内容的提问来源于stack exchange，提问作者user9355680