有時候,我們需要將TensorFlow的模型導出為單個文件(同時包含模型架構定義與權重),方便在其他地方使用(如在c++中部署網絡)。利用tf.train.write_graph()默認情況下只導出了網絡的定義(沒有權重),而利用tf.train.Saver().save()導出的文件graph_def與權重是分離的,因此需要采用別的方法。
我們知道,graph_def文件中沒有包含網絡中的Variable值(通常情況存儲了權重),但是卻包含了constant值,所以如果我們能把Variable轉換為constant,即可達到使用一個文件同時存儲網絡架構與權重的目標。
我們可以采用以下方式凍結權重并保存網絡:
import tensorflow as tffrom tensorflow.python.framework.graph_util import convert_variables_to_constants# 構造網絡a = tf.Variable([[3],[4]], dtype=tf.float32, name='a')b = tf.Variable(4, dtype=tf.float32, name='b')# 一定要給輸出tensor取一個名字!!output = tf.add(a, b, name='out')# 轉換Variable為constant,并將網絡寫入到文件with tf.Session() as sess: sess.run(tf.global_variables_initializer()) # 這里需要填入輸出tensor的名字 graph = convert_variables_to_constants(sess, sess.graph_def, ["out"]) tf.train.write_graph(graph, '.', 'graph.pb', as_text=False)
當恢復網絡時,可以使用如下方式:
import tensorflow as tfwith tf.Session() as sess: with open('./graph.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) output = tf.import_graph_def(graph_def, return_elements=['out:0']) print(sess.run(output))輸出結果為:
[array([[ 7.],
[ 8.]], dtype=float32)]
可以看到之前的權重確實保存了下來!!
問題來了,我們的網絡需要能有一個輸入自定義數據的接口啊!不然這玩意有什么用。。別急,當然有辦法。
import tensorflow as tffrom tensorflow.python.framework.graph_util import convert_variables_to_constantsa = tf.Variable([[3],[4]], dtype=tf.float32, name='a')b = tf.Variable(4, dtype=tf.float32, name='b')input_tensor = tf.placeholder(tf.float32, name='input')output = tf.add((a+b), input_tensor, name='out')with tf.Session() as sess: sess.run(tf.global_variables_initializer()) graph = convert_variables_to_constants(sess, sess.graph_def, ["out"]) tf.train.write_graph(graph, '.', 'graph.pb', as_text=False)
用上述代碼重新保存網絡至graph.pb,這次我們有了一個輸入placeholder,下面來看看怎么恢復網絡并輸入自定義數據。
import tensorflow as tfwith tf.Session() as sess: with open('./graph.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) output = tf.import_graph_def(graph_def, input_map={'input:0':4.}, return_elements=['out:0'], name='a') print(sess.run(output))
新聞熱點
疑難解答