TensorFlow 模型固化及生成TF Lite模型

convert_variables_to_constants

saver = tf.train.import_meta_graph(checkpoint + '.meta', clear_devices=True) #得到图、clear_devices :Whether or not to clear the device field for an `Operation` or `Tensor` during import.
sess = tf.InteractiveSession()
saver.restore(sess, checkpoint)

graph = tf.get_default_graph() #获得默认的图
input_graph_def = graph.as_graph_def()  #返回一个序列化的图代表当前的图

output_node_names="concat,Reshape_1"
output_graph_def = graph_util.convert_variables_to_constants(  #模型持久化,将变量值固定
  sess,
  input_graph_def,
  output_node_names.split(",") #如果有多个输出节点,以逗号隔开
)

with tf.gfile.GFile("lane.pb", "wb") as f: #保存模型
  f.write(output_graph_def.SerializeToString()) #序列化输出

测试模型

with tf.gfile.GFile('lane.pb', 'rb') as f:
  graph_def = tf.GraphDef()
  graph_def.ParseFromString(f.read())
  output = tf.import_graph_def(graph_def,
                               return_elements=['concat:0', 'Reshape_1:0'])

生成TF Lite模型

tflite_model = tf.contrib.lite.toco_convert(output_graph_def, [img], [output])


猜你喜欢

转载自blog.csdn.net/ywcpig/article/details/80456152