TensorFlow带Embedding层Seq2seq模型预测报错及掩码失效问题咨询

• 你遇到的LSTM报错本质是：将仅适配训练流程的端到端耦合计算图直接用于自回归逐词预测，没有拆分独立的推理用编码、解码子模型，导致调用模型时输入张量和计算图预期的输入数量、维度不匹配。
  你当前写的model是纯Teacher Forcing模式的训练模型：训练阶段Decoder的初始状态是和Encoder的输出硬绑定的，整个计算图是为了“一次性喂入完整的源序列、完整的带移位的目标序列，直接计算全序列损失”设计的，根本不支持逐词自回归的推理流程，直接调用predict自然会出现输入对不上的问题。哪怕你调整单条样本的维度适配模型输入格式，直接调用训练模型也只能输出全序列的一次性结果，无法实现逐词自回归生成。
• 你感知到的mask_zero未生效的问题不存在本质逻辑错误：该参数的作用是生成“跳过值为0的填充位”的掩码，顺着计算图自动传递到下游连接的LSTM层，在序列计算、损失计算时忽略0值填充位置。你看到的“未生效”的观感，是因为错误的推理链路中断了计算图的正常掩码传递，不是Embedding层的配置出了问题。只要计算图连通、用的是Keras内置层，掩码就会正常工作。

Seq2seq的标准结构必须拆成3个独立的计算图才能兼顾训练效率和推理正确性：

1. 训练用端到端模型 #

逻辑和你现有代码基本一致，仅需调整命名区分用途，不需要改动已有的训练流程：

# Encoder部分定义
encoder_inputs = tf.keras.Input(shape=(None,))
enc_emb_layer = tf.keras.layers.Embedding(
    input_dim=max_eng_vocabulary,
    output_dim=embedding_output_length,
    mask_zero=True
)(encoder_inputs)
encoder_lstm_layer = tf.keras.layers.LSTM(50, return_state=True)
encoder_outputs, state_h, state_c = encoder_lstm_layer(enc_emb_layer)
encoder_states = [state_h, state_c]

# Decoder部分定义（训练用，初始状态硬绑定Encoder输出）
decoder_inputs = tf.keras.Input(shape=(None,))
dec_emb_layer = tf.keras.layers.Embedding(
    input_dim=max_chi_vocabulary,
    output_dim=embedding_output_length,
    mask_zero=True
)(decoder_inputs)
decoder_lstm_layer = tf.keras.layers.LSTM(
    50,
    return_state=True,
    return_sequences=True
)
decoder_outputs, _, _ = decoder_lstm_layer(dec_emb_layer, initial_state=encoder_states)
decoder_dense_layer = tf.keras.layers.Dense(max_chi_vocabulary, activation='softmax')
decoder_outputs = decoder_dense_layer(decoder_outputs)

# 可直接用于fit的训练模型
train_model = tf.keras.Model([encoder_inputs, decoder_inputs], decoder_outputs)

2. 推理用独立Encoder模型 #

用来把输入的源序列编码成初始解码状态，输出为LSTM的两个记忆状态张量：

inference_encoder = tf.keras.Model(encoder_inputs, encoder_states)

3. 推理用独立Decoder模型 #

需要把之前硬绑定的状态输入改成显式的入参，复用训练阶段已经初始化好权重的层，避免权重重新随机初始化：

# 显式定义解码状态的输入张量
decoder_state_input_h = tf.keras.Input(shape=(50,))
decoder_state_input_c = tf.keras.Input(shape=(50,))
decoder_states_inputs = [decoder_state_input_h, decoder_state_input_c]

# 复用训练阶段的已训练层权重，不要新建层
decoder_outputs, out_h, out_c = decoder_lstm_layer(
    dec_emb_layer(decoder_inputs),
    initial_state=decoder_states_inputs
)
decoder_out_states = [out_h, out_c]
decoder_outputs = decoder_dense_layer(decoder_outputs)

# 推理用Decoder的输入：当前步的词token + 上一步传递的LSTM状态
# 推理用Decoder的输出：当前步的词概率 + 要向下一步传递的更新后LSTM状态
inference_decoder = tf.keras.Model(
    [decoder_inputs] + decoder_states_inputs,
    [decoder_outputs] + decoder_out_states
)

不要直接调用训练用模型做预测，按照自回归逻辑逐词生成：

1. 将待预测的单条源序列调整为带batch维度的形状后喂给inference_encoder，拿到初始的state_h和state_c
2. 给Decoder传入初始的起始符对应token，调整形状为(1, 1)（批量大小为1，当前时间步长度为1）
3. 循环生成结果：每一步将当前token、上一步输出的两个状态喂给inference_decoder，拿到当前步的预测概率、更新后的状态；将概率对应的最高概率词作为下一步的输入，直到预测出结束符，或者达到预设的最大序列长度

关于mask_zero的额外说明：
该参数在你当前的配置下是正常生效的，如果后续你在模型中加入自定义计算层，必须在自定义层中实现compute_mask方法，否则掩码在层间传递时会中断，才会出现掩码失效的问题。Keras内置的Embedding、LSTM、Dense层都自带掩码传递逻辑，不需要额外配置。

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