如何在未知模型类型时获取TensorFlow .pb模型的输出节点名称

嘿，这个问题在处理未知结构的TensorFlow .pb模型时太常见了！我来给你分享几个实用的方法，不用额外加参数就能自动识别输出节点：

输出节点的核心特征是没有后续节点引用它，也就是出度为0。我们可以通过构建节点间的依赖关系来找出这类节点：

import tensorflow as tf
from tensorflow.core.framework import graph_pb2

def get_output_nodes(graph_def):
    # 先构建所有节点的输出依赖映射
    node_output_map = {}
    for node in graph_def.node:
        # 遍历当前节点的所有输入，记录输入节点的输出指向
        for input_name in node.input:
            # 处理节点名可能带的张量后缀（比如 `input:0` 只取 `input`）
            input_node_base_name = input_name.split(':')[0]
            if input_node_base_name not in node_output_map:
                node_output_map[input_node_base_name] = []
            node_output_map[input_node_base_name].append(node.name)
    
    # 找出没有被任何节点引用的节点（出度为0）
    output_nodes = []
    for node in graph_def.node:
        if node.name not in node_output_map or len(node_output_map[node.name]) == 0:
            output_nodes.append(node.name)
    return output_nodes

# 加载你的.pb模型
with tf.io.gfile.GFile('your_uploaded_model.pb', 'rb') as f:
    graph_def = graph_pb2.GraphDef()
    graph_def.ParseFromString(f.read())

# 获取输出节点
output_node_names = get_output_nodes(graph_def)
print("识别到的输出节点:", output_node_names)

这个方法几乎适用于所有TensorFlow模型，不管是分类、检测还是其他类型——毕竟输出节点肯定是整个计算图的终点。

很多模型的输出节点会使用特定的操作类型（Op），比如分类模型常用Softmax/ArgMax，检测模型常用Identity（导出时用来封装输出）、Squeeze，或者直接用业务相关的Op名比如detection_boxes。你可以扩展你之前的输入节点识别逻辑：

def get_output_nodes_by_op(graph_def):
    # 这里可以根据常见模型类型补充更多Op
    common_output_ops = ('Softmax', 'ArgMax', 'Identity', 'Squeeze', 'detection_boxes', 'detection_scores', 'detection_classes')
    output_nodes = [n.name for n in graph_def.node if n.op in common_output_ops]
    return output_nodes

# 使用示例
output_node_names = get_output_nodes_by_op(graph_def)
print("基于Op识别的输出节点:", output_node_names)

这个方法对公开的TensorFlow模型（比如Object Detection API、TensorFlow Hub的模型）特别好用，能快速定位到业务相关的输出节点。

为了避免误判（比如某些中间辅助节点也可能出度为0），可以把两种方法结合起来，优先选择同时满足出度为0和属于常见输出Op的节点：

def get_reliable_output_nodes(graph_def):
    # 步骤1：找出度为0的节点
    node_output_map = {}
    for node in graph_def.node:
        for input_name in node.input:
            input_node_base_name = input_name.split(':')[0]
            if input_node_base_name not in node_output_map:
                node_output_map[input_node_base_name] = []
            node_output_map[input_node_base_name].append(node.name)
    out_degree_zero_nodes = [n.name for n in graph_def.node if n.name not in node_output_map or len(node_output_map[n.name]) == 0]
    
    # 步骤2：筛选常见输出Op的节点
    common_output_ops = ('Softmax', 'ArgMax', 'Identity', 'Squeeze')
    op_matched_nodes = [n.name for n in graph_def.node if n.op in common_output_ops]
    
    # 取交集，交集为空则 fallback 到出度为0的节点
    reliable_outputs = list(set(out_degree_zero_nodes) & set(op_matched_nodes))
    if not reliable_outputs:
        reliable_outputs = out_degree_zero_nodes
    return reliable_outputs

• 很多模型会有多个输出节点（比如检测模型会同时输出detection_boxes、detection_scores、num_detections），上面的方法会全部识别出来，推理时可以一次性获取这些输出。
• 推理时可能需要给节点名加上张量后缀（比如:0），你可以在获取到节点名后拼接成{node_name}:0来获取对应的张量。

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