TensorFlow 2.* 如何获取Keras模型内部指定层的输出值

你构建的TensorFlow模型结构如下：

要获取特定输入下红色标记的Attention层（输出为5个浮点值）的计算结果，不需要改动原端到端模型的训练逻辑和结构，以下两种方案可直接落地：

方案1：构建专用特征提取子模型（推荐，适合反复验证场景） #

这是Keras官方推荐的中间层输出获取方式，逻辑清晰，推理稳定：

• 第一步先定位目标Attention层：如果建模时你手动给该层设置过name参数，可直接通过层名索引；如果没命名，先遍历所有层打印信息，找到输出形状匹配的Attention层即可：

# 遍历打印所有层的信息，找到输出形状最后一维为5的Attention层
for layer in original_model.layers:
    print(f"层名: {layer.name}, 层类型: {layer.__class__.__name__}, 输出形状: {layer.output_shape}")

• 第二步基于原模型构建独立的特征提取模型，输入和原模型保持一致，输出直接绑定目标Attention层的输出张量：

from tensorflow import keras

# 将引号内的内容替换为你上一步查到的目标Attention层名称
target_layer = original_model.get_layer(name="替换为你的目标Attention层名")
# 构建提取模型
attention_extractor = keras.Model(
    inputs=original_model.input,
    outputs=target_layer.output
)

• 第三步传入预处理好的特定测试输入，直接推理即可拿到目标结果：

# 注意测试输入要和模型训练时的预处理逻辑、输入维度完全一致，记得补batch维度
test_input =  # 你的验证用输入数据
# 推理得到Attention层输出的5个浮点值
attention_result = attention_extractor.predict(test_input, verbose=0)

多输入模型场景下，只需将inputs参数按原模型输入顺序传入所有输入层即可，其余逻辑不变。

方案2：用后端函数临时抓取输出（适合单次调试场景） #

如果只是临时调试拿一次结果，不需要留存子模型，可以直接用TensorFlow后端的函数接口绑定张量：

import tensorflow as tf

target_layer = original_model.get_layer(name="替换为你的目标Attention层名")
# 绑定输入输出
fetch_attention = tf.keras.backend.function(
    inputs=[original_model.input],
    outputs=[target_layer.output]
)
# 传入测试输入拿结果
attention_result = fetch_attention([test_input])[0]

验证逻辑正确性的小提示 #

拿到输出后如果要校验Attention层的特征提取逻辑是否符合预期，可以用同样的方法把该层的输入张量（query、key、value对应输入）也抓取出来，手动实现注意力分数计算、softmax归一化、加权求和的全流程，逐位和模型输出的5个值对比即可确认逻辑是否正确。

• 以上方案全部适配TensorFlow 2.x的eager执行模式，不需要额外启动Session，也不会改动原模型的权重参数。

内容的提问来源于stack exchange，提问作者S M Abrar Jahin