Tensorflow SavedModel模型的保存與加載

這兩天搜索了很多關於Tensorflow模型保存與加載的資料，發現不少資料都是關於checkpoints模型格式的，而最新的SavedModel模型格式則資料較少，爲此總結一下TensorFlow如何保存SavedModel模型，並加載之。

爲何要採用SavedModel格式呢？其主要優勢是SaveModel與語言無關，好比可使用python語言訓練模型，而後在Java中很是方便的加載模型。固然這也不是說checkpoints模型格式作不到，只是在跨語言時比較麻煩。另外若是使用Tensorflow Serving server來部署模型，必須選擇SavedModel格式。

SavedModel包含啥？

一個比較完整的SavedModel模型包含如下內容：

assets/
assets.extra/
variables/
    variables.data-*****-of-*****
    variables.index
saved_model.pb
複製代碼

saved_model.pb是MetaGraphDef，它包含圖形結構。variables文件夾保存訓練所習得的權重。assets文件夾能夠添加可能須要的外部文件，assets.extra是一個庫能夠添加其特定assets的地方。

MetaGraph是一個數據流圖，加上其相關的變量、assets和簽名。MetaGraphDef是MetaGraph的Protocol Buffer表示。

assets和assets.extra是可選的，好比本文示例代碼保存的模型只包含如下的內容：

variables/
    variables.data-*****-of-*****
    variables.index
saved_model.pb
複製代碼

保存

爲了簡單起見，咱們使用一個很是簡單的手寫識別代碼做爲示例，代碼以下：

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}))
複製代碼

這段代碼很簡單，一個簡單的梯度遞減迴歸模型。要保存該模型，咱們還須要對代碼做一點小小的改動。

添加命名

在輸入和輸出Ops中添加名稱，這樣咱們在加載時能夠方便的按名稱引用操做。將上述的x賦值語句修改成：

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

固然你也能夠不給名稱，系統會默認給一個名稱，好比上面的x系統會給一個"Placeholder"，當咱們須要引用多個op的時候，給每一個op一個命名，確實方便給咱們後面使用。

你也可使用tf.identity給tensor命名，好比在上述代碼上添加一行：

tf.identity(y, name="myOutput")
複製代碼

給輸出也命一個名。

保存到文件

最簡單的保存方法是使用tf.saved_model.simple_save函數，代碼以下：

tf.saved_model.simple_save(sess,
            "./model",
            inputs={"myInput": x},
            outputs={"myOutput": y})
複製代碼

這段代碼將模型保存在**./model**目錄。

固然你也能夠採用比較複雜的寫法：

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，在簽名的定義上更加靈活。這裏說說tag的用途吧。

一個模型能夠包含不一樣的MetaGraphDef，何時須要多個MetaGraphDef呢？也許你想保存圖形的CPU版本和GPU版本，或者你想區分訓練和發佈版本。這個時候tag就能夠用來區分不一樣的MetaGraphDef，加載的時候可以根據tag來加載模型的不一樣計算圖。

在simple_save方法中，系統會給一個默認的tag: 「serve」，也能夠用tag_constants.SERVING這個常量。

加載

對不一樣語言而言，加載過程有些相似，這裏仍是以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)))
複製代碼

須要注意，load函數中第二個參數是tag，須要和保存模型時的參數一致，第三個參數是模型保存的文件夾。

調用load函數後，不只加載了計算圖，還加載了訓練中習得的變量值，有了這二者，咱們就能夠調用其進行推斷新給的測試數據。

小結

將過程捋順了以後，你會發覺保存和加載SavedModel其實很簡單。但在摸索過程當中，也走了很多的彎路，主要緣由是如今搜索到的大部分資料仍是用tf.train.Saver()來保存模型，還有的是用tf.gfile.FastGFile來序列化模型圖。

本文的完整代碼請參考：github.com/mogoweb/aie…

但願這篇文章對您有幫助，感謝閱讀！