Mask-RCNN应用 - 训练生成的.h5模型转换为opencv-dnn或tensorflow可用的pb模型

MaskRCNN入门路径–> Mask-RCNN应用研究方法 - 持续更新中

如有问题或需要指导，请私聊留下联系方式或用手机打开https://m.tb.cn/h.fINaraE?tk=PCzA2jPp4V0进行咨询

本篇文章介绍如何将训练生成的.h5文件转换为pb文件，用于tensorflow再训练、调用或者opencv-dnn调用

本文代码测试环境 tensorflow-1.13.0, keras-2.1.2

前言

1、matterport/Mask_RCNN的训练模型保存方式

• matterport/Mask_RCNN是在MaskRCNN基于Python 3, Keras, and TensorFlow的实现，参照源码以及训练过程可以发现，最终生成的.h5训练文件实际上是keras生成的训练文件
• keras对于模型的保存有两种方式：model.save_weights()以及model.save()两种方式。model.save_weights()只保存权重而不保存网络，model.save()保存完整的网络模型与权重
• 由于matterport/Mask_RCNN中使用了自己定义的Lamba层[点击可查看]，并不是完全基于keras定义的网络，因此matterport/Mask_RCNN无法使用model.save()保存模型，选用model.save_weights()[点击可查看]保存模型，这就意味着保存的模型只有权重而没有网络模型，无法通过通常的keras模型转换为pb模型方法进行转换
• matterport/Mask_RCNN只保存权重不保存模型的问题通过加载网络要先进行config加载模型的网络也可以看出

2、keras通常保存模型的方法

• keras保存模型具有多种方法

3、关于keras中引入Lamba函数问题

• keras里建立网络模型不止可以使用keras已经定义好的网络层，同样可以通过Lamba函数加入自己定义的网络层
• 引入Lamba之后将无法通过2中所述方法进行保存

转换的具体步骤

以coco.py为示例

1、加载模型的config文件

• 在使用本文所述的MasKRCNN模型时，无论train、evalute或是test，在进入网络之前都必须加载config文件，这一步在转换时不可缺少

from mrcnn.config import Config

class CocoConfig(Config):
    """Configuration for training on MS COCO.
    Derives from the base Config class and overrides values specific
    to the COCO dataset.
    """
    # Give the configuration a recognizable name
    NAME = "coco"

    # We use a GPU with 12GB memory, which can fit two images.
    # Adjust down if you use a smaller GPU.
    IMAGES_PER_GPU = 2

    # Uncomment to train on 8 GPUs (default is 1)
    # GPU_COUNT = 8

    # Number of classes (including background)
    NUM_CLASSES = 1 + 80  # COCO has 80 classes

config = CocoConfig()

2、生成MaskRCNN模型

• 与上一步介绍的一样，这一步时MasRCNN无论train、evalute或是test必须的步骤，旨在建立MaskRCNN的model

model = modellib.MaskRCNN(mode='train', config=config,
                              model_dir=model_dir)

3、加载训练好的.h5文件生成keras_model

• 这一步结束后，将生成完整的keras_model

# model_path is the absolute path saving the .h5 model
keras_model = model.load_weights(model_path, by_name=True)

4、通过tensorflow进行转换

out_prefix = "output_"
for i in range(len(keras_model.outputs)):
    out_nodes.append(out_prefix + str(i + 1))
    tf.identity(keras_model.output[i], out_prefix + str(i + 1))
model_name = os.path.splitext(os.path.split(model_path)[-1])[0] + '_' + mode + '.pb'
sess = K.get_session()
init_graph = sess.graph.as_graph_def()
main_graph = graph_util.convert_variables_to_constants(sess, init_graph, out_nodes)
graph_io.write_graph(main_graph, output_dir, name=model_name, as_text=False)

至此可以通过Netron查看生成的pb模型

如需完整代码，请私信留言