TensorFlow:将多个pb文件模型合并成一个

由于计算机计算能力的增强和相关AI芯片的产出,深度学习中的神经网络结构也是朝着更深、更宽等方向发展。总之就是相应的网络结构参数越来越多,结构越来越复杂。这样无论对于工作还是学习上,想要训练一个网络变得越来越困难。迁移学习的概念很好的解决了这个问题。在之前的视频场景分类中介绍过,将一张图片输入到一个网络,在具体的分类层的前一层输出的数据可以看成是这张图片的全局特征。我们只需对相应的特征进行处理。因此,我们可以利用前人训练好的模型的输出数据当做自己设计模型的输入,这样只需训练自己的网络模型结构中的参数既可。由此引来了,如何把自己的和别人的网络模型合并成一个网络模型。

结构图graph

        TensorFlow的特点就是图和数据项分离,这样增加了网络的灵活性。我们在编写整个网络结构,就类似于画结构图,最后将数据喂给整个网络。但在读别人的程序中,经常会看到Graph,GraphDef,MetaGraph这三个关键词。同样都是图,这三个却又很大的区别。这篇文章对于三种‘图’介绍的不错。在这里做一下总结,Graph是op和tensor的集合。因此在训练阶段画出的图形都是以这种形式保存的,很好的解释了图和数据的分离。GraphDef是序列化后的图,以这种方式存在,相当于数据烧进了图里面,不再有变量的存在。因此,我们在加载或生成pb文件时经常看到GraphDef的定义。MetaGraph相当于是由计算图和相关的元数据构成,这相比于Graph多了些在训练过程中产生的一些数据(比如说训练迭代次数)。以上是我对着三种图的理解,可能存在理解偏差,如有错误请指出。

由此三种图对应三种读取和存储方法:

tf.train.Saver()/saver.restore()

tf.train.export_meta_graph/import_meta_graph

tf.train.write_graph()/tf.import_graph_def()

       本文主要是对两个pb文件进行合并,因此,可能看到GraphDef存在的形式比较多。首先这里给出一段遍历pb文件中所有节点的程序,因为在生成还是加载pb文件时,需要知道输入和输出的节点名称。

graph_def = tf.GraphDef()
    with open(model_path, "rb") as f:
        graph_def.ParseFromString(f.read())
        for i,n in enumerate(graph_def.node):
			print("node name: %s" % n.name)

模型合并

        该篇文章写的算是不错的。但是,如果对于‘图’的理解不是那么深,则很容易让思维陷入局限。总之,该篇文章很好的介绍了,把第一个模型的输出作为第二个模型输入,最后保存成一个模型。核心代码就直接拿来用了:

with tf.Graph().as_default() as g_combined:
  with tf.Session(graph=g_combined) as sess:

    x = tf.placeholder(tf.string, name="base64_input")

    y, = tf.import_graph_def(g1def, input_map={"input_string:0": x}, return_elements=["DecodeJPGOutput:0"])

    z, = tf.import_graph_def(g2def, input_map={"myInput:0": y}, return_elements=["myOutput:0"])
    tf.identity(z, "myOutput")

    tf.saved_model.simple_save(sess,
              "./modelbase64",
              inputs={"base64_input": x},
              outputs={"myOutput": z})

但是如果是这样的场景该怎么办?比如:第一个模型的输出按照条件的不同有可能是第二个模型的输入,也有可能不是(第二个模型的输入和第一个模型的输入相同);第一个模型多个输出的集合作为第二个模型的输入,等等。另外:tf.saved_model.simple_save对于该代码保存的模型,在用常规pb文件加载方式进行加载时出现错误,后面我会介绍另外一种保存方式。

      其实,陷入的思维局限就是:模型的合并是两个模型的串联,第一个模型的输出必须紧跟着第二个模型。可能有人会说,当然是这样啊,不然我合并这两个模型干嘛。但是,如果在工程化上,如果你的图是两个以上甚至十几个,这样维护起来就太麻烦了。

     真正的模型合并其实并不是几张模型图串联起来,想象成并联的说法也不是那么准确。确切的说法感觉应该是,众多子图构成一张大图。如何控制子图之间的联系,不是在生成大图时决定的,而是在加载大图后。以两个子图合成一个大图为例:

    with com_graph.as_default() as g_combined:
        with tf.Session(graph=g_combined) as sess:
            graph_def_one = load_def1(first_pb_path)
            graph_def_two = load_def2(second_pb_path)
            rgbs = tf.placeholder(dtype=tf.uint8,name="video_images")
            features = tf.placeholder(dtype=tf.float32,name="image_feature")
            #first graph load
            Frame_Features = tf.import_graph_def(graph_def_one, input_map={'DecodeJpeg:0': rgbs[:, :, ::-1]},
                                                  return_elements=['pool_3/_reshape:0'])

            tf.identity(Frame_Features, "pool_3/_reshape:0")
            # second graph load
            z, = tf.import_graph_def(graph_def_two, input_map={"Placeholder:0": features}, return_elements=["model/Mul_2:0"])
            tf.identity(z, "model/Mul_2")
            # freeze combined graph
            g_combined_def = graph_util.convert_variables_to_constants(sess, sess.graph_def, ["model/Mul_2","pca_final_feature"])
            tf.train.write_graph(g_combined_def, out_path, 'my_model.pb', as_text=False)

