保存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了