def testCustomMainOp(self):
export_dir = os.path.join(test.get_temp_dir(), "test_main_op")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
v1 = variables.Variable(1, name="v1")
ops.add_to_collection("v", v1)
v2 = variables.Variable(2, name="v2")
ops.add_to_collection("v", v2)
# Initialize another variable `v3` to 42.
v3 = variables.Variable(42, name="v3")
ops.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the main_op.
with ops.control_dependencies([main_op.main_op()]):
add_v1_v2 = math_ops.add(v1._ref(), v2._ref())
custom_main_op = control_flow_ops.group(
state_ops.assign(v3, add_v1_v2))
sess.run(custom_main_op)
builder.add_meta_graph_and_variables(sess, ["foo"],
main_op=custom_main_op)
# Save the SavedModel to disk.
builder.save()
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertEqual(1, ops.get_collection("v")[0].eval())
self.assertEqual(2, ops.get_collection("v")[1].eval())
# Evaluates to the sum of the first two variables and assigned as part of
# the main_op, following a restore.
self.assertEqual(3, ops.get_collection("v")[2].eval())
def testCustomMainOp(self):
export_dir = self._get_export_dir("test_main_op")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
v1 = variables.Variable(1, name="v1")
ops.add_to_collection("v", v1)
v2 = variables.Variable(2, name="v2")
ops.add_to_collection("v", v2)
# Initialize another variable `v3` to 42.
v3 = variables.Variable(42, name="v3")
ops.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the main_op.
with ops.control_dependencies([main_op.main_op()]):
add_v1_v2 = math_ops.add(v1._ref(), v2._ref())
custom_main_op = control_flow_ops.group(state_ops.assign(v3, add_v1_v2))
sess.run(custom_main_op)
builder.add_meta_graph_and_variables(
sess, ["foo"], main_op=custom_main_op)
# Save the SavedModel to disk.
builder.save()
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertEqual(1, ops.get_collection("v")[0].eval())
self.assertEqual(2, ops.get_collection("v")[1].eval())
# Evaluates to the sum of the first two variables and assigned as part of
# the main_op, following a restore.
self.assertEqual(3, ops.get_collection("v")[2].eval())
def testCustomMainOp(self):
export_dir = os.path.join(tf.test.get_temp_dir(), "test_main_op")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=tf.Graph()) as sess:
# Add `v1` and `v2` variables to the graph.
v1 = tf.Variable(1, name="v1")
tf.add_to_collection("v", v1)
v2 = tf.Variable(2, name="v2")
tf.add_to_collection("v", v2)
# Initialize another variable `v3` to 42.
v3 = tf.Variable(42, name="v3", trainable=False, collections=[])
tf.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the main_op.
assign_v3 = tf.assign(v3, tf.add(v1, v2))
custom_main_op = tf.group(main_op.main_op(), assign_v3)
sess.run(tf.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["foo"],
main_op=custom_main_op)
# Save the SavedModel to disk.
builder.save()
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertEqual(1, tf.get_collection("v")[0].eval())
self.assertEqual(2, tf.get_collection("v")[1].eval())
# Evaluates to the sum of the first two variables and assigned as part of
# the main_op, following a restore.
self.assertEqual(3, tf.get_collection("v")[2].eval())
def _generate_saved_model_for_half_plus_two(export_dir,
as_text=False,
use_main_op=False):
"""Generates SavedModel for half plus two.
Args:
export_dir: The directory to which the SavedModel should be written.
as_text: Writes the SavedModel protocol buffer in text format to disk.
use_main_op: Whether to supply a main op during SavedModel build time.
"""
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
with tf.Session(graph=tf.Graph()) as sess:
# Set up the model parameters as variables to exercise variable loading
# functionality upon restore.
a = tf.Variable(0.5, name="a")
b = tf.Variable(2.0, name="b")
c = tf.Variable(3.0, name="c")
# Create a placeholder for serialized tensorflow.Example messages to be fed.
serialized_tf_example = tf.placeholder(tf.string, name="tf_example")
# Parse the tensorflow.Example looking for a feature named "x" with a single
# floating point value.
feature_configs = {
"x": tf.FixedLenFeature(
[1], dtype=tf.float32),
"x2": tf.FixedLenFeature(
[1], dtype=tf.float32, default_value=[0.0])
}
tf_example = tf.parse_example(serialized_tf_example, feature_configs)
# Use tf.identity() to assign name
x = tf.identity(tf_example["x"], name="x")
y = tf.add(tf.multiply(a, x), b, name="y")
y2 = tf.add(tf.multiply(a, x), c, name="y2")
x2 = tf.identity(tf_example["x2"], name="x2")
y3 = tf.add(tf.multiply(a, x2), c, name="y3")
# Create an assets file that can be saved and restored as part of the
# SavedModel.
original_assets_directory = "/tmp/original/export/assets"
original_assets_filename = "foo.txt"
original_assets_filepath = _write_assets(original_assets_directory,
original_assets_filename)
# Set up the assets collection.
assets_filepath = tf.constant(original_assets_filepath)
tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, assets_filepath)
filename_tensor = tf.Variable(
original_assets_filename,
name="filename_tensor",
trainable=False,
collections=[])
assign_filename_op = filename_tensor.assign(original_assets_filename)
# Set up the signature for Predict with input and output tensor
# specification.
predict_input_tensor = tf.saved_model.utils.build_tensor_info(x)
predict_signature_inputs = {"x": predict_input_tensor}
predict_output_tensor = tf.saved_model.utils.build_tensor_info(y)
predict_signature_outputs = {"y": predict_output_tensor}
predict_signature_def = (
tf.saved_model.signature_def_utils.build_signature_def(
predict_signature_inputs, predict_signature_outputs,
tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
# Initialize all variables and then save the SavedModel.
sess.run(tf.global_variables_initializer())
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
"classify_x2_to_y3":
_build_classification_signature(x2, y3),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
if use_main_op:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(main_op.main_op(), assign_filename_op))
else:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=tf.group(assign_filename_op))
builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir,
as_text=False,
use_main_op=False):
"""Generates SavedModel for half plus two.
Args:
export_dir: The directory to which the SavedModel should be written.
as_text: Writes the SavedModel protocol buffer in text format to disk.
use_main_op: Whether to supply a main op during SavedModel build time.
"""
builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
with tf.Session(graph=tf.Graph()) as sess:
# Set up the model parameters as variables to exercise variable loading
# functionality upon restore.
a = tf.Variable(0.5, name="a")
b = tf.Variable(2.0, name="b")
c = tf.Variable(3.0, name="c")
# Create a placeholder for serialized tensorflow.Example messages to be fed.
serialized_tf_example = tf.placeholder(tf.string, name="tf_example")
# Parse the tensorflow.Example looking for a feature named "x" with a single
# floating point value.
feature_configs = {
"x": tf.FixedLenFeature([1], dtype=tf.float32),
"x2": tf.FixedLenFeature([1],
dtype=tf.float32,
default_value=[0.0])
}
tf_example = tf.parse_example(serialized_tf_example, feature_configs)
# Use tf.identity() to assign name
x = tf.identity(tf_example["x"], name="x")
y = tf.add(tf.multiply(a, x), b, name="y")
y2 = tf.add(tf.multiply(a, x), c, name="y2")
x2 = tf.identity(tf_example["x2"], name="x2")
y3 = tf.add(tf.multiply(a, x2), c, name="y3")
# Create an assets file that can be saved and restored as part of the
# SavedModel.
original_assets_directory = "/tmp/original/export/assets"
original_assets_filename = "foo.txt"
original_assets_filepath = _write_assets(original_assets_directory,
original_assets_filename)
# Set up the assets collection.
assets_filepath = tf.constant(original_assets_filepath)
tf.add_to_collection(tf.GraphKeys.ASSET_FILEPATHS, assets_filepath)
filename_tensor = tf.Variable(original_assets_filename,
name="filename_tensor",
trainable=False,
collections=[])
assign_filename_op = filename_tensor.assign(original_assets_filename)
# Set up the signature for Predict with input and output tensor
# specification.
predict_input_tensor = tf.saved_model.utils.build_tensor_info(x)
predict_signature_inputs = {"x": predict_input_tensor}
predict_output_tensor = tf.saved_model.utils.build_tensor_info(y)
predict_signature_outputs = {"y": predict_output_tensor}
predict_signature_def = (
tf.saved_model.signature_def_utils.build_signature_def(
predict_signature_inputs, predict_signature_outputs,
tf.saved_model.signature_constants.PREDICT_METHOD_NAME))
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
# Initialize all variables and then save the SavedModel.
sess.run(tf.global_variables_initializer())
signature_def_map = {
"regress_x_to_y":
_build_regression_signature(serialized_tf_example, y),
"regress_x_to_y2":
_build_regression_signature(serialized_tf_example, y2),
"regress_x2_to_y3":
_build_regression_signature(x2, y3),
"classify_x_to_y":
_build_classification_signature(serialized_tf_example, y),
"classify_x2_to_y3":
_build_classification_signature(x2, y3),
tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
predict_signature_def
}
if use_main_op:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(
tf.GraphKeys.ASSET_FILEPATHS),
main_op=tf.group(main_op.main_op(), assign_filename_op))
else:
builder.add_meta_graph_and_variables(
sess, [tf.saved_model.tag_constants.SERVING],
signature_def_map=signature_def_map,
assets_collection=tf.get_collection(
tf.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=tf.group(assign_filename_op))
builder.save(as_text)