def load_def2(model_path):
    graph_def1 = tf.GraphDef()
    with open(model_path, "rb") as f:
        graph_def1.ParseFromString(f.read())
        return graph_def1

上面就是将两个模型合成一个模型的代码,可以看到,其实并没有将第一张图的输出当成第二张图的输入。他们彼此没有多大关系,只是每张子图在进行加载的时候,该函数:

def import_graph_def(graph_def,
                     input_map=None,
                     return_elements=None,
                     name=None,
                     op_dict=None,
                     producer_op_list=None):

需要指名input(input_map)和输出return_elements(其实也不是必须有,可以输入为None)。因为是输入的pb文件模型,所以加载的图都是graph_def形式。

通过以上代码就可以将两个模型合并,生成一个总的pb文件。加载和单个pb文件类似。主要是通过

sess.graph.get_tensor_by_name()

获取节点名字。

global graph
    graph = tf.get_default_graph()
    with graph.as_default():

        graph_def = tf.GraphDef.FromString(open(model_path, 'rb').read())
        tf.import_graph_def(graph_def, name="")
        with tf.Session() as sess:
            sess.run(tf.global_variables_initializer())
            input_image_tensor = sess.graph.get_tensor_by_name("video_images:0")
            feature = sess.graph.get_tensor_by_name("pool_3/_reshape:0")
            image_feature = sess.graph.get_tensor_by_name("image_feature:0")
            predict_result = sess.graph.get_tensor_by_name("model/Mul_2:0")
            features = []
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess=sess, coord=coord)
            for rgb in frame_iterator(video_file):
                feature_value, = sess.run(
                    [pca_feature],
                    feed_dict={input_image_tensor: rgb[:, :, ::-1]})
                features.append(feature_value)
            features = tf.concat(features, axis=0)
            features = tf.reshape(features,[-1,1024])
            features = tf.expand_dims(features,0)

            predict_result_value, = sess.run(
                [predict_result],
                feed_dict={image_feature: sess.run(features)})
            print('the result is')
            print(predict_result_value)

            coord.request_stop()
            coord.join(threads)

model_path是新pb文件路径。一定要注意的是,新的pb文件的输入,是要按照在生成pb文件的placehold名字,而不是子图的。比如video_images是合成后pb文件的输入节点。其等价于子图节点DecodeJpeg。本文的例子就是,子图一通过frame_iterator()循环函数决定调用次数(假设300次)。然后将这三百次的结果进行append,随后reshape新的矩阵形式的Tensor送入子图二。

总之,在之前合并模型时都是将子图串起来合成,而这次并不是单纯的串联。所以对之前合并模型的认识有一定的局限。在合成模型的程序时,循环子图三百次得出结果再把结果输送到子图二,结果造成生成的pb文件大小超过上限而导致失败。我想这是因为:在合并的时候,由于这个循环而导致是合成的是301个子模型,并不是两个。那么这张图确实是太大。

  • 0
    点赞
  • 17
    收藏
    觉得还不错? 一键收藏
  • 10
    评论
要将训练好的 TensorFlow 模型保存为 .pb 文件,您可以按照以下步骤进行操作: 1. 定义模型结构:在保存模型之前,您需要定义模型的结构,包括输入和输出节点的名称、形状和数据类型。您可以使用 TensorFlow 的高级 API(如 Keras)或自定义模型来定义模型结构。 2. 加载模型权重:将训练好的模型权重加载到定义的模型结构中。这可以通过加载已保存的模型权重文件(如 .h5、.ckpt 等)或通过重新训练模型来实现。 3. 创建 SavedModel:使用 TensorFlow 的 `tf.saved_model.save` 函数将模型保存为 SavedModel 格式。SavedModel 是 TensorFlow 的一种标准模型保存格式,可以包含模型的计算图和变量值。 ```python import tensorflow as tf # 定义和加载模型权重 model = ... # 定义模型结构 model.load_weights('model_weights.h5') # 加载模型权重 # 保存为 SavedModel 格式 tf.saved_model.save(model, 'saved_model') ``` 这将会在指定路径下创建一个名为 `saved_model` 的文件夹,其中包含了模型的计算图和变量值。 4. 导出为 .pb 文件:从 SavedModel 中导出所需的 .pb 文件。可以使用 TensorFlow 的 `tf.compat.v1.graph_util.convert_variables_to_constants` 函数将 SavedModel 的计算图和变量值转换为常量,并保存为 .pb 文件。 ```python from tensorflow.python.framework import graph_util # 加载 SavedModel saved_model_dir = 'saved_model' saved_model = tf.saved_model.load(saved_model_dir) # 将 SavedModel 转换为 .pb 文件 output_pb_file = 'model.pb' graph_def = graph_util.convert_variables_to_constants( saved_model.sess, saved_model.sess.graph_def, ['output_node_name']) with tf.io.gfile.GFile(output_pb_file, 'wb') as f: f.write(graph_def.SerializeToString()) ``` 将上述代码中的 `'output_node_name'` 替换为模型输出节点的名称。 现在,您应该已经功将训练好的 TensorFlow 模型保存为 .pb 文件。请注意,这只是一个基本示例,具体的实现细节可能因您的模型结构和需求而有所不同。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 10
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值