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模型-只有模型
  这种方式只保存了模型的图结构,可以保留隐私的公布到网上。
  感觉一些水的论文会用这种方式。
 

  模型直接下载下来是saved model形式的,需要依据他的脚本将saved model格式的模型转换为pb格式的
  现在就可以讲解pb模型的加载了

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


 事件1、从git上下载预训练模型,进行使用  

  模型下载:https://github.com/openvinotoolkit/open_model_zoo/blob/master/models/public/resnet-50-tf/resnet-50-tf.md

  下载模型的形式:
    saved_model形式
  自己走的弯路
    这里面有一个frozen_saved_model.py这么一个脚本,
    因为这个脚本的存在,令自己在pb和saved model这两种格式之间产生了混淆
  弯路解析
    frozen_saved_model.py这个脚本的作用是将saved model 格式的模型转换成 pb格式的。
    因为pb格式是嵌入式应用的,saved model这个格式,同事保存计算图和变量,
  执行frozen_saved_model.py
    将saved model 转换为frozen的pb格式,会在文件下生成一个文件
    

import argparse
import tensorflow.compat.v1 as tf

from tensorflow.python.framework import graph_io
from tensorflow.python.tools import optimize_for_inference_lib

def parse_args():
    """Parse input arguments"""

    parser = argparse.ArgumentParser(description='Freeze saved model')

    parser.add_argument('--saved_model_dir', type=str, #required=True,
                        default= './model',
                        help='Path to saved model directory.')
    parser.add_argument('--save_file', type=str, #required=True,
                        default='frozen_model.pb',
                        help='Path to resulting frozen model.')
    return parser.parse_args()

def freeze(saved_model_dir, input_nodes, output_nodes, save_file):
    graph_def = tf.Graph()
    with tf.Session(graph=graph_def) as sess:
        tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], saved_model_dir)
        frozen_graph_def = tf.graph_util.convert_variables_to_constants(
            sess,
            sess.graph_def,
            output_nodes
        )
        frozen_graph_def = optimize_for_inference_lib.optimize_for_inference(
            frozen_graph_def,
            input_nodes,
            output_nodes,
            tf.float32.as_datatype_enum
        )
        with open(save_file, 'wb') as f:
            f.write(frozen_graph_def.SerializeToString())

def main():
    args = parse_args()
    input_nodes = ['map/TensorArrayStack/TensorArrayGatherV3']
    output_nodes = ['softmax_tensor']
    freeze(args.saved_model_dir, input_nodes, output_nodes, args.save_file)

if __name__ == '__main__':
    main()
View Code

  载入 pb 模型执行推理

'''
功能:读取.pb格式的模型,并使用模型进行预测。注意这里不是saved_model的格式
'''
import tensorflow as tf
from tensorflow.python.platform import gfile
from PIL import Image
import numpy as np

#将tf网络reset确保清除缓存,不然有可能这个计算图还是之前没有修改的
tf.reset_default_graph()

#导入计算图
sess = tf.Session()
with gfile.FastGFile('frozen_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())

    # 获取这个模型的输入输出的tensor的名字
    input_img = sess.graph.get_tensor_by_name('map/TensorArrayStack/TensorArrayGatherV3:0')
    output_tenosr = sess.graph.get_tensor_by_name('softmax_tensor:0')

    #读入一张图片,
    #shape:必须是[B, 224,224,3]
    #PNG格式的图片通道是4通道的,还有一个透明度通道,自己用Windows的绘制进行转换
    #输入的格式必须是numpy的array格式。验证表示tf的tensor 会报错
    #上面这几个需要根据自己的实际情况来改动。
    img = np.array(Image.open('./dog.jpg').resize((224,224)))
    img = img[np.newaxis, :]
    print(img.shape)
    
    #打印输出
    ret = sess.run(output_tenosr,  feed_dict={input_img: img})
    print(ret)#输出一个tensor分别对应imageNet的1000个类别
pb模型加载

 

saved model模型加载

  注意这里就不再用pb这个模型文件了,而是使用下saved_model.pb所在的文件夹,

import tensorflow as tf
import numpy as np
from PIL import Image

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ["serve"], "./model")
    #这个model下面有variable文件夹和saved_model.pb文件
    sess.run(tf.global_variables_initializer())

    # 获取这个模型的输入输出的tensor的名字
    input_img = sess.graph.get_tensor_by_name('map/TensorArrayStack/TensorArrayGatherV3:0')
    output_tenosr = sess.graph.get_tensor_by_name('softmax_tensor:0')

    # 读入一张图片,
    # shape:必须是[B, 224,224,3]
    # PNG格式的图片通道是4通道的,还有一个透明度通道,自己用Windows的绘制进行转换
    # 输入的格式必须是numpy的array格式。验证表示tf的tensor 会报错
    # 上面这几个需要根据自己的实际情况来改动。
    img = np.array(Image.open('./dog.jpg').resize((224, 224)))
    img = img[np.newaxis, :]
    print(img.shape)

    # 打印输出
    ret = sess.run(output_tenosr, feed_dict={input_img: img})
    print(np.array(ret).sum())
