TensorFlow模型保存/載入的方法有哪些

這篇文章主要介紹了TensorFlow模型保存/載入的方法有哪些，具有一定借鑒價(jià)值，感興趣的朋友可以參考下，希望大家閱讀完這篇文章之后大有收獲，下面讓小編帶著大家一起了解一下。

TensorFlow 模型保存/載入

我們?cè)谏暇€使用一個(gè)算法模型的時(shí)候，首先必須將已經(jīng)訓(xùn)練好的模型保存下來。tensorflow保存模型的方式與sklearn不太一樣，sklearn很直接，一個(gè)sklearn.externals.joblib的dump與load方法就可以保存與載入使用。而tensorflow由于有g(shù)raph, operation 這些概念，保存與載入模型稍顯麻煩。

一、基本方法

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

保存

• 定義變量

• 使用saver.save()方法保存
  

載入

• 定義變量

• 使用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")

這種方法不方便的在于，在使用模型的時(shí)候，必須把模型的結(jié)構(gòu)重新定義一遍，然后載入對(duì)應(yīng)名字的變量的值。但是很多時(shí)候我們都更希望能夠讀取一個(gè)文件然后就直接使用模型，而不是還要把模型重新定義一遍。所以就需要使用另一種方法。

二、不需重新定義網(wǎng)絡(luò)結(jié)構(gòu)的方法

tf.train.import_meta_graph

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

這個(gè)方法可以從文件中將保存的graph的所有節(jié)點(diǎn)加載到當(dāng)前的default graph中，并返回一個(gè)saver。也就是說，我們?cè)诒４娴臅r(shí)候，除了將變量的值保存下來，其實(shí)還有將對(duì)應(yīng)graph中的各種節(jié)點(diǎn)保存下來，所以模型的結(jié)構(gòu)也同樣被保存下來了。

比如我們想要保存計(jì)算最后預(yù)測(cè)結(jié)果的y，則應(yīng)該在訓(xùn)練階段將它添加到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, h2_dim], stddev=0.1), name='w1')
b1 = tf.Variable(tf.zeros([h2_dim]), name='b1')
w2 = tf.Variable(tf.zeros([h2_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ù)測(cè)目標(biāo)
y = tf.nn.softmax(tf.matmul(hidden1_drop, w2) + b2)
# 創(chuàng)建saver
saver = tf.train.Saver(...variables...)
# 假如需要保存y，以便在預(yù)測(cè)時(shí)使用
tf.add_to_collection('pred_network', y)
sess = tf.Session()
for step in xrange(1000000):
 sess.run(train_op)
 if step % 1000 == 0:
  # 保存checkpoint, 同時(shí)也默認(rèn)導(dǎo)出一個(gè)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() 返回一個(gè)list. 但是這里只要第一個(gè)參數(shù)即可
 y = tf.get_collection('pred_network')[0]

 graph = tf.get_default_graph()

 # 因?yàn)閥中有placeholder，所以sess.run(y)的時(shí)候還需要用實(shí)際待預(yù)測(cè)的樣本以及相應(yīng)的參數(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進(jìn)行預(yù)測(cè) 
 sess.run(y, feed_dict={input_x:...., keep_prob:1.0})

這里有兩點(diǎn)需要注意的：

一、saver.restore()時(shí)填的文件名，因?yàn)樵趕aver.save的時(shí)候，每個(gè)checkpoint會(huì)保存三個(gè)文件，如
my-model-10000.meta, my-model-10000.index, my-model-10000.data-00000-of-00001
在import_meta_graph時(shí)填的就是meta文件名，我們知道權(quán)值都保存在my-model-10000.data-00000-of-00001這個(gè)文件中，但是如果在restore方法中填這個(gè)文件名，就會(huì)報(bào)錯(cuò)，應(yīng)該填的是前綴，這個(gè)前綴可以使用tf.train.latest_checkpoint(checkpoint_dir)這個(gè)方法獲取。

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

感謝你能夠認(rèn)真閱讀完這篇文章，希望小編分享的“TensorFlow模型保存/載入的方法有哪些”這篇文章對(duì)大家有幫助，同時(shí)也希望大家多多支持億速云，關(guān)注億速云行業(yè)資訊頻道，更多相關(guān)知識(shí)等著你來學(xué)習(xí)!