def testSignatureDefs(self):
export_dir = os.path.join(
compat.as_bytes(tf.test.get_temp_dir()),
compat.as_bytes("signature_defs"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable and a single entry in the signature def map.
# SavedModel is invoked to add with weights.
with self.test_session(graph=tf.Graph()) as sess:
v = tf.Variable(42, name="v")
sess.run(tf.initialize_all_variables())
self.assertEqual(42, v.eval())
# Build and populate an empty SignatureDef for testing.
foo_signature = utils.build_signature_def(dict(), dict(), "foo")
builder.add_meta_graph_and_variables(
sess, ["foo"], signature_def_map={"foo_key": foo_signature})
# Graph with the same single variable and multiple entries in the signature
# def map. No weights are saved by SavedModel.
with self.test_session(graph=tf.Graph()) as sess:
v = tf.Variable(43, name="v")
sess.run(tf.initialize_all_variables())
self.assertEqual(43, v.eval())
# Build and populate a different SignatureDef for testing.
bar_signature = utils.build_signature_def(dict(), dict(), "bar")
# Also, build a different SignatureDef corresponding to "foo_key" defined
# in the previous graph.
foo_new_signature = utils.build_signature_def(dict(), dict(), "foo_new")
builder.add_meta_graph(
["bar"],
signature_def_map={"bar_key": bar_signature,
"foo_key": foo_new_signature})
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo". The single entry in the SignatureDef map
# corresponding to "foo_key" should exist.
with self.test_session(graph=tf.Graph()) as sess:
foo_graph = loader.load(sess, ["foo"], export_dir)
self.assertEqual(42, tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
foo_signature = foo_graph.signature_def
self.assertEqual(len(foo_signature), 1)
self.assertEqual("foo", foo_signature["foo_key"].method_name)
# Restore the graph with tag "bar". The SignatureDef map should have two
# entries. One corresponding to "bar_key" and another corresponding to the
# new value of "foo_key".
with self.test_session(graph=tf.Graph()) as sess:
bar_graph = loader.load(sess, ["bar"], export_dir)
self.assertEqual(42, tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
bar_signature = bar_graph.signature_def
self.assertEqual(len(bar_signature), 2)
self.assertEqual("bar", bar_signature["bar_key"].method_name)
self.assertEqual("foo_new", bar_signature["foo_key"].method_name)
def testSignatureDefs(self):
export_dir = os.path.join(tf.test.get_temp_dir(),
"test_signature_defs")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable and a single entry in the signature def map.
# SavedModel is invoked to add with weights.
with self.test_session(graph=tf.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Build and populate an empty SignatureDef for testing.
foo_signature = utils.build_signature_def(dict(), dict(), "foo")
builder.add_meta_graph_and_variables(
sess, ["foo"], signature_def_map={"foo_key": foo_signature})
# Graph with the same single variable and multiple entries in the signature
# def map. No weights are saved by SavedModel.
with self.test_session(graph=tf.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
# Build and populate a different SignatureDef for testing.
bar_signature = utils.build_signature_def(dict(), dict(), "bar")
# Also, build a different SignatureDef corresponding to "foo_key" defined
# in the previous graph.
foo_new_signature = utils.build_signature_def(
dict(), dict(), "foo_new")
builder.add_meta_graph(["bar"],
signature_def_map={
"bar_key": bar_signature,
"foo_key": foo_new_signature
})
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo". The single entry in the SignatureDef map
# corresponding to "foo_key" should exist.
with self.test_session(graph=tf.Graph()) as sess:
foo_graph = loader.load(sess, ["foo"], export_dir)
self.assertEqual(
42,
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)[0].eval())
foo_signature = foo_graph.signature_def
self.assertEqual(len(foo_signature), 1)
self.assertEqual("foo", foo_signature["foo_key"].method_name)
# Restore the graph with tag "bar". The SignatureDef map should have two
# entries. One corresponding to "bar_key" and another corresponding to the
# new value of "foo_key".
with self.test_session(graph=tf.Graph()) as sess:
bar_graph = loader.load(sess, ["bar"], export_dir)
self.assertEqual(
42,
tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)[0].eval())
bar_signature = bar_graph.signature_def
self.assertEqual(len(bar_signature), 2)
self.assertEqual("bar", bar_signature["bar_key"].method_name)
self.assertEqual("foo_new", bar_signature["foo_key"].method_name)
def testBuildSignatureDef(self):
x = tf.placeholder(tf.float32, 1, name="x")
x_tensor_info = utils.build_tensor_info(x)
inputs = dict()
inputs["foo-input"] = x_tensor_info
y = tf.placeholder(tf.float32, name="y")
y_tensor_info = utils.build_tensor_info(y)
outputs = dict()
outputs["foo-output"] = y_tensor_info
signature_def = utils.build_signature_def(inputs, outputs,
"foo-method-name")
self.assertEqual("foo-method-name", signature_def.method_name)
# Check inputs in signature def.
self.assertEqual(1, len(signature_def.inputs))
x_tensor_info_actual = signature_def.inputs["foo-input"]
self.assertEqual("x:0", x_tensor_info_actual.name)
self.assertEqual(types_pb2.DT_FLOAT, x_tensor_info_actual.dtype)
self.assertEqual(1, len(x_tensor_info_actual.tensor_shape.dim))
self.assertEqual(1, x_tensor_info_actual.tensor_shape.dim[0].size)
# Check outputs in signature def.
self.assertEqual(1, len(signature_def.outputs))
y_tensor_info_actual = signature_def.outputs["foo-output"]
self.assertEqual("y:0", y_tensor_info_actual.name)
self.assertEqual(types_pb2.DT_FLOAT, y_tensor_info_actual.dtype)
self.assertEqual(0, len(y_tensor_info_actual.tensor_shape.dim))
def testBuildSignatureDef(self):
x = tf.placeholder(tf.float32, 1, name="x")
x_tensor_info = utils.build_tensor_info(x)
inputs = dict()
inputs["foo-input"] = x_tensor_info
y = tf.placeholder(tf.float32, name="y")
y_tensor_info = utils.build_tensor_info(y)
outputs = dict()
outputs["foo-output"] = y_tensor_info
signature_def = utils.build_signature_def(inputs, outputs,
"foo-method-name")
self.assertEqual("foo-method-name", signature_def.method_name)
# Check inputs in signature def.
self.assertEqual(1, len(signature_def.inputs))
x_tensor_info_actual = signature_def.inputs["foo-input"]
self.assertEqual("x:0", x_tensor_info_actual.name)
self.assertEqual(types_pb2.DT_FLOAT, x_tensor_info_actual.dtype)
self.assertEqual(1, len(x_tensor_info_actual.tensor_shape.dim))
self.assertEqual(1, x_tensor_info_actual.tensor_shape.dim[0].size)
# Check outputs in signature def.
self.assertEqual(1, len(signature_def.outputs))
y_tensor_info_actual = signature_def.outputs["foo-output"]
self.assertEqual("y:0", y_tensor_info_actual.name)
self.assertEqual(types_pb2.DT_FLOAT, y_tensor_info_actual.dtype)
self.assertEqual(0, len(y_tensor_info_actual.tensor_shape.dim))
def _generate_saved_model_for_half_plus_two(export_dir, as_text=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.
"""
builder = 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")
# 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),}
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.mul(a, x), b, name="y")
# 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)
# Set up the signature for regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = serialized_tf_example.name
signature_inputs = {signature_constants.REGRESS_INPUTS: input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = tf.identity(y).name
signature_outputs = {signature_constants.REGRESS_OUTPUTS: output_tensor}
signature_def = utils.build_signature_def(
signature_inputs, signature_outputs,
signature_constants.REGRESS_METHOD_NAME)
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [constants.TAG_SERVING],
signature_def_map={
signature_constants.REGRESS_METHOD_NAME:
signature_def
},
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS))
builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir, as_text=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.
"""
builder = 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")
# 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),
}
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.mul(a, x), b, name="y")
# Set up the signature for regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = serialized_tf_example.name
signature_inputs = {signature_constants.REGRESS_INPUTS: input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = tf.identity(y).name
signature_outputs = {
signature_constants.REGRESS_OUTPUTS: output_tensor
}
signature_def = utils.build_signature_def(
signature_inputs, signature_outputs,
signature_constants.REGRESS_METHOD_NAME)
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [constants.TAG_SERVING],
signature_def_map={
signature_constants.REGRESS_METHOD_NAME: signature_def
})
builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir, as_text=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.
"""
builder = 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")
# 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),}
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.mul(a, x), b, name="y")
# Set up the signature for regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = serialized_tf_example.name
signature_inputs = {"input": input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = tf.identity(y).name
signature_outputs = {"output": output_tensor}
signature_def = utils.build_signature_def(signature_inputs,
signature_outputs, "regression")
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [constants.TAG_SERVING],
signature_def_map={"regression": signature_def})
builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir, as_text=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.
"""
builder = 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")
# Set up placeholders.
x = tf.placeholder(tf.float32, name="x")
y = tf.add(tf.mul(a, x), b, name="y")
# Set up the signature for regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = x.name
signature_inputs = {"input": input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = y.name
signature_outputs = {"output": output_tensor}
signature_def = utils.build_signature_def(signature_inputs,
signature_outputs, "regression")
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [constants.TAG_SERVING],
signature_def_map={"regression": signature_def})
builder.save(as_text)
def _generate_saved_model_for_half_plus_two(export_dir):
"""Generates SavedModel for half plus two.
Args:
export_dir: The directory to which the SavedModel should be written.
"""
builder = 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")
# Set up placeholders.
x = tf.placeholder(tf.float32, name="x")
y = tf.add(tf.mul(a, x), b, name="y")
# Set up the signature for regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = x.name
signature_inputs = {"input": input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = y.name
signature_outputs = {"output": output_tensor}
signature_def = utils.build_signature_def(signature_inputs,
signature_outputs,
"regression")
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [constants.TAG_SERVING],
signature_def_map={"regression": signature_def})
builder.save()
def _generate_saved_model_for_half_plus_two(export_dir, as_text=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.
"""
builder = 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")
# 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),
}
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.mul(a, x), b, name="y")
# 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 regression with input and output tensor
# specification.
input_tensor = meta_graph_pb2.TensorInfo()
input_tensor.name = serialized_tf_example.name
signature_inputs = {signature_constants.REGRESS_INPUTS: input_tensor}
output_tensor = meta_graph_pb2.TensorInfo()
output_tensor.name = tf.identity(y).name
signature_outputs = {
signature_constants.REGRESS_OUTPUTS: output_tensor
}
signature_def = utils.build_signature_def(
signature_inputs, signature_outputs,
signature_constants.REGRESS_METHOD_NAME)
# Initialize all variables and then save the SavedModel.
sess.run(tf.initialize_all_variables())
builder.add_meta_graph_and_variables(
sess, [tag_constants.SERVING],
signature_def_map={
signature_constants.REGRESS_METHOD_NAME: signature_def
},
assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
legacy_init_op=tf.group(assign_filename_op))
builder.save(as_text)
def add_graph_and_variables(self,
input_tensors,
output_tensors,
assets_collection=None,
legacy_init_op=None,
main_op=None):
"""添加当前训练模型的 meta graph 和参数。
Args:
input_tensors: 导出模型的输入的别名和 Tensors 之间的字典。
output_tensors: 导出模型的输出的别名和 Tensors 之间的字。，
assets_collection: 附加资产文件列表，可选。
资产文件会在模型导出和导入时被当作模型的一部分进行处理。
资产文件主要应用场景：训练模型的某些操作需要外部附加文件进行初始化等。
在导出模型的时候，资产文件会被拷贝到模型导出路径的 assets 目录下。
legacy_init_op: 在导出模型被加载要被执行的初始化操作，可选。
main_op: 导出模型在被加载时执行的操作，可选。
"""
# Set up the signature for input and output tensorflow specification.
signature_inputs = {}
for (alias_name, input_tensor) in input_tensors.items():
input_tensor_info = meta_graph_pb2.TensorInfo()
input_tensor_info.name = input_tensor.name
signature_inputs[alias_name] = input_tensor_info
signature_outputs = {}
for (alias_name, output_tensor) in output_tensors.items():
output_tensor_info = meta_graph_pb2.TensorInfo()
output_tensor_info.name = output_tensor.name
signature_outputs[alias_name] = output_tensor_info
signature_def = utils.build_signature_def(
signature_inputs, signature_outputs,
caicloud_constants.MODEL_METHOD_NAME)
signature_def_map = {
caicloud_constants.MODEL_METHOD_NAME: signature_def
}
# Save asset files and write them to disk, if any.
self._save_and_write_assets(assets_collection)
if main_op is None:
# Add legacy init op to the SavedModel.
self._maybe_add_legacy_init_op(legacy_init_op)
else:
self._add_main_op(main_op)
# Initialize a saver to generate a sharded output for all variables in the
# current scope.
self._saver = tf_saver.Saver(variables.global_variables(),
sharded=True,
write_version=saver_pb2.SaverDef.V2)
# Export the meta graph def.
meta_graph_def = self._saver.export_meta_graph(clear_devices=True)
# Tag the meta graph def and add it to the SavedModel.
self._tag_and_add_meta_graph(meta_graph_def,
[caicloud_constants.MODEL_TAG],
signature_def_map)
self._has_added_graph_and_variables = True
