1. 程式人生 > >Tensorflow 訓練模型資料freeze固話儲存在Graph中

Tensorflow 訓練模型資料freeze固話儲存在Graph中

在tensorflow中,graph是訓練的核心,當一個模型訓練完成後,需要將模型儲存下來,一個通常的操作是:

variables = tf.all_variables()
                saver = tf.train.Saver(variables)
                saver.save(sess, "data/data.ckpt")
tf.train.write_graph(sess.graph_def, 'graph', 'model.ph', False)

這樣就可以將model儲存在model.ph檔案中,然而使用的時候不僅要載入模型檔案model.ph,還要載入儲存的data.ckpt資料檔案才能使用。這樣保持了資料與模型的分離,確實是個不錯的方法。
當我們把一個訓練模型完整的訓練好上線時候,我們期待的場景是:將一張圖片喂進去,然後得出結果。 這時候再這樣載入或許有些不必要,特別是在一些變數”不明”的時候特別麻煩.這時候一個比較好的方法就是將變數(偏執,權重等)固化到模型資料中。

建立圖

在檔案開頭增加如下程式碼
這裡寫圖片描述

宣告tensor

在需要的操作新增
這裡寫圖片描述

固化儲存

這裡寫圖片描述

固化操作中最重要的函式是:

tf.graph_util.convert_variables_to_constants(sess, input_graph_def, output_node_names, variable_names_whitelist=None, variable_names_blacklist=None)

程式碼執行後控制檯列印:
這裡寫圖片描述
這樣在我們使用的時候就不要再進行data.ckpt的資料恢復。直接通過:

sess.graph.get_tensor_by_name()

就可以獲取一個tensor,是不是很方便。

小報錯:

  • assert d in name_to_node_map, “%s is not in graph” % d
    AssertionError: A is not in graph
    解決:仔細檢視tensor名稱是否正確。