TensorFlow训练SimpleRNN聊天Bot报错：logits与labels维度不匹配

问题重现 #

训练时触发以下错误：

logits and labels must have the same first dimension, got logits shape [3,21] and labels shape [15]

错误发生在model.fit执行阶段，已完成padding和分词处理，但维度仍不兼容。

错误原因 #

当前模型结构最终输出为单个输入样本对应一个单token预测结果（形状[3, 21]，3为样本数，21为词汇表大小），但标签使用的a_sequences[:, 1:]是每个输入对应一个多token序列（形状[3, 5]，展平后为[15]），两者维度完全不匹配，导致损失函数计算失败。

修正方案 #

要实现聊天机器人的序列生成，需将模型调整为多步序列预测结构，让模型每个时间步输出一个token，与标签序列的每个位置对应：

• 修改模型结构：让最后一层RNN返回序列，并用TimeDistributed包裹Dense层，使每个时间步输出对应词汇分布。
• 统一序列长度：将问句的padding长度对齐到答案的最长序列长度，保证输入输出序列维度匹配。
• 标签保持序列形状：无需展平标签，保持[样本数, 序列长度]的结构，与模型输出[样本数, 序列长度, 词汇表大小]兼容。

完整修正代码 #

import tensorflow as tf
import numpy as np
import json

# Step 1: Prepare the Data
questions = ["What is your name?", "How are you?", "What is the time?"]
answers = ["My name is ChatBot.", "I am fine, thank you.", "I do not have a watch."]

with open('q.txt', 'w') as q_file, open('a.txt', 'w') as a_file:
    for q, a in zip(questions, answers):
        q_file.write(q + '\n')
        a_file.write(a + '\n')

# Step 2: Preprocess the Data
with open('q.txt', 'r') as q_file:
    questions = q_file.readlines()
with open('a.txt', 'r') as a_file:
    answers = a_file.readlines()

tokenizer = tf.keras.preprocessing.text.Tokenizer()
tokenizer.fit_on_texts(questions + answers)

q_sequences = tokenizer.texts_to_sequences(questions)
a_sequences = tokenizer.texts_to_sequences(answers)

# 统一padding长度为最长答案序列的长度
max_answer_len = max(len(seq) for seq in a_sequences)
q_sequences = tf.keras.preprocessing.sequence.pad_sequences(q_sequences, padding='post', maxlen=max_answer_len)
a_sequences = tf.keras.preprocessing.sequence.pad_sequences(a_sequences, padding='post', maxlen=max_answer_len)

q_sequences = np.array(q_sequences)
a_sequences = np.array(a_sequences)

# Step 3: Build the Model（调整为序列输出结构）
model = tf.keras.Sequential([
    tf.keras.layers.Embedding(input_dim=len(tokenizer.word_index) + 1, output_dim=64, input_length=q_sequences.shape[1]),
    tf.keras.layers.SimpleRNN(64, return_sequences=True),
    tf.keras.layers.SimpleRNN(64, return_sequences=True),  # 修改为返回序列
    tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(len(tokenizer.word_index) + 1, activation='softmax'))  # 时间分布Dense层
])

model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])
model.summary()

# Step 4: Train the Model
model.fit(q_sequences, a_sequences[:, 1:], epochs=50)

# Step 5: Save the Weights and Biases
model.save('chatbot_model.h5')

weights_and_biases = {}
for layer in model.layers:
    if isinstance(layer, tf.keras.layers.SimpleRNN):
        weights_and_biases[layer.name] = {
            'weights': layer.get_weights()[0].tolist(),
            'biases': layer.get_weights()[1].tolist()
        }
    elif isinstance(layer, tf.keras.layers.TimeDistributed):
        dense_layer = layer.layer
        weights_and_biases[dense_layer.name] = {
            'weights': dense_layer.get_weights()[0].tolist(),
            'biases': dense_layer.get_weights()[1].tolist()
        }

with open('weights_and_biases.json', 'w') as json_file:
    json.dump(weights_and_biases, json_file)

# Step 6: Custom Feedforward Function（适配序列输出）
def custom_feedforward(input_sequence, weights_and_biases):
    def relu(x):
        return np.maximum(0, x)

    def softmax(x):
        e_x = np.exp(x - np.max(x))
        return e_x / e_x.sum(axis=-1, keepdims=True)

    # 读取两层RNN和Dense层权重
    w_rnn1 = np.array(weights_and_biases['simple_rnn']['weights'])
    b_rnn1 = np.array(weights_and_biases['simple_rnn']['biases'])
    w_rnn2 = np.array(weights_and_biases['simple_rnn_1']['weights'])
    b_rnn2 = np.array(weights_and_biases['simple_rnn_1']['biases'])
    w_dense = np.array(weights_and_biases['dense']['weights'])
    b_dense = np.array(weights_and_biases['dense']['biases'])

    # 第一层RNN计算
    hidden_state1 = np.zeros((w_rnn1.shape[1],))
    rnn1_outputs = []
    for t in range(input_sequence.shape[0]):
        hidden_state1 = relu(np.dot(input_sequence[t], w_rnn1) + b_rnn1 + np.dot(hidden_state1, w_rnn1.T))
        rnn1_outputs.append(hidden_state1)
    rnn1_outputs = np.array(rnn1_outputs)

    # 第二层RNN计算（返回序列）
    hidden_state2 = np.zeros((w_rnn2.shape[1],))
    rnn2_outputs = []
    for t in range(rnn1_outputs.shape[0]):
        hidden_state2 = relu(np.dot(rnn1_outputs[t], w_rnn2) + b_rnn2 + np.dot(hidden_state2, w_rnn2.T))
        rnn2_outputs.append(hidden_state2)
    rnn2_outputs = np.array(rnn2_outputs)

    # Dense层计算（每个时间步）
    outputs = []
    for t in range(rnn2_outputs.shape[0]):
        output = softmax(np.dot(rnn2_outputs[t], w_dense) + b_dense)
        outputs.append(output)
    outputs = np.array(outputs)
    
    return outputs

# Example usage of the custom feedforward function
input_sequence = q_sequences[0]
output = custom_feedforward(input_sequence, weights_and_biases)
predicted_tokens = np.argmax(output, axis=-1)
print("Predicted token indices:", predicted_tokens)
print("Predicted text:", tokenizer.sequences_to_texts([predicted_tokens]))

额外说明 #

如果要实现更标准的聊天机器人seq2seq结构（encoder-decoder），可将模型拆分为编码器（处理问句）和解码器（生成答案），使用LSTM或SimpleRNN构建，并配合教师强制（teacher forcing）训练，这种结构更适合多轮对话和长文本生成。

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