tf serving

我们平时使用tf.Saver()保存的模型是checkpoint格式的，

# construct graph!
...
# add save/restore ops
saver = tf.train.Saver()
...
# save after training
save_path = saver.save(sess, "/tmp/model.ckpt")

 

 

但是在TensorFlow Serving中一个servable的模型目录中是一个pb格式文件和一个名为variables的目录，因此需要在模型保存时就保存好可部署的模型格式

The Saver API saves the variables in checkpoint files and requires you to reconstruct the graph in order to load the variables. This is desired when you are splitting your training into separate sessions and want a quick way to resume training. But for loading a model in a different language or to “package” a complete model the SavedModel API is recommended.

 

-ckpt_model
        -checkpoint
        -***.ckpt.data-00000-of-00001
        -***.ckpt.index
        -***.ckpt.meta
#转换为
-servable_model
        -version
                -saved_model.pb
                -variables

 saved_model.pb包含图形结构SignatureDefs。variables文件夹保存训练所习得的权重。

SignatureDefs被用于定义一个计算（computation）的signature. 通常用于input keys，output keys以及method names

 

 

为了简单起见，我们使用一个非常简单的手写识别代码作为示例，代码如下：

from tensorflow.examples.tutorials.mnist import input_data
import tensorflow as tf

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

sess = tf.InteractiveSession()
x = tf.placeholder(tf.float32, [None, 784])
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

y = tf.nn.softmax(tf.matmul(x, W) + b)
y_ = tf.placeholder(tf.float32, [None, 10])
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), 1))

train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
tf.global_variables_initializer().run()

for i in range(1000):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    train_step.run({x: batch_xs, y_: batch_ys})

correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

print(accuracy.eval({x: mnist.test.images, y_: mnist.test.labels}))

 

要保存该模型，我们还需要对代码作一点小小的改动。

添加命名

 This step isn’t really necessary but it does make it a lot easier to load.

x = tf.placeholder(tf.float32, [None, 784], name="myInput")

y = tf.nn.softmax(tf.matmul(x, W) + b, name="myOutput") 

There are cases where you may have a final output node/tensor but didn’t get the chance to add a name for whatever reason. In those cases you could use tf.identity which allows you to add a name given a Tensor.

 

 

def addNameToTensor(someTensor, theName):
    return tf.identity(someTensor, name=theName)

 

save 

 

builder = tf.saved_model.builder.SavedModelBuilder("./model")

signature = predict_signature_def(inputs={'myInput': x},
                                  outputs={'myOutput': y})
builder.add_meta_graph_and_variables(sess=sess,
                                     tags=[tag_constants.SERVING],
                                     signature_def_map={'predict': signature})
builder.save()

 自己定义tag，在签名的定义上更加灵活。一个模型可以包含不同的MetaGraphDef，什么时候需要多个MetaGraphDef呢？也许你想保存图形的CPU版本和GPU版本，或者你想区分训练和发布版本。这个时候tag就可以用来区分不同的MetaGraphDef，加载的时候能够根据tag来加载模型的不同计算图。

The tag is used to distinguish different MetaGraphDef saved and is needed when loading the model. 

 

 

加载

对不同语言而言，加载过程有些类似，这里还是以python为例：

 

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, ["serve"], "./model")
  graph = tf.get_default_graph()

  input = np.expand_dims(mnist.test.images[0], 0)
  x = sess.graph.get_tensor_by_name('myInput:0')
  y = sess.graph.get_tensor_by_name('myOutput:0')
  batch_xs, batch_ys = mnist.test.next_batch(1)
  scores = sess.run(y,
           feed_dict={x: batch_xs})
  print("predict: %d, actual: %d" % (np.argmax(scores, 1), np.argmax(batch_ys, 1)))

 

Loading in Java

 

import org.tensorflow.*
SavedModelBundle savedModelBundle = SavedModelBundle.load("./export_path", "serve");
Graph graph = savedModelBundle.graph();
printOperations(graph);
Tensor result = savedModelBundle.session().runner()
                .feed("myInput", tensorInput)
                .fetch("myOutput")
                .run().get(0);