如何保存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

 output = sess.graph.get_tensor_by_name("predict:0")

然后就是和训练的时候一样，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的值：

 pre = sess.run(tf.argmax(output, 1), feed_dict={input_x:data})

就报错了：

很明显，你需要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)))

测试结果：