如何获取TensorFlow Keras模型中包含的tf.Operation操作

1. 提取模型计算图中的所有操作 #

首先拿到模型对应的计算图，筛选出你需要的目标算子：

import tensorflow as tf

# 你的示例模型
x = inp = tf.keras.Input((5,))
x = tf.keras.layers.Dense(7, activation="relu")(x)
x = tf.keras.layers.Dense(1)(x)
model = tf.keras.Model(inp, x)

# 获取模型前向计算的静态图
forward_func = tf.function(model).get_concrete_function(tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype))
graph = forward_func.graph

# 获取所有tf.Operation对象
all_ops = graph.get_operations()

# 示例：筛选所有Relu算子的操作
target_ops = [op for op in all_ops if op.type == "Relu"]
# 提取操作对应的输出张量
target_tensors = [op.outputs[0] for op in target_ops]

2. 构建提取输出的新模型 #

直接基于原模型的输入和筛选得到的目标张量构建新的Keras模型，调用时就可以直接输出所有目标张量的值：

# 构建多输出提取模型
extract_model = tf.keras.Model(
    inputs=model.inputs,
    outputs=target_tensors
)

# 测试调用，输入任意符合形状的张量即可
test_input = tf.random.normal((2, 5))
result = extract_model(test_input)

# 输出为所有目标算子的计算结果，顺序和你筛选的target_ops顺序一致
print(f"共提取到{len(result)}个Relu算子的输出，每个输出形状为{result[0].shape}")

• 筛选算子时也可以通过op.name匹配自定义关键词，适用于你知道目标操作名称前缀的场景
• 该方法完全复用原模型的权重，不需要重新训练，对结构复杂的嵌套子层模型同样有效
• 如果你的模型是多输入结构，在调用get_concrete_function时需要按顺序传入所有输入对应的tf.TensorSpec即可
• TensorFlow 2.6.0版本完全兼容上述用法，无需额外安装依赖

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