最近接觸了ｔｅｎｓｏｒｆｌｏｗ的ｏｂｊｅｃｔ　ｄｅｔｅｃｔｉｏｎ　ＡＰＩ發現裡面讀取的預先訓練模型都是ｐｂ格式。

谷歌推薦的儲存模型的方式是儲存模型為 PB 檔案，它具有語言獨立性，可獨立執行，封閉的序列化格式，任何語言都可以解析它，它允許其他語言和深度學習框架讀取、繼續訓練和遷移 TensorFlow 的模型。

它的主要使用場景是實現建立模型與使用模型的解耦， 使得前向推導 inference的程式碼統一。

另外的好處是儲存為 PB 檔案時候，模型的變數都會變成固定的，導致模型的大小會大大減小，適合在手機端執行。

還有一個就是，真正離線測試使用的時候，ｐｂ格式的資料能夠保證資料不會更新變動，就是不會進行反饋調節啦。

儲存 PB 檔案的程式碼：

import tensorflow as tf
import os
from tensorflow.python.framework import graph_util

pb_file_path = os.getcwd()

with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf
 
.multiply(x, y)
    # 這裡的輸出需要加上name屬性
    op = tf.add(xy, b, name='op_to_store')

    sess.run(tf.global_variables_initializer())

    # convert_variables_to_constants 需要指定output_node_names，list()，可以多個
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'
 
])

    # 測試 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op, feed_dict))

    # 寫入序列化的 PB 檔案
    with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

    # 輸出
    # INFO:tensorflow:Froze 1 variables.
    # Converted 1 variables to const ops.
    # 31

載入 PB 模型檔案典型程式碼：

from tensorflow.python.platform import gfile

sess = tf.Session()
with gfile.FastGFile(pb_file_path+'model.pb', 'rb') as f:
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    sess.graph.as_default()
    tf.import_graph_def(graph_def, name='') # 匯入計算圖

# 需要有一個初始化的過程    
sess.run(tf.global_variables_initializer())

# 需要先復原變數
print(sess.run('b:0'))
# 1

# 輸入
input_x = sess.graph.get_tensor_by_name('x:0')
input_y = sess.graph.get_tensor_by_name('y:0')

op = sess.graph.get_tensor_by_name('op_to_store:0')

ret = sess.run(op,  feed_dict={input_x: 5, input_y: 5})
print(ret)
# 輸出 26

儲存為 save model 格式也可以生成模型的 PB 檔案，並且更加簡單。

儲存好以後到saved_model_dir目錄下，會有一個saved_model.pb檔案以及variables資料夾。顧名思義，variables儲存所有變數，saved_model.pb用於儲存模型結構等資訊。

import tensorflow as tf
import os
from tensorflow.python.framework import graph_util

pb_file_path = os.getcwd()

with tf.Session(graph=tf.Graph()) as sess:
    x = tf.placeholder(tf.int32, name='x')
    y = tf.placeholder(tf.int32, name='y')
    b = tf.Variable(1, name='b')
    xy = tf.multiply(x, y)
    # 這裡的輸出需要加上name屬性
    op = tf.add(xy, b, name='op_to_store')

    sess.run(tf.global_variables_initializer())

    # convert_variables_to_constants 需要指定output_node_names，list()，可以多個
    constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

    # 測試 OP
    feed_dict = {x: 10, y: 3}
    print(sess.run(op, feed_dict))

    # 寫入序列化的 PB 檔案
    with tf.gfile.FastGFile(pb_file_path+'model.pb', mode='wb') as f:
        f.write(constant_graph.SerializeToString())

    # INFO:tensorflow:Froze 1 variables.
    # Converted 1 variables to const ops.
    # 31
    
    
    # 官網有誤，寫成了 saved_model_builder  
    builder = tf.saved_model.builder.SavedModelBuilder(pb_file_path+'savemodel')
    # 構造模型儲存的內容，指定要儲存的 session，特定的 tag, 
    # 輸入輸出資訊字典，額外的資訊
    builder.add_meta_graph_and_variables(sess,
                                       ['cpu_server_1'])


# 新增第二個 MetaGraphDef 
#with tf.Session(graph=tf.Graph()) as sess:
#  ...
#  builder.add_meta_graph([tag_constants.SERVING])
#...

builder.save()  # 儲存 PB 模型

這種方法對應的匯入模型的方法：

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ['cpu_1'], pb_file_path+'savemodel')
    sess.run(tf.global_variables_initializer())

    input_x = sess.graph.get_tensor_by_name('x:0')
    input_y = sess.graph.get_tensor_by_name('y:0')

    op = sess.graph.get_tensor_by_name('op_to_store:0')

    ret = sess.run(op,  feed_dict={input_x: 5, input_y: 5})
    print(ret)
# 只需要指定要恢復模型的 session，模型的 tag，模型的儲存路徑即可,使用起來更加簡單

這樣和之前的匯入 PB 模型一樣，也是要知道tensor的name。那麼如何可以在不知道tensor name的情況下使用呢，實現徹底的解耦呢？ 給add_meta_graph_and_variables方法傳入第三個引數，signature_def_map即可。

參考：

https://zhuanlan.zhihu.com/p/32887066