View Code

下面是从其他博客里面获取的一个完整的例子,可以作为辅助参考
    先定义一个网络

 1 from tensorflow.examples.tutorials.mnist import input_data
 2 import tensorflow as tf
 3 
 4 mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
 5 
 6 sess = tf.InteractiveSession()
 7 x = tf.placeholder(tf.float32, [None, 784])
 8 W = tf.Variable(tf.zeros([784, 10]))
 9 b = tf.Variable(tf.zeros([10]))
10 
11 y = tf.nn.softmax(tf.matmul(x, W) + b)
12 y_ = tf.placeholder(tf.float32, [None, 10])
13 cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), 1))
14 
15 train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
16 tf.global_variables_initializer().run()
17 
18 for i in range(1000):
19     batch_xs, batch_ys = mnist.train.next_batch(100)
20     train_step.run({x: batch_xs, y_: batch_ys})
21 
22 correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
23 accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
24 
25 print(accuracy.eval({x: mnist.test.images, y_: mnist.test.labels}))

  因为我们使用这个模型进行预测,我们肯定是输入一个东西,然后求的我们的输出。
  为了方便找到我们网络中的输入和输出,因此我们可以对我们输入的tensor 进行命名
  命名方式有两种

x = tf.placeholder(tf.float32, [None, 784], name="myInput")#name 就是我们的名字
1 tf.identity(y, name="myOutput")#将identity加到y上面就表示对y进行重命名

  简单方式保存到文件

1 tf.saved_model.simple_save(sess,
2             "./model",
3             inputs={"myInput": x},
4             outputs={"myOutput": y})

  比较复杂的保存模型的方法
    他的优点是:可以自己定义tag,在签名的定义上更加灵活。
  这里说说tag的用途吧。

    一个模型可以包含不同的MetaGraphDef,
  什么时候需要多个MetaGraphDef呢?
      也许你想保存图形的CPU版本和GPU版本,或者你想区分训练和发布版本。
    这个时候tag就可以用来区分不同的MetaGraphDef,加载的时候能够根据tag来加载模型的不同计算图。

  在simple_save方法中,系统会给一个默认的tag: “serve”,也可以用tag_constants.SERVING这个常量。

1 builder = tf.saved_model.builder.SavedModelBuilder("./model")
2 
3 signature = predict_signature_def(inputs={'myInput': x},
4                                   outputs={'myOutput': y})
5 builder.add_meta_graph_and_variables(sess=sess,
6                                      tags=[tag_constants.SERVING],
7                                      signature_def_map={'predict': signature})
8 builder.save()

模型的加载
  tf.saved_model.loader.load中的serve对应的是tag

 1 mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)
 2 
 3 with tf.Session(graph=tf.Graph()) as sess:
 4   tf.saved_model.loader.load(sess, ["serve"], "./model")
 5   graph = tf.get_default_graph()
 6 
 7   input = np.expand_dims(mnist.test.images[0], 0)
 8   x = sess.graph.get_tensor_by_name('myInput:0')
 9   y = sess.graph.get_tensor_by_name('myOutput:0')
10   batch_xs, batch_ys = mnist.test.next_batch(1)
11   scores = sess.run(y,
12            feed_dict={x: batch_xs})
13   print("predict: %d, actual: %d" % (np.argmax(scores, 1), np.argmax(batch_ys, 1)))

 

