如何儲存tensorflow完整模型，并加載已儲存模型

自己在使用tensorflow的時候，想要儲存下訓練好的的模型，以供下次使用。到網上看了很多教程，大多數使用的是tf.train.Saver(),這種方法還是太麻煩，沒法直接像其他架構一樣的儲存成一個黑盒，你隻要給輸入就行。後來找了很多的部落格，總算是找到了一種比較簡單的方法，就是使用tf.saved_model.builder。

接下來以一個CNN訓練mnist手寫數字識别的例子介紹這種方法

模型儲存

訓練及儲存的代碼如下：

import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)

batch_size = 100
n_batch = mnist.train.num_examples // batch_size

def weight_variable(shape):
    return tf.Variable(tf.truncated_normal(shape,stddev=0.1))

def bias_vairable(shape):
    return tf.Variable(tf.constant(0.1, shape=shape))

def conv2d(x,W):
    return tf.nn.conv2d(x,W,strides=[1,1,1,1],padding='SAME')

def max_pool_2x2(x):
    return tf.nn.max_pool(x,ksize=[1,2,2,1],strides=[1,2,2,1],padding='SAME')

x = tf.placeholder(tf.float32,[None,784], name="input_x")
y = tf.placeholder(tf.float32,[None,10], name="input_y")
keep_prob = tf.placeholder(tf.float32)

x_image = tf.reshape(x,[-1,28,28,1])

W_conv1 = weight_variable([5,5,1,32]) # 5*5的采樣視窗，32個卷積核從1個平面抽取特征
b_conv1 = bias_vairable([32]) #每個卷積核一個偏置值

# 28*28*1 的圖檔卷積之後變為28*28*32
h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
# 池化之後變為 14*14*32
h_pool1 = max_pool_2x2(h_conv1)

# 第二次卷積之後變為 14*14*64
W_conv2 = weight_variable([5,5,32,64])
b_conv2 = bias_vairable([64])
h_conv2 = tf.nn.relu(conv2d(h_pool1,W_conv2) + b_conv2)
# 第二次池化之後變為 7*7*64
h_pool2 = max_pool_2x2(h_conv2)


# 第一個全連接配接層
W_fc1 = weight_variable([7*7*64,1024])
b_fc1 = bias_vairable([1024])
# 7*7*64的圖像變成1維向量
h_pool2_flat = tf.reshape(h_pool2,[-1,7*7*64])
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

# 第二個全連接配接層
W_fc2 = weight_variable([1024,10])
b_fc2 = bias_vairable([10])
logits = tf.matmul(h_fc1_drop,W_fc2) + b_fc2
prediction = tf.nn.sigmoid(logits, name="predict")

loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(labels=y, logits=logits))
train_step = tf.train.AdamOptimizer(0.001).minimize(loss)

prediction_2 = tf.nn.softmax(prediction)
correct_prediction = (tf.equal(tf.argmax(prediction_2,1), tf.argmax(y,1)))
accuracy = tf.reduce_mean(tf.cast(correct_prediction,tf.float32))


with tf.Session() as sess:
    sess.run(tf.global_variables_initializer())
    writer = tf.summary.FileWriter("./log/", sess.graph)
    for epoch in range(21):
        for batch in range(n_batch):
            batch_xs, batch_ys = mnist.train.next_batch(batch_size)
            sess.run(train_step, feed_dict={x:batch_xs,y:batch_ys,keep_prob:0.7})
        acc = sess.run(accuracy, feed_dict={x:mnist.test.images, y:mnist.test.labels, keep_prob:1.0})
        print("Iter: " + str(epoch) + ", acc: " + str(acc))
    builder = tf.saved_model.builder.SavedModelBuilder("F:/xh/CNN/dd")
    builder.add_meta_graph_and_variables(sess,[tf.saved_model.tag_constants.TRAINING])
    builder.save()
           

儲存模型的代碼就是下面三行：

builder = tf.saved_model.builder.SavedModelBuilder("F:/xh/CNN/dd")
builder.add_meta_graph_and_variables(sess,[tf.saved_model.tag_constants.TRAINING])
builder.save()
           

複用的話改一下路徑就行了，注意，儲存的時候最好你最後一級路徑是空的。就比如"F:/xh/CNN/dd" dd這個檔案夾其實是不存在的，系統會自動建立。

這樣模型的圖和變量值都被存了下來，儲存完之後的檔案長這樣：

加載模型

1.将模型加載到Sessiong中

import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)


with tf.Session(graph= tf.Graph()) as sess:
   tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], "F:/xh/CNN/dd")
           

這一步比較簡單

2.擷取模型的輸入tensor

我們使用一個已經訓練好的模型，關鍵就在于找到輸入點和輸出點。把我們的測試資料賦給輸入點，擷取輸出點的預測值，就算是使用完成了。因為tensorflow的基本單元是tensor, 是以我們需要擷取到所有需要輸入的tensor名：

input_x = sess.graph.get_tensor_by_name("input_x:0")
   #input_y = sess.graph.get_tensor_by_name("input_y:0")
   keep_prob = sess.graph.get_tensor_by_name("Placeholder:0")
           

這裡括号中的tensor名怎麼檢視呢？

你可以自己在訓練的時候直接命名，像這樣：

x = tf.placeholder(tf.float32,[None,784], name="input_x")
y = tf.placeholder(tf.float32,[None,10], name="input_y")


           

然後加一個：0就可以了

那如果沒有自己命名怎麼辦？

利用tensorboard的檢視圖

檢視的模型：

放大一下，發現我們一共有三個可以手動輸入的tensor：

分别是

input_x,input_y和Placeholder

以擷取input_x的名字為例：

檢視input_x的後一個節點，發現輸入名就是input_x（這也是我在定義時直接就命名的關系，沒命名可能不是這個）,于是直接輸入

input_x = sess.graph.get_tensor_by_name(“input_x:0”) 即可

訓練模型，并擷取輸出

首先擷取一個輸出tensor

然後就是和訓練的時候一樣，feed一下所需的值,擷取輸出就行了：

data = mnist.train.images[0:1,:] #這個data就可以替換成需要的資料
   pre = sess.run(tf.argmax(output, 1), feed_dict={input_x:data, keep_prob:1.0})
   print(sess.run(tf.argmax(mnist.train.labels[0:10,:],1)))
           

這個地方注意，你必須要feed所有需要的值才行（input_y就可以不feed，因為測試的時候不需要标簽），比如一開始我隻賦了input_x的值：

就報錯了：

很明顯，你需要feed “Placeholder”這個tensor的值才行，加上之後，成功運作

載入加測試的整體代碼

import tensorflow as tf
import numpy as np
from tensorflow.examples.tutorials.mnist import input_data
mnist = input_data.read_data_sets('MNIST_data', one_hot=True)


with tf.Session(graph= tf.Graph()) as sess:
   tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.TRAINING], "F:/xh/CNN/dd")
   input_x = sess.graph.get_tensor_by_name("input_x:0")
   #input_y = sess.graph.get_tensor_by_name("input_y:0")
   keep_prob = sess.graph.get_tensor_by_name("Placeholder:0")
   output = sess.graph.get_tensor_by_name("predict:0")

   data = mnist.train.images[0:1,:]
   pre = sess.run(tf.argmax(output, 1), feed_dict={input_x:data, keep_prob:1.0})
   print("prediction is :", pre)
   print("true  is:",sess.run(tf.argmax(mnist.train.labels[0:1,:],1)))
           

測試結果：