Keras 3中Attention层无法接收LSTM层输出的问题求助

问题原因 #

这个错误通常是因为Keras 3中tf.keras.layers.Attention的实现与KerasTensor兼容性不佳，或是mask传递不完整导致内部触发了TensorFlow原生函数调用，而KerasTensor无法直接作为其输入。

修复方案 #

1. 切换到Keras 3原生Attention层 #

替换tf.keras.layers.Attention为Keras 3原生的keras.layers.Attention，原生层对KerasTensor的支持更完善：

# 替换原有Attention层定义
attention_layer = keras.layers.Attention()

2. 显式传递mask参数 #

由于Embedding层设置了mask_zero=True，encoder和decoder的输出张量带有mask信息，但Attention层可能无法自动识别，需显式传入：

attention_outputs = attention_layer(
    [decoder_outputs, encoder_outputs],
    mask=[decoder_outputs._keras_mask, encoder_outputs._keras_mask]
)

3. 规范TextVectorization层使用 #

确保text_vec_layer_en和text_vec_layer_es是Keras 3的keras.layers.TextVectorization层，避免返回非KerasTensor类型的输出：

# 示例创建方式（若尚未创建）
text_vec_layer_en = keras.layers.TextVectorization(
    max_tokens=vocab_size_en,
    output_mode="int",
    output_sequence_length=max_seq_len_en
)
text_vec_layer_es = keras.layers.TextVectorization(
    max_tokens=vocab_size_es,
    output_mode="int",
    output_sequence_length=max_seq_len_es
)

完整修复后代码片段 #

import keras  # 确保导入Keras 3核心库

encoder_inputs = keras.layers.Input(shape=[], dtype="string")
decoder_inputs = keras.layers.Input(shape=[], dtype="string")

embed_size = 128
encoder_inputs_ids = text_vec_layer_en(encoder_inputs)
decoder_inputs_ids = text_vec_layer_es(decoder_inputs)

# 英西翻译需使用各自独立的词汇表，修正为对应词汇表大小
encoder_embedding_layer = keras.layers.Embedding(vocab_size_en, embed_size, mask_zero=True)
decoder_embedding_layer = keras.layers.Embedding(vocab_size_es, embed_size, mask_zero=True)
encoder_embeddings = encoder_embedding_layer(encoder_inputs_ids)
decoder_embeddings = decoder_embedding_layer(decoder_inputs_ids)

encoder = keras.layers.LSTM(512, return_sequences=True, return_state=True)
encoder_outputs, state_h, state_c = encoder(encoder_embeddings)
encoder_state = [state_h, state_c]

decoder = keras.layers.LSTM(512, return_sequences=True)
decoder_outputs = decoder(decoder_embeddings, initial_state=encoder_state)

# 使用Keras 3原生Attention层并传入mask
attention_layer = keras.layers.Attention()
attention_outputs = attention_layer(
    [decoder_outputs, encoder_outputs],
    mask=[decoder_outputs._keras_mask, encoder_outputs._keras_mask]
)

output_layer = keras.layers.Dense(vocab_size_es, activation="softmax")
Y_probas = output_layer(attention_outputs)

# 构建完整模型
model = keras.Model(inputs=[encoder_inputs, decoder_inputs], outputs=Y_probas)

额外提示 #

• Keras 3中建议统一使用keras.layers而非tf.keras.layers，以获得最佳多后端兼容性和KerasTensor支持。
• 英西翻译的encoder和decoder需使用各自独立的词汇表，避免共用同一个vocab_size导致后续输出错误。

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