這兩天搜索了不少關於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
希望這篇文章對您有幫助，感謝閱讀！
