tensorflow 三种模型:ckpt、pb、pb-savemodel

1、CKPT
　　目录结构
　　　　　 checkpoint：
　　　　　model.ckpt-1000.index

　　　　model.ckpt-1000.data-00000-of-00001

　　　　model.ckpt-1000.meta

　　特点：
　　　　首先这种模型文件是依赖 TensorFlow 的，只能在其框架下使用；
　　　　数据和图是分开的
　　　　这种在训练的时候用的比较多。
　　代码：就省略了

2、pb模型-只有模型
　　这种方式只保存了模型的图结构，可以保留隐私的公布到网上。
　　感觉一些水的论文会用这种方式。
　　代码：
　　　　thanks：https://www.jianshu.com/p/9221fbf52c55
　　·　　

 import os

 import tensorflow as tf

 from tensorflow.python.saved_model import builder as saved_model_builder

 from tensorflow.python.saved_model import (signature_constants, signature_def_utils, tag_constants, utils)

 class model():

     def __init__(self):

         self.a = tf.placeholder(tf.float32, [None])

         self.w = tf.Variable(tf.constant(2.0, shape=[1]), name="w")

         b = tf.Variable(tf.constant(0.5, shape=[1]), name="b")

         self.y = self.a * self.w + b

 #模型保存为ckpt

 def save_model():

     graph1 = tf.Graph()

     with graph1.as_default():

         m = model()

     with tf.Session(graph=graph1) as session:

         session.run(tf.global_variables_initializer())

         update = tf.assign(m.w, [10])

         session.run(update)

         predict_y = session.run(m.y,feed_dict={m.a:[3.0]})

         print(predict_y)

         saver = tf.train.Saver()

         saver.save(session,"model_pb/model.ckpt")

 #保存为pb模型

 def export_model(session, m):

    #只需要修改这一段，定义输入输出，其他保持默认即可

     model_signature = signature_def_utils.build_signature_def(

         inputs={"input": utils.build_tensor_info(m.a)},

         outputs={

             "output": utils.build_tensor_info(m.y)},

         method_name=signature_constants.PREDICT_METHOD_NAME)

     export_path = "pb_model/1"

     if os.path.exists(export_path):

         os.system("rm -rf "+ export_path)

     print("Export the model to {}".format(export_path))

     try:

         legacy_init_op = tf.group(

             tf.tables_initializer(), name='legacy_init_op')

         builder = saved_model_builder.SavedModelBuilder(export_path)

         builder.add_meta_graph_and_variables(

             session, [tag_constants.SERVING],

             clear_devices=True,

             signature_def_map={

                 signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:

                     model_signature,

             },

             legacy_init_op=legacy_init_op)

         builder.save()

     except Exception as e:

         print("Fail to export saved model, exception: {}".format(e))

 #加载pb模型

 def load_pb():

     session = tf.Session(graph=tf.Graph())

     model_file_path = "pb_model/1"

     meta_graph = tf.saved_model.loader.load(session, [tf.saved_model.tag_constants.SERVING], model_file_path)

     model_graph_signature = list(meta_graph.signature_def.items())[0][1]

     output_tensor_names = []

     output_op_names = []

     for output_item in model_graph_signature.outputs.items():

         output_op_name = output_item[0]

         output_op_names.append(output_op_name)

         output_tensor_name = output_item[1].name

         output_tensor_names.append(output_tensor_name)

     print("load model finish!")

     sentences = {}

     # 测试pb模型

     for test_x in [[1],[2],[3],[4],[5]]:

         sentences["input"] = test_x

         feed_dict_map = {}

         for input_item in model_graph_signature.inputs.items():

             input_op_name = input_item[0]

             input_tensor_name = input_item[1].name

             feed_dict_map[input_tensor_name] = sentences[input_op_name]

         predict_y = session.run(output_tensor_names, feed_dict=feed_dict_map)

         print("predict pb y:",predict_y)

 if __name__ == "__main__":

     save_model()

     graph2 = tf.Graph()

     with graph2.as_default():

         m = model()

         saver = tf.train.Saver()

     with tf.Session(graph=graph2) as session:

         saver.restore(session, "model_pb/model.ckpt") #加载ckpt模型

         export_model(session, m)

     load_pb()

