Last active
November 5, 2022 17:10
-
-
Save kristijanbartol/7bef5cdd1f0c58c894bb492f92bf9937 to your computer and use it in GitHub Desktop.
Add serving metadata tag to a custom saved TensorFlow model.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import tensorflow as tf | |
# Saved TF model path | |
model_path = 'saved_model.pb' | |
# Directory to export new model | |
target_dir = 'saved_model' | |
with tf.gfile.FastGFile(model_path, 'rb') as f: | |
graph_def = tf.GraphDef() | |
graph_def.ParseFromString(f.read()) | |
tf.import_graph_def(graph_def, name='') | |
# Input tensor name | |
input_name = 'image' | |
# Output tensor name | |
output_name = 'Openpose/concat_stage7' | |
with tf.Session() as sess: | |
model_input = tf.saved_model.utils.build_tensor_info(sess.graph.get_tensor_by_name(input_name + ':0')) | |
model_output = tf.saved_model.utils.build_tensor_info(sess.graph.get_tensor_by_name(output_name + ':0')) | |
signature_definition = tf.saved_model.signature_def_utils.build_signature_def( | |
inputs={input_name: model_input}, | |
outputs={output_name: model_output}, | |
method_name=tf.saved_model.signature_constants.PREDICT_METHOD_NAME) | |
builder = tf.saved_model.builder.SavedModelBuilder(target_dir) | |
builder.add_meta_graph_and_variables(sess, [tf.saved_model.tag_constants.SERVING], | |
signature_def_map={ | |
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature_definition | |
}, clear_devices=True) | |
builder.save() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment