Reputation: 1586
I am sorry in advance if the title does not reflect exactly my problem (I think it does, but I'm not sure), which I describe below.
I am working on converting a Yolo object detection model to a TensorFlow frozen model .pb
and then to use that model for prediction on mobile phones.
I have successfully obtained a working .pb
model (i.e. a frozen graph from Yolo's graph). But since the outputs of the network (there are two of them) are not the bounding boxes, I have to write a function for conversion (this part is not my question, I already have a working function for this task):
def get_boxes_from_output(outputs_of_the_graph, anchors,
num_classes, input_image_shape,
score_threshold=score, iou_threshold=iou)
"""
Apply some operations on the outputs_of_the_graph to obtain bounding boxes information
"""
return boxes, scores, classes
So the pipeline is simple: I have to load the pb
model, then throw the image data to it to get two outputs, then from these two outputs, I apply the above function (that contains tensor operations) to obtain bounding boxes information. The code look like this:
model_path = 'model_data/yolo.pb'
class_names = _get_class('model_data/classes.txt')
anchors = _get_anchors('model_data/yolo_anchors.txt')
score = 0.25
iou = 0.5
# Load the Tensorflow model into memory.
detection_graph = tf.Graph()
with detection_graph.as_default():
graph_def = tf.GraphDef()
with tf.gfile.GFile(model_path, 'rb') as fid:
graph_def.ParseFromString(fid.read())
tf.import_graph_def(graph_def, name='')
# Get the input and output nodes (there are two outputs)
l_input = detection_graph.get_tensor_by_name('input_1:0')
l_output = [detection_graph.get_tensor_by_name('conv2d_10/BiasAdd:0'),
detection_graph.get_tensor_by_name('conv2d_13/BiasAdd:0')]
#initialize_all_variables
tf.global_variables_initializer()
# Generate output tensor targets for filtered bounding boxes.
input_image_shape = tf.placeholder(dtype=tf.float32,shape=(2, ))
training = tf.placeholder(tf.bool, name='training')
boxes, scores, classes = get_boxes_from_output(l_output, anchors,
len(class_names), input_image_shape,
score_threshold=score, iou_threshold=iou)
image = Image.open('./data/image1.jpg')
image = preprocess_image(image)
image_data = np.array(image, dtype='float32')
image_data = np.expand_dims(image_data, 0) # Add batch dimension.
sess = tf.Session(graph=detection_graph)
# Run the session to get the output bounding boxes
out_boxes, out_scores, out_classes = sess.run(
[boxes, scores, classes],
feed_dict={
l_input: image_data,
input_image_shape: [image.size[1], image.size[0]],
training: False
})
# Now how do I save a new model that outputs directly [boxes, scores, classes]
Now my question is how do I save a new .pb
model from the session, so that I can load it again elsewhere and it can directly outputs boxes, scores, classes
?
I hope the question is clear enough.
Thank you very much in advance for your help!
Upvotes: 9
Views: 12204
Reputation: 17201
Once you have added the new ops, you need to write the new graph using tf.train.write_graph
:
boxes, scores, classes = get_boxes_from_output()
tf.train.write_graph(sess.graph_def,save_dir,'new_cnn_weights.pb')
Then you need to freeze the above graph using the freeze_graph
utility. Make sure the output_node_names
are set to boxes, scores, classes
as shown below:
# Freeze graph
from tensorflow.python.tools import freeze_graph
import os
input_graph_path = os.path.join(save_dir, 'new_cnn_weights.pb')
input_saver_def_path = ''
input_binary = False
output_node_names = 'boxes, scores, classes'
restore_op_name = ''
filename_tensor_name = ''
output_graph_path = os.path.join(save_dir, 'new_frozen_cnn_weights.pb')
clear_devices = False
checkpoint_path = os.path.join(save_dir, 'test_model')
freeze_graph.freeze_graph(input_graph_path, input_saver_def_path,
input_binary, checkpoint_path, output_node_names,
restore_op_name, filename_tensor_name,
output_graph_path, clear_devices, '')
#Load the new optimized graph and check whether the output is consistent,
tf.reset_default_graph()
with tf.gfile.GFile(save_dir+'new_frozen_cnn_weights.pb', 'rb') as f:
graph_def_optimized = tf.GraphDef()
graph_def_optimized.ParseFromString(f.read())
G = tf.Graph()
with tf.Session(graph=G) as sess:
boxes,scores,classes = tf.import_graph_def(graph_def_optimized, return_elements=['boxes:0', 'scores:0', 'classes:0'])
print('Operations in Optimized Graph:')
print([op.name for op in G.get_operations()])
x = G.get_tensor_by_name('import/import/input:0')
print(sess.run([boxes, scores, classes], feed_dict={x: np.expand_dims(img, 0)}))
Upvotes: 3