Keras 3.7.0中Attention层集成到自定义TCAN模块时无法返回注意力分数的问题

你遇到的问题核心在于Keras的Eager执行模式与符号计算图模式下，Attention层的输出形式存在本质差异：单独测试时能正常解包，但在自定义模型块的函数式API构建中会触发报错，下面为你详细分析并给出解决方案。

问题原因分析 #

• 单独测试的Eager模式：你的测试代码是在Eager执行模式下运行的，此时层调用返回的是实际Tensor对象组成的Python元组，因此可以直接用output, attention_scores = ...的语法解包。
• 自定义块的符号计算图模式：在构建模型的函数式API中，所有操作都是在创建符号化的计算图，此时Attention层返回的是一个KerasTensor复合结构（而非Python元组）。KerasTensor是符号化张量，不支持Python的迭代解包语法（即未实现__iter__方法），因此直接解包会触发NotImplementedError。

解决方案：通过索引访问复合输出 #

在函数式API的模型构建中，不要直接解包返回结果，而是将Attention层的输出视为可索引的结构，通过索引提取输出张量和注意力分数。修改你的tcan_block函数中注意力层的调用逻辑即可：

修改后的TCAN块代码 #

from tensorflow.keras.layers import Attention, Dense, Conv1D, Lambda, Add
import tensorflow.keras.backend as K

def tcan_block(inputs, filters, kernel_size, activation, dilation_rate, d_k, atn_dropout):
    """
    A single block of TCAN.
    Arguments:
        inputs: Tensor, input sequence.
        filters: Integer, number of filters for the convolution.
        kernel_size: Integer, size of the convolution kernel.
        dilation_rate: Integer, dilation rate for the convolution.
        d_k: Integer, dimensionality of the attention keys/queries.
    Returns:
        Tensor, output of the TCAN block.
    """
    # Temporal Attention
    query = Dense(d_k)(inputs)
    key = Dense(d_k)(inputs)
    value = Dense(d_k)(inputs)

    # 1. 初始化Attention层并调用，获取复合输出结构
    attention_layer = Attention(use_scale=True, dropout=atn_dropout)
    attention_outputs = attention_layer(
        [query, value, key],
        use_causal_mask=True,
        return_attention_scores=True,
    )
    
    # 2. 通过索引提取输出张量和注意力分数（替代直接解包）
    attention_output = attention_outputs[0]
    attention_scores = attention_outputs[1]

    # Dilated Convolution
    conv_output = Conv1D(
        filters, kernel_size, dilation_rate=dilation_rate, padding="causal", activation=activation
    )(attention_output)

    # Enhanced Residual
    importance = Lambda(lambda x: K.cumsum(x, axis=1))(attention_scores)
    enhanced_residual = Lambda(lambda x: x[0] * x[1])([inputs, importance])

    # Add residual connection
    output = Add()([inputs, conv_output, enhanced_residual])
    return output

验证修改有效性 #

你可以通过构建一个简单的测试模型来验证修改后的TCAN块是否正常工作：

from tensorflow.keras.models import Model
from tensorflow.keras.layers import Input

# 构建测试模型
test_input = Input(shape=(8, 16))
test_output = tcan_block(
    test_input, 
    filters=32, 
    kernel_size=3, 
    activation='relu', 
    dilation_rate=1, 
    d_k=16, 
    atn_dropout=0.1
)
model = Model(inputs=test_input, outputs=test_output)
model.summary()

该代码应该能正常生成模型结构摘要，不会触发之前的NotImplementedError。

额外扩展：暴露注意力分数作为模型输出 #

如果你需要将注意力分数作为模型的输出（比如后续用于可视化或分析），可以让tcan_block返回多输出结果：

def tcan_block(inputs, filters, kernel_size, activation, dilation_rate, d_k, atn_dropout):
    # ... 前面的代码保持不变 ...
    # 返回主输出和注意力分数的列表
    return [output, attention_scores]

# 构建多输出模型
test_input = Input(shape=(8, 16))
tcan_main_output, attention_scores_output = tcan_block(
    test_input, 
    filters=32, 
    kernel_size=3, 
    activation='relu', 
    dilation_rate=1, 
    d_k=16, 
    atn_dropout=0.1
)
model = Model(inputs=test_input, outputs=[tcan_main_output, attention_scores_output])

备注：内容来源于stack exchange，提问作者Furkan Öztürk