美文网首页tensorflow
保存keras模型为pb

保存keras模型为pb

作者: 夕一啊 | 来源:发表于2019-06-17 16:32 被阅读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了

    相关文章

      网友评论

        本文标题:保存keras模型为pb

        本文链接:https://www.haomeiwen.com/subject/chfzfctx.html