TensorFlow:读取冻结模型,添加操作,然后保存到新的冻结模型

时间:2018-05-29 17:53:24

标签: tensorflow

我很抱歉,如果标题不能完全反映我的问题(我认为确实如此,但我不确定),我将在下面介绍。

我正在努力将Yolo对象检测模型转换为TensorFlow冻结模型company_name | group | country | status com3 | 1 | UK | deleted com4 | 2 | ES | deleted com5 | 2 | FR | deleted com6 | 3 | RO | deleted com7 | 3 | BG | deleted ,然后使用该模型在手机上进行预测。

我已经成功获得了一个有效的.pb模型(即Yolo图中的冻结图)。但由于网络的输出(有两个)不是边界框,我必须编写一个转换函数(这部分不是我的问题,我已经有了这个任务的工作函数):

.pb

所以管道很简单:我必须加载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 模型,然后将图像数据抛给它以获得两个输出,然后从这两个输出中,我应用上面的函数(包含张量操作)获取边界框信息。代码如下所示:

pb

现在我的问题是如何从会话中保存新的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] 模型,以便我可以在其他地方重新加载它,它可以直接输出.pb < / p>

我希望这个问题很清楚。

非常感谢您的帮助!

1 个答案:

答案 0 :(得分:0)

添加节点并保存到冻结模型

添加新操作后,您需要使用tf.train.write_graph编写新图表:

boxes, scores, classes = get_boxes_from_output()
tf.train.write_graph(sess.graph_def,save_dir,'new_cnn_weights.pb')

然后,您需要使用freeze_graph实用程序冻结上图。确保output_node_names设置为boxes, scores, classes,如下所示:

# 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)}))