Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Tensorflow serve on Sagemaker
def export_h5_to_pb(path_to_h5, export_path):
# Set the learning phase to Test since the model is already trained.
K.set_learning_phase(0)
# Load the Keras model
keras_model = load_model(path_to_h5)
# Build the Protocol Buffer SavedModel at 'export_path'
builder = saved_model_builder.SavedModelBuilder(export_path)
# Create prediction signature to be used by TensorFlow Serving Predict API
signature = predict_signature_def(inputs={"images": keras_model.input},
outputs={"scores": keras_model.output})
with K.get_session() as sess:
# Save the meta graph and the variables
builder.add_meta_graph_and_variables(sess=sess, tags=[tag_constants.SERVING],
signature_def_map={"serving_default": signature})
builder.save()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment