1. 程式人生 > >TensorFlow 模型固化及生成TF Lite模型

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])