3、pb模型-Saved model
　　这是一种简单格式pb模型保存方式
　　目录结构
　　　　└── 1
　　　　···├── saved_model.pb
　　　　···└── variables
　　　　·········├── variables.data-00000-of-00001
　　　　·········└── variables.index
　　特点：
　　　　对于训练好的模型，我们都是用来进行使用的，也就是进行inference。
　　　　这个时候就不模型变化了。这种方式就将变量的权重变成了一个常亮。
　　　　这样方式模型会变小
　　　　在一些嵌入式吗，用C或者C++的系统中，我们也是常用.pb格式的。
　　代码：
　　　　thanks to https://www.jianshu.com/p/9221fbf52c55

 import os

 import tensorflow as tf

 from tensorflow.python.saved_model import builder as saved_model_builder

 from tensorflow.python.saved_model import (signature_constants, signature_def_utils, tag_constants, utils)

 class model():

     def __init__(self):

         self.a = tf.placeholder(tf.float32, [None])

         self.w = tf.Variable(tf.constant(2.0, shape=[1]), name="w")

         b = tf.Variable(tf.constant(0.5, shape=[1]), name="b")

         self.y = self.a * self.w + b

 #模型保存为ckpt

 def save_model():

     graph1 = tf.Graph()

     with graph1.as_default():

         m = model()

     with tf.Session(graph=graph1) as session:

         session.run(tf.global_variables_initializer())

         update = tf.assign(m.w, [10])

         session.run(update)

         predict_y = session.run(m.y,feed_dict={m.a:[3.0]})

         print(predict_y)

         saver = tf.train.Saver()

         saver.save(session,"model_pb/model.ckpt")

 #保存为pb模型

 def export_model(session, m):

    #只需要修改这一段，定义输入输出，其他保持默认即可

     model_signature = signature_def_utils.build_signature_def(

         inputs={"input": utils.build_tensor_info(m.a)},

         outputs={

             "output": utils.build_tensor_info(m.y)},

         method_name=signature_constants.PREDICT_METHOD_NAME)

     export_path = "pb_model/1"

     if os.path.exists(export_path):

         os.system("rm -rf "+ export_path)

     print("Export the model to {}".format(export_path))

     try:

         legacy_init_op = tf.group(

             tf.tables_initializer(), name='legacy_init_op')

         builder = saved_model_builder.SavedModelBuilder(export_path)

         builder.add_meta_graph_and_variables(

             session, [tag_constants.SERVING],

             clear_devices=True,

             signature_def_map={

                 signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:

                     model_signature,

             },

             legacy_init_op=legacy_init_op)

         builder.save()

     except Exception as e:

         print("Fail to export saved model, exception: {}".format(e))

 #加载pb模型

 def load_pb():

     session = tf.Session(graph=tf.Graph())

     model_file_path = "pb_model/1"

     meta_graph = tf.saved_model.loader.load(session, [tf.saved_model.tag_constants.SERVING], model_file_path)

     model_graph_signature = list(meta_graph.signature_def.items())[0][1]

     output_tensor_names = []

     output_op_names = []

     for output_item in model_graph_signature.outputs.items():

         output_op_name = output_item[0]

         output_op_names.append(output_op_name)

         output_tensor_name = output_item[1].name

         output_tensor_names.append(output_tensor_name)

     print("load model finish!")

     sentences = {}

     # 测试pb模型

     for test_x in [[1],[2],[3],[4],[5]]:

         sentences["input"] = test_x

         feed_dict_map = {}

         for input_item in model_graph_signature.inputs.items():

             input_op_name = input_item[0]

             input_tensor_name = input_item[1].name

             feed_dict_map[input_tensor_name] = sentences[input_op_name]

         predict_y = session.run(output_tensor_names, feed_dict=feed_dict_map)

         print("predict pb y:",predict_y)

 if __name__ == "__main__":

     save_model()

     graph2 = tf.Graph()

     with graph2.as_default():

         m = model()

         saver = tf.train.Saver()

     with tf.Session(graph=graph2) as session:

         saver.restore(session, "model_pb/model.ckpt") #加载ckpt模型

         export_model(session, m)

     load_pb()

巴特西

tensorflow 三种模型:ckpt、pb、pb-savemodel

最新文章

热门文章