TensorFlow 模型保存/載入
我們在上線使用一個算法模型的時候，首先必須將已經(jīng)訓練好的模型保存下來。tensorflow保存模型的方式與sklearn不太一樣，sklearn很直接，一個sklearn.externals.joblib的dump與load方法就可以保存與載入使用。而tensorflow由于有g(shù)raph, operation 這些概念，保存與載入模型稍顯麻煩。
一、基本方法
網(wǎng)上搜索tensorflow模型保存，搜到的大多是基本的方法。即
保存
載入
如 保存 代碼如下
import tensorflow as tf import numpy as np W = tf.Variable([[1,1,1],[2,2,2]],dtype = tf.float32,name='w') b = tf.Variable([[0,1,2]],dtype = tf.float32,name='b') init = tf.initialize_all_variables() saver = tf.train.Saver() with tf.Session() as sess: sess.run(init) save_path = saver.save(sess,"save/model.ckpt")
載入代碼如下
import tensorflow as tf import numpy as np W = tf.Variable(tf.truncated_normal(shape=(2,3)),dtype = tf.float32,name='w') b = tf.Variable(tf.truncated_normal(shape=(1,3)),dtype = tf.float32,name='b') saver = tf.train.Saver() with tf.Session() as sess: saver.restore(sess,"save/model.ckpt")
這種方法不方便的在于，在使用模型的時候，必須把模型的結(jié)構(gòu)重新定義一遍，然后載入對應名字的變量的值。但是很多時候我們都更希望能夠讀取一個文件然后就直接使用模型，而不是還要把模型重新定義一遍。所以就需要使用另一種方法。
二、不需重新定義網(wǎng)絡結(jié)構(gòu)的方法
tf.train.import_meta_graphimport_meta_graph( meta_graph_or_file, clear_devices=False, import_scope=None, **kwargs)
這個方法可以從文件中將保存的graph的所有節(jié)點加載到當前的default graph中，并返回一個saver。也就是說，我們在保存的時候，除了將變量的值保存下來，其實還有將對應graph中的各種節(jié)點保存下來，所以模型的結(jié)構(gòu)也同樣被保存下來了。
比如我們想要保存計算最后預測結(jié)果的y，則應該在訓練階段將它添加到collection中。具體代碼如下
保存
### 定義模型input_x = tf.placeholder(tf.float32, shape=(None, in_dim), name='input_x')input_y = tf.placeholder(tf.float32, shape=(None, out_dim), name='input_y')w1 = tf.Variable(tf.truncated_normal([in_dim, h1_dim], stddev=0.1), name='w1')b1 = tf.Variable(tf.zeros([h1_dim]), name='b1')w2 = tf.Variable(tf.zeros([h1_dim, out_dim]), name='w2')b2 = tf.Variable(tf.zeros([out_dim]), name='b2')keep_prob = tf.placeholder(tf.float32, name='keep_prob')hidden1 = tf.nn.relu(tf.matmul(self.input_x, w1) + b1)hidden1_drop = tf.nn.dropout(hidden1, self.keep_prob)### 定義預測目標y = tf.nn.softmax(tf.matmul(hidden1_drop, w2) + b2)# 創(chuàng)建saversaver = tf.train.Saver(...variables...)# 假如需要保存y，以便在預測時使用tf.add_to_collection('pred_network', y)sess = tf.Session()for step in xrange(1000000): sess.run(train_op) if step % 1000 == 0: # 保存checkpoint, 同時也默認導出一個meta_graph # graph名為'my-model-{global_step}.meta'. saver.save(sess, 'my-model', global_step=step)載入
with tf.Session() as sess: new_saver = tf.train.import_meta_graph('my-save-dir/my-model-10000.meta') new_saver.restore(sess, 'my-save-dir/my-model-10000') # tf.get_collection() 返回一個list. 但是這里只要第一個參數(shù)即可 y = tf.get_collection('pred_network')[0] graph = tf.get_default_graph() # 因為y中有placeholder，所以sess.run(y)的時候還需要用實際待預測的樣本以及相應的參數(shù)來填充這些placeholder，而這些需要通過graph的get_operation_by_name方法來獲取。 input_x = graph.get_operation_by_name('input_x').outputs[0] keep_prob = graph.get_operation_by_name('keep_prob').outputs[0] # 使用y進行預測 sess.run(y, feed_dict={input_x:...., keep_prob:1.0})這里有兩點需要注意的：
一、saver.restore()時填的文件名，因為在saver.save的時候，每個checkpoint會保存三個文件，如
my-model-10000.meta, my-model-10000.index, my-model-10000.data-00000-of-00001
在import_meta_graph時填的就是meta文件名，我們知道權(quán)值都保存在my-model-10000.data-00000-of-00001這個文件中，但是如果在restore方法中填這個文件名，就會報錯，應該填的是前綴，這個前綴可以使用tf.train.latest_checkpoint(checkpoint_dir)這個方法獲取。
二、模型的y中有用到placeholder，在sess.run()的時候肯定要feed對應的數(shù)據(jù)，因此還要根據(jù)具體placeholder的名字，從graph中使用get_operation_by_name方法獲取。
以上就是本文的全部內(nèi)容，希望對大家的學習有所幫助，也希望大家多多支持VEVB武林網(wǎng)。
新聞熱點
疑難解答
圖片精選
