Kaggle中Keras双向LSTM训练触发输入维度不匹配报错

报错信息 #

ValueError: Input 0 of layer bidirectional_1 is incompatible with the layer: expected ndim=3, found ndim=2. Full shape received: (13, 64)

<ipython-input-25-f223adf5c5a7> in <module>
     12   for i in range(len(scores_transformed_train)):
     13     # Optimize the model
---> 14     loss_value, grads = grad(temp_model, title_train[i], 
scores_transformed_train[i])
     15     optimizer.apply_gradients(zip(grads, temp_model.trainable_variables))
     16 

核心问题定位 #

• 报错本质是输入张量维度缺失：Bidirectional包装的LSTM层强制要求输入为3维结构，格式为(批次大小, 序列长度, 特征维度)。
• 当前代码逐样本训练时，直接将单条样本title_train[i]传入模型，没有保留批次维度。这条单样本经过TextVectorization、Embedding层处理后，输出形状为(13, 64)的2维张量，完全达不到LSTM的3维输入要求。
• 现有模型的层配置本身无逻辑错误：默认配置的Bidirectional(LSTM(64))会输出2维的最终序列隐藏状态，和后续Dense层的输入要求是匹配的，问题完全出在输入侧的维度缺失。

原有模型代码参考 #

encoder = tf.keras.layers.experimental.preprocessing.TextVectorization()
encoder.adapt(title_train)

model = tf.keras.Sequential([
    encoder,
    tf.keras.layers.Embedding(
        input_dim=len(encoder.get_vocabulary()),
        output_dim=64,
        mask_zero=True),
    tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(64)),
    tf.keras.layers.Dense(64, activation='relu'),
    tf.keras.layers.Dense(1, 'sigmoid')
])

修复方案 #

方案1：逐样本训练时手动补全批次维度 #

在自定义训练循环里，给传入的单条输入、标签增加长度为1的批次轴，保证经过层处理后的张量维度符合LSTM要求：

for i in range(len(scores_transformed_train)):
    # 为单样本增加批次维度，将原形状(序列长度,)转为(1, 序列长度)
    single_title = tf.expand_dims(title_train[i], axis=0)
    single_score = tf.expand_dims(scores_transformed_train[i], axis=0)
    loss_value, grads = grad(temp_model, single_title, single_score)
    optimizer.apply_gradients(zip(grads, temp_model.trainable_variables))

补全维度后，Embedding层输出形状会变为(1, 13, 64)，完全匹配LSTM的3维输入标准。

方案2：改为批量训练（推荐） #

放弃逐样本迭代的写法，用tf.data构建带批次的数据集，既避免手动处理维度的问题，也能大幅提升训练效率：

# 构建带shuffle、批次划分的训练数据集
train_dataset = tf.data.Dataset.from_tensor_slices(
    (title_train, scores_transformed_train)
).shuffle(buffer_size=len(title_train)).batch(32)

# 编译后直接调用内置fit训练，无需手动写梯度计算逻辑
model.compile(optimizer=optimizer, loss=loss_object)
model.fit(train_dataset, epochs=10)

避坑提示 #

• 不要随意给单层LSTM添加return_sequences=True参数：该参数会让LSTM输出3维张量，反而会和后续Dense层的2维输入要求冲突，只有堆叠多层LSTM时才需要开启该参数给下一层传递3维输入。
• 自定义训练循环中手动传入单样本时，必须手动补全批次维度，TensorFlow不会自动为这种手动调用补维度。

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