如何将文档序列传入Keras Hub文本层后接入LSTM完成标签预测

错误根源 #

你遇到的形状不匹配问题核心来自两个错误设计：

• 输入层形状定义错误：你的单样本输入是长度为2的字符串序列，不是单个字符串，初始Input(shape=())完全不符合输入数据结构
• 硬编码BATCH_SIZE的reshape操作：训练/验证/测试阶段的batch大小可能不一致，最后一个batch样本数不足BATCH_SIZE时必然出现形状不匹配，且多余的reshape操作打乱了张量维度对应关系

调整后的代码实现 #

import tensorflow as tf
import tensorflow_hub as hub

BATCH_SIZE = 32 # 替换为你实际使用的batch大小
y = ... # 替换为你提前处理好的标签矩阵，y.shape[1]为分类任务类别数

hub_model = 'https://tfhub.dev/google/nnlm-en-dim50/2'
hub_layer = hub.KerasLayer(hub_model, dtype='string', trainable=False)

def build_model():
    # 单样本输入是长度为2的字符串数组，所以shape设置为(2,)
    inputs = tf.keras.Input(shape=(2,), dtype='string')
    # 对序列内每个文本单独调用Hub层生成embedding，输出形状自动为 (None, 2, 50)
    x = tf.keras.layers.TimeDistributed(hub_layer)(inputs)
    # 输出格式已符合LSTM的输入要求，无需手动reshape
    x = tf.keras.layers.LSTM(32, activation='relu')(x)
    outputs = tf.keras.layers.Dense(y.shape[1], activation='sigmoid')(x)
    return tf.keras.Model(inputs, outputs)

数据传入要求 #

• 训练输入张量的形状必须为 (总样本数, 2)，每个位置存储对应的字符串文本
• 标签张量形状为 (总样本数, 类别数)，和输入样本一一对应
• 后续适配变长序列时，只需把输入形状改为shape=(None,)，配合序列填充、mask机制即可，无需修改Hub层调用逻辑

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