tensorflow model save and restore

TensorFlow 模型保存/載入

我們在上線使用一個算法模型的時候，首先必須將已經訓練好的模型保存下來。tensorflow保存模型的方式與sklearn不太一樣，sklearn很直接，一個sklearn.externals.joblib的dump與load方法就可以保存與載入使用。而tensorflow由於有graph, operation 這些概念，保存與載入模型稍顯麻煩。由於TensorFlow 的版本一直在更新, 保存模型的方法也發生了改變，在python 環境,和在C++ 環境(移動端) 等不同的平台需要的模型文件也是不也一樣的。

(https://stackoverflow.com/questions/44516609/tensorflow-what-is-the-relationship-between-ckpt-file-and-ckpt-meta-and-ckp有一個解釋如下)

• the .ckpt file is the old version output of saver.save(sess), which is the equivalent of your .ckpt-data (see below)

• the "checkpoint" file is only here to tell some TF functions which is the latest checkpoint file.

• .ckpt-meta contains the metagraph, i.e. the structure of your computation graph, without the values of the variables (basically what you can see in tensorboard/graph).

• .ckpt-data contains the values for all the variables, without the structure. To restore a model in python, you'll usually use the meta and data files with (but you can also use the .pb file):

 

saver = tf.train.import_meta_graph(path_to_ckpt_meta)
saver.restore(sess, path_to_ckpt_data)

 

• I don't know exactly for .ckpt-index, I guess it's some kind of index needed internally to map the two previous files correctly. Anyway it's not really necessary usually, you can restore a model with only .ckpt-meta and .ckpt-data.

• the .pb file can save your whole graph (meta + data). To load and use (but not train) a graph in c++ you'll usually use it, created with freeze_graph, which creates the .pb file from the meta and data. Be careful, (at least in previous TF versions and for some people) the py function provided by freeze_graph did not work properly, so you'd have to use the script version. Tensorflow also provides a tf.train.Saver.to_proto() method, but I don't know what it does exactly.

一、基本方法

網上搜索tensorflow模型保存，搜到的大多是基本的方法。即

保存

1. 定義變量
2. 使用saver.save()方法保存

載入

1. 定義變量
2. 使用saver.restore()方法載入

如 保存 代碼如下

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

二、不需重新定義網絡結構的方法

tf.train.import_meta_graph(
    meta_graph_or_file,
    clear_devices=False,
    import_scope=None,
    **kwargs
)

這個方法可以從文件中將保存的graph的所有節點加載到當前的default graph中，並返回一個saver。也就是說，我們在保存的時候，除了將變量的值保存下來，其實還有將對應graph中的各種節點保存下來，所以模型的結構也同樣被保存下來了。

比如我們想要保存計算最后預測結果的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)
# 創建saver
saver = 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. 但是這里只要第一個參數即可
  y = tf.get_collection('pred_network')[0]

  graph = tf.get_default_graph()

  # 因為y中有placeholder，所以sess.run(y)的時候還需要用實際待預測的樣本以及相應的參數來填充這些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文件名，我們知道權值都保存在my-model-10000.data-00000-of-00001這個文件中，但是如果在restore方法中填這個文件名，就會報錯，應該填的是前綴，這個前綴可以使用tf.train.latest_checkpoint(checkpoint_dir)這個方法獲取。

二、模型的y中有用到placeholder，在sess.run()的時候肯定要feed對應的數據，因此還要根據具體placeholder的名字，從graph中使用get_operation_by_name方法獲取。