下面代码是第二份资料
    thanks to https://www.jianshu.com/p/9221fbf52c55

  1 import os
  2 import tensorflow as tf
  3 from tensorflow.python.saved_model import builder as saved_model_builder
  4 from tensorflow.python.saved_model import (signature_constants, signature_def_utils, tag_constants, utils)
  5 
  6 class model():
  7     def __init__(self):
  8         self.a = tf.placeholder(tf.float32, [None])
  9         self.w = tf.Variable(tf.constant(2.0, shape=[1]), name="w")
 10         b = tf.Variable(tf.constant(0.5, shape=[1]), name="b")
 11         self.y = self.a * self.w + b
 12 
 13 #模型保存为ckpt
 14 def save_model():
 15     graph1 = tf.Graph()
 16     with graph1.as_default():
 17         m = model()
 18     with tf.Session(graph=graph1) as session:
 19         session.run(tf.global_variables_initializer())
 20         update = tf.assign(m.w, [10])
 21         session.run(update)
 22         predict_y = session.run(m.y,feed_dict={m.a:[3.0]})
 23         print(predict_y)
 24 
 25         saver = tf.train.Saver()
 26         saver.save(session,"model_pb/model.ckpt")
 27 
 28 
 29 #保存为pb模型
 30 def export_model(session, m):
 31 
 32 
 33    #只需要修改这一段,定义输入输出,其他保持默认即可
 34     model_signature = signature_def_utils.build_signature_def(
 35         inputs={"input": utils.build_tensor_info(m.a)},
 36         outputs={
 37             "output": utils.build_tensor_info(m.y)},
 38 
 39         method_name=signature_constants.PREDICT_METHOD_NAME)
 40 
 41     export_path = "pb_model/1"
 42     if os.path.exists(export_path):
 43         os.system("rm -rf "+ export_path)
 44     print("Export the model to {}".format(export_path))
 45 
 46     try:
 47         legacy_init_op = tf.group(
 48             tf.tables_initializer(), name='legacy_init_op')
 49         builder = saved_model_builder.SavedModelBuilder(export_path)
 50         builder.add_meta_graph_and_variables(
 51             session, [tag_constants.SERVING],
 52             clear_devices=True,
 53             signature_def_map={
 54                 signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
 55                     model_signature,
 56             },
 57             legacy_init_op=legacy_init_op)
 58 
 59         builder.save()
 60     except Exception as e:
 61         print("Fail to export saved model, exception: {}".format(e))
 62 
 63 #加载pb模型
 64 def load_pb():
 65     session = tf.Session(graph=tf.Graph())
 66     model_file_path = "pb_model/1"
 67     meta_graph = tf.saved_model.loader.load(session, [tf.saved_model.tag_constants.SERVING], model_file_path)
 68 
 69     model_graph_signature = list(meta_graph.signature_def.items())[0][1]
 70     output_tensor_names = []
 71     output_op_names = []
 72     for output_item in model_graph_signature.outputs.items():
 73         output_op_name = output_item[0]
 74         output_op_names.append(output_op_name)
 75         output_tensor_name = output_item[1].name
 76         output_tensor_names.append(output_tensor_name)
 77     print("load model finish!")
 78     sentences = {}
 79     # 测试pb模型
 80     for test_x in [[1],[2],[3],[4],[5]]:
 81         sentences["input"] = test_x
 82         feed_dict_map = {}
 83         for input_item in model_graph_signature.inputs.items():
 84             input_op_name = input_item[0]
 85             input_tensor_name = input_item[1].name
 86             feed_dict_map[input_tensor_name] = sentences[input_op_name]
 87         predict_y = session.run(output_tensor_names, feed_dict=feed_dict_map)
 88         print("predict pb y:",predict_y)
 89 
 90 if __name__ == "__main__":
 91 
 92     save_model()
 93 
 94     graph2 = tf.Graph()
 95     with graph2.as_default():
 96         m = model()
 97         saver = tf.train.Saver()
 98     with tf.Session(graph=graph2) as session:
 99         saver.restore(session, "model_pb/model.ckpt") #加载ckpt模型
100         export_model(session, m)
101 
102     load_pb()

 

你可能感兴趣的:(tensorflow 三种模型:ckpt、pb、pb-savemodel)