tensorflow 之模型的保存與加載(三)
阿新 • • 發佈:2018-06-28
返回 toc lena time 第二篇 運算 element shape utf
前面的兩篇博文
第一篇:簡單的模型保存和加載,會包含所有的信息:神經網絡的op,node,args等;
第二篇:選擇性的進行模型參數的保存與加載。
本篇介紹,只保存和加載神經網絡的計算圖,即前向傳播的過程。
#!/usr/bin/env python3 #-*- coding:utf-8 -*- ############################ #File Name: save_restore.py #Brief: #Author: frank #Mail: [email protected] #Created Time:2018-06-26 20:30:09 ############################import tensorflow as tf from tensorflow.python.framework import graph_util #TF提供了convert_variables_to_constants函數,通過這個函數可以將計算圖中的變量及其取值通過常量的方式保存. v1 = tf.Variable(tf.constant(1.0, shape=[1]), name="v1") v2 = tf.Variable(tf.constant(2.0, shape=[1]), name="v2") result = v1 + v2 init_op = tf.global_variables_initializer() with tf.Session() as sess: sess.run(init_op)#導出當前計算圖的GraphDef部分,只需要這一步就可以完成從輸入層到輸出層的計算過程 graph_def = tf.get_default_graph().as_graph_def() #將圖中的變量及其取值轉化為常量,同時將圖中不必要的節點去掉.本程序只關心加法運算,所以這裏只保存[‘add‘]節點,其他和該計算無關的節點就不保存了 output_graph_def = graph_util.convert_variables_to_constants(sess, graph_def, [‘add‘]) with tf.gfile.GFile("./combine/combined_model.pb", "wb") as f: | f.write(output_graph_def.SerializeToString())
1 #!/usr/bin/env python3 2 #-*- coding:utf-8 -*- 3 ############################ 4 #File Name: save_restore2.py 5 #Brief: 6 #Author: frank 7 #Mail: [email protected] 8 #Created Time:2018-06-26 20:30:09 9 ############################ 10 11 import tensorflow as tf 12 from tensorflow.python.platform import gfile 13 14 with tf.Session() as sess: 15 model_filename = "combine/combined_model.pb" 16 #讀取保存的模型文件,並將文件解析成對應的GraphDef Protocol Buffer. 17 with gfile.FastGFile(model_filename, "rb") as f: 18 | graph_def = tf.GraphDef() 19 | graph_def.ParseFromString(f.read()) 20 21 #將graph_def中保存的圖加載到當前圖中.return_elements=["add:0"]給出了返回的張量的名稱.在保存的時候給出的是計算節點的名稱,所以為"add".在加載的時候給出的是張量的名稱,所以是add:0. 22 result = tf.import_graph_def(graph_def, return_elements=["add:0"]) 23 print(sess.run(result))
tensorflow 之模型的保存與加載(三)