Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
tensorflow saved_model, `tags` is the `key` to the metagraph
import tensorflow as tf
import numpy as np
# graph
tf.reset_default_graph()
x = tf.placeholder(tf.float32, [], name='x')
y = tf.multiply(x, 2, name='y')
z = tf.add(x, 1, name='z')
# signature
model_inputs = tf.saved_model.utils.build_tensor_info(x)
model_outputs_y = tf.saved_model.utils.build_tensor_info(y)
model_outputs_z = tf.saved_model.utils.build_tensor_info(z)
model_xy_sig = tf.saved_model.signature_def_utils.build_signature_def(
inputs={
'x':
model_inputs
},
outputs={
'y':
model_outputs_y
},
method_name='xy')
model_xz_sig = tf.saved_model.signature_def_utils.build_signature_def(
inputs={
'x':
model_inputs
},
outputs={
'z':
model_outputs_z
},
method_name='xz')
# export
EXPORT_DIR = './exported_model'
if tf.gfile.Exists(EXPORT_DIR):
tf.gfile.DeleteRecursively(EXPORT_DIR)
builder = tf.saved_model.builder.SavedModelBuilder(EXPORT_DIR)
with tf.Session() as sess:
builder.add_meta_graph_and_variables(
sess,
tags=[tf.saved_model.tag_constants.SERVING, 'xy', 'xz'],
signature_def_map={
'signature-xy': model_xy_sig,
# 'signature-xz': model_xz_sig
}
)
# builder.add_meta_graph(
# sess,
# tags=[tf.saved_model.tag_constants.SERVING, 'xz'],
# signature_def_map={
# 'signature-xz': model_xz_sig
# }
# )
builder.save()
# load
tf.reset_default_graph()
with tf.Session() as sess:
gd = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING, 'xy', 'xz'], EXPORT_DIR)
# tf.saved_model.loader.load(sess, ['xt', tf.saved_model.tag_constants.SERVING, 'xz'], EXPORT_DIR)
g = tf.get_default_graph()
print(g.get_operations())
x = g.get_tensor_by_name('x:0')
y = g.get_tensor_by_name('y:0')
z = g.get_tensor_by_name('z:0')
print(sess.run([y, z], feed_dict={x: 100}))
# check loaded graph definition
sig1 = gd.signature_def['signature-xy']
sig1.inputs
sig1.outputs
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment