CODEWITHSUNDEEP
pythonubuntutensorflowmachine-learningpre-trained-model
Dan D.
Dan D.

Reputation: 8567

Losing output node names in TensorFlow model after fine-tuning from pre-trained model

I follow the tutorial at https://tensorflow-object-detection-api-tutorial.readthedocs.io to fine-tune a pre-trained model to detect new objects in images. The pre-trained model is ssd_inception_v2_coco.

I trained and evaluated the model successfully after a few thousand steps with loss from 26 down to 1. However, I failed to create the frozen model with this code:

#this code runs in model dir
import tensorflow as tf

#make .pb file from model at step 1000
saver = tf.train.import_meta_graph(
        './model.ckpt-1000.meta', clear_devices=True)

graph = tf.get_default_graph()
input_graph_def = graph.as_graph_def()
sess = tf.Session()
saver.restore(sess, "./model.ckpt-1000")

#node names
i=0
for n in tf.get_default_graph().as_graph_def().node:
  print(n.name,i);    
  i+=1
#end for
print("total:",i);

output_node_names=[
  "detection_boxes","detection_classes",
  "detection_scores","num_detections"
];
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess,input_graph_def,output_node_names);

#save to .pb file
output_graph="./model.pb"
with tf.gfile.GFile(output_graph, "wb") as f:
  f.write(output_graph_def.SerializeToString());
#end with

sess.close();

The error is:

enter image description here

It seems the fine-tuned model has lost its output node names. There are these output node names in the original pre-trained model (change the checkpoint files in the code above to the ones in the original trained model): detection_boxes, detection_classes, detection_scores, and num_detections. The output node names are exactly in the original one, here are their indices (from the node name 'for' loop above):

enter image description here

My question is how to keep the output node names from the original pre-trained model? The node names are defined in code, but there's no code here, only some configs and the file 'train.py'.

PS. There's something called summary_op after total_loss, but I don't know whether it is the output(?):

enter image description here

Upvotes: 0

Views: 1153

Answers (1)

Dan D.
Dan D.

Reputation: 8567

In order to have 'image_tensor' (input), and other output node names 'detection_boxes', 'detection_classes', 'detection_scores', 'num_detections', use the utility script in tensorflow/models/research/object_detection named 'export_inference_graph.py'. This script even optimises the frozen graph (frozen model) for inference. As checked on my test model, the number of nodes reduced from 26,000 down to 5,000; this is great for inference speed.

Here's the link to export_inference_graph.py: https://github.com/tensorflow/models/blob/0558408514dacf2fe2860cd72ac56cbdf62a24c0/research/object_detection/export_inference_graph.py

How to run:

#bash command
python3 export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path PATH_TO_PIPELINE.config \
--trained_checkpoint_prefix PATH_TO/model.ckpt-NUMBER \
--output_directory PATH_TO_NEW_DIR

The .pb creating code in question only works on model created from scratch with node names defined manually, for a model checkpoint fine-tuned from a pre-trained model downloaded from TensorFlow Model Zoo https://github.com/tensorflow/models/blob/master/research/object_detection/g3doc/detection_model_zoo.md, it won't work!

Upvotes: 1

Related Questions