Somil Shah
Somil Shah

Reputation: 109

How do you find the output_node_names from a frozen model (a pb file) in Tensorflow?

I am trying to convert my frozen_model.pb to a tensorflow JS compatible (.pb) file, which is based on SSD Mobilenet V2 COCO pretrained model by Tensorflow. I am stuck at how to get the output_node_names parameter which is needed while using the tensorflowjs_converter. How do I get to know the output node names?

I have tried to get the operation names by using the below Python script, but am not able to understand which one is the output node.

def load_graph(model_file):
  graph = tf.Graph()
  graph_def = tf.GraphDef()

  with open(model_file, "rb") as f:
    graph_def.ParseFromString(f.read())
  with graph.as_default():
    tf.import_graph_def(graph_def)

  return graph

graph = load_graph('frozen_model.pb')
ops = graph.get_operations()

Upvotes: 3

Views: 4000

Answers (1)

hangqi
hangqi

Reputation: 141

Firstly, you can inspect all the nodes in your graph_def as follows:

for node in graph_def.node
    print(node.name)

Alternatively, if you want to visually see the graph and determine which node to be used as output, using TensorBoard is the way to go. There is a tool called import_pb_to_tensorboard. It's essentially using a handful of lines to write the graph to a log_dir, which you can point tensorboard to. You can simply copy those lines into your own script to achieve the same thing without building from the tensorflow repo.

Thirdly, there is another tool called summarize_graph tool:

bazel build tensorflow/tools/graph_transforms:summarize_graph
bazel-bin/tensorflow/tools/graph_transforms/summarize_graph --in_graph=/path/to/your/graph.pb

Upvotes: 5

Related Questions