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來序列化模型圖。
本文的完整程式碼請參考：https://github.com/mogoweb/aiexamples/tree/master/tensorflow/saved_model
希望這篇文章對您有幫助，感謝閱讀！
