tensorflow

保存keras模型为pb

2019-06-17  本文已影响0人  夕一啊

keras的模型如果要发布成tfserving,也需要保存为pb格式

# new_model 为加载keras 保存模型的 h5 文件之后的模型
def export_pb(new_model, export_path):

    # 在这里放输入输出
    model_signature = predict_signature_def(inputs={'images': new_model.input},outputs={'scores': new_model.output})

   # 其他地方不用修改
    with K.get_session() as sess:
        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(
                sess, [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))

之后就可以像tf模型一样直接用tfserving了

上一篇下一篇

猜你喜欢

热点阅读