def testLegacyInitOp(self):
export_dir = self._get_export_dir("test_legacy_init_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", trainable=False, collections=[])
ops.add_to_collection("v", v3)
# Set up an assignment op to be run as part of the legacy_init_op.
assign_v3 = state_ops.assign(v3, math_ops.add(v1, v2))
legacy_init_op = control_flow_ops.group(assign_v3, name="legacy_init_op")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(
sess, ["foo"], legacy_init_op=legacy_init_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 legacy_init_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 testGraphWithoutVariables(self):
export_dir = self._get_export_dir("test_graph_has_variables")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with no variables.
with self.test_session(graph=ops.Graph()) as sess:
constant_5_name = constant_op.constant(5.0).name
builder.add_meta_graph_and_variables(sess, ["foo"])
# Second graph with no variables
with self.test_session(graph=ops.Graph()) as sess:
constant_6_name = constant_op.constant(6.0).name
builder.add_meta_graph(["bar"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo".
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
# Read the constant a from the graph.
a = ops.get_default_graph().get_tensor_by_name(constant_5_name)
b = constant_op.constant(6.0)
c = a * b
self.assertEqual(30.0, sess.run(c))
# Restore the graph with tag "bar".
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
# Read the constant a from the graph.
a = ops.get_default_graph().get_tensor_by_name(constant_6_name)
b = constant_op.constant(5.0)
c = a * b
self.assertEqual(30.0, sess.run(c))
def testTrainOpGroup(self):
export_dir = self._get_export_dir("test_train_op_group")
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)
sess.run(variables.global_variables_initializer())
train_op = control_flow_ops.group()
sess.run(train_op)
# TODO(karmel): remove explicit call when in the public method.
builder._add_train_op(train_op)
builder.add_meta_graph_and_variables(sess, ["foo"])
# 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())
self.assertIsInstance(
ops.get_collection(constants.TRAIN_OP_KEY)[0], ops.Operation)
def testTrainOpAfterVariables(self):
export_dir = self._get_export_dir("test_train_op_after_variables")
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)
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["pre_foo"])
train_op = state_ops.assign_add(v1, v2)
sess.run(train_op)
# TODO(karmel): remove explicit call when in the public method.
builder._add_train_op(train_op)
builder.add_meta_graph(["foo"])
# Save the SavedModel to disk.
builder.save()
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertIsInstance(
ops.get_collection(constants.TRAIN_OP_KEY)[0], ops.Tensor)
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["pre_foo"], export_dir)
self.assertFalse(ops.get_collection(constants.TRAIN_OP_KEY))
def testSaveAsText(self):
export_dir = os.path.join(
compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("astext"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable. SavedModel 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())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Graph with the same single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
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())
builder.add_meta_graph(["bar"])
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Restore the graph with tag "foo", whose variables were saved.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertEqual(42, tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
# Restore the graph with tag "bar", whose variables were not saved.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
self.assertEqual(42, tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
def testCustomSaveable(self):
export_dir = self._get_export_dir("custom_saveable")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with session.Session(
graph=ops.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
# CheckpointedOp is a key-value table that can be saved across sessions.
# The table register itself in SAVEABLE_OBJECTS collection.
v1 = saver_test_utils.CheckpointedOp(name="v1")
variables.global_variables_initializer().run()
v1.insert("k1", 3.0).run()
# Once the table is restored, we can access it through this reference.
ops.add_to_collection("table_ref", v1.table_ref)
builder.add_meta_graph_and_variables(sess, ["foo"])
# Save the SavedModel to disk.
builder.save()
with session.Session(
graph=ops.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
loader.load(sess, ["foo"], export_dir)
# Instantiate a wrapper object from the checkpointed reference.
v1 = saver_test_utils.CheckpointedOp(
name="v1", table_ref=ops.get_collection("table_ref")[0])
self.assertEqual(b"k1", v1.keys().eval())
self.assertEqual(3.0, v1.values().eval())
def testSaveAsText(self):
export_dir = self._get_export_dir("test_astext")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable. SavedModel invoked to:
# - add with weights.
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess, ["foo"])
# Graph with the same single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
builder.add_meta_graph(["bar"])
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Restore the graph with tag "foo", whose variables were saved.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
self.assertEqual(
42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Restore the graph with tag "bar", whose variables were not saved.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
self.assertEqual(
42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
def export_fn(estimator, export_dir, checkpoint_path=None, eval_result=None):
"""A wrapper to export to SavedModel, and convert it to other formats."""
result_dir = base_strategy.export(estimator, export_dir,
checkpoint_path,
eval_result)
with ops.Graph().as_default() as graph:
with tf_session.Session(graph=graph) as sess:
saved_model_loader.load(
sess, [tag_constants.SERVING], result_dir)
# Note: This is GTFlow internal API and might change.
ensemble_model = graph.get_operation_by_name(
"ensemble_model/TreeEnsembleSerialize")
_, dfec_str = sess.run(ensemble_model.outputs)
dtec = tree_config_pb2.DecisionTreeEnsembleConfig()
dtec.ParseFromString(dfec_str)
# Export the result in the same folder as the saved model.
if convert_fn:
convert_fn(dtec, sorted_feature_names,
len(dense_floats),
len(sparse_float_indices),
len(sparse_int_indices), result_dir, eval_result)
feature_importances = _get_feature_importances(
dtec, sorted_feature_names,
len(dense_floats),
len(sparse_float_indices), len(sparse_int_indices))
sorted_by_importance = sorted(
feature_importances.items(), key=lambda x: -x[1])
assets_dir = os.path.join(result_dir, "assets.extra")
gfile.MakeDirs(assets_dir)
with gfile.GFile(os.path.join(assets_dir, "feature_importances"),
"w") as f:
f.write("\n".join("%s, %f" % (k, v) for k, v in sorted_by_importance))
return result_dir
def testSignatureDefs(self):
export_dir = self._get_export_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=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Build and populate an empty SignatureDef for testing.
foo_signature = signature_def_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=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
# Build and populate a different SignatureDef for testing.
bar_signature = signature_def_utils.build_signature_def(dict(),
dict(), "bar")
# Also, build a different SignatureDef corresponding to "foo_key" defined
# in the previous graph.
foo_new_signature = signature_def_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=ops.Graph()) as sess:
foo_graph = loader.load(sess, ["foo"], export_dir)
self.assertEqual(
42, ops.get_collection(ops.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=ops.Graph()) as sess:
bar_graph = loader.load(sess, ["bar"], export_dir)
self.assertEqual(
42, ops.get_collection(ops.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 testVariables(self):
export_dir = os.path.join(
compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("variables"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with two variables. SavedModel invoked to:
# - add with weights.
with self.test_session(graph=tf.Graph()) as sess:
v1 = tf.Variable(1, name="v1")
v2 = tf.Variable(2, name="v2")
sess.run(tf.initialize_all_variables())
self.assertEqual(1, v1.eval())
self.assertEqual(2, v2.eval())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Graph with a single variable (subset of the variables from the previous
# graph whose weights were saved). SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=tf.Graph()) as sess:
v2 = tf.Variable(3, name="v2")
sess.run(tf.initialize_all_variables())
self.assertEqual(3, v2.eval())
builder.add_meta_graph(["bar"])
# Graph with a single variable (disjoint set of variables from the previous
# graph whose weights were saved). SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=tf.Graph()) as sess:
v3 = tf.Variable(4, name="v3")
sess.run(tf.initialize_all_variables())
self.assertEqual(4, v3.eval())
builder.add_meta_graph(["baz"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo", whose variables were saved.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
collection_vars = tf.get_collection(tf.GraphKeys.VARIABLES)
self.assertEqual(len(collection_vars), 2)
self.assertEqual(1, collection_vars[0].eval())
self.assertEqual(2, collection_vars[1].eval())
# Restore the graph with tag "bar", whose variables were not saved. Only the
# subset of the variables added to the graph will be restored with the
# checkpointed value.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
collection_vars = tf.get_collection(tf.GraphKeys.VARIABLES)
self.assertEqual(len(collection_vars), 1)
self.assertEqual(2, collection_vars[0].eval())
# Try restoring the graph with tag "baz", whose variables were not saved.
# Since this graph has a disjoint set of variables from the set that was
# saved, this should raise an error.
with self.test_session(graph=tf.Graph()) as sess:
self.assertRaises(errors.NotFoundError, loader.load, sess, ["baz"],
export_dir)
def testClearExtraneousSavers(self):
export_dir = os.path.join(test.get_temp_dir(),
"test_clear_extraneous_savers")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Create a variable and a Saver.
with ops.Graph().as_default() as graph:
with session.Session(
target="",
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Add two Savers, which should be removed in
# add_meta_graph_and_variables() in favor of the locally added one.
saver1 = tf_saver.Saver()
graph.add_to_collection(ops.GraphKeys.SAVERS, saver1)
saver2 = tf_saver.Saver()
graph.add_to_collection(ops.GraphKeys.SAVERS, saver2)
# Confirm there are two SaverDefs.
savers = graph.get_collection(ops.GraphKeys.SAVERS)
self.assertEqual(2, len(savers))
# Confirm there are two Save and two Restore ops.
save_op_names = set([x.name for x in graph.get_operations()
if x.type == "SaveV2"])
self.assertSetEqual(set(["save/SaveV2", "save_1/SaveV2"]),
save_op_names)
restore_op_names = set([x.name for x in graph.get_operations()
if x.type == "RestoreV2"])
self.assertSetEqual(set(["save/RestoreV2", "save_1/RestoreV2"]),
restore_op_names)
# The SavedModel builder adds its own Saver' for a total of three.
builder.add_meta_graph_and_variables(
sess, [tag_constants.TRAINING], clear_devices=True)
# Save the SavedModel to disk.
builder.save()
# Restore the graph.
with ops.Graph().as_default() as graph:
with self.test_session(graph=graph) as sess:
loader.load(sess, [tag_constants.TRAINING], export_dir)
self.assertEqual(
42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Confirm that the reloaded graph has only one SaverDef.
savers = ops.get_collection(ops.GraphKeys.SAVERS)
self.assertEqual(1, len(savers))
# The reloaded graph should have exactly one Save and one Restore op.
save_op_names = set([x.name for x in graph.get_operations()
if x.type == "SaveV2"])
self.assertSetEqual(set(["save_2/SaveV2"]), save_op_names)
restore_op_names = set([x.name for x in graph.get_operations()
if x.type == "RestoreV2"])
self.assertSetEqual(set(["save_2/RestoreV2"]), restore_op_names)
def testStripDefaultAttrsInconsistentConsumerDefaults(self):
if ops._USE_C_API: return # TODO(skyewm): get this working
export_dir = self._get_export_dir(
"test_strip_default_attrs_no_consumer_defaults")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Add a graph with two float32 variables and a Complex Op composing them
# with strip_default_attrs enabled. This must remove the following
# defaults for the "Complex" Op:
# o "T" : float32. (input type)
# o "Tout" : complex64. (output type)
with session.Session(graph=ops.Graph()) as sess:
real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(
sess, ["foo"], strip_default_attrs=True)
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Update the Op registry to remove defaults for all attrs("T", "Tout") from
# the "Complex" OpDef.
complex_op_def = op_def_registry.get_registered_ops()["Complex"]
original_complex_op_def = op_def_pb2.OpDef()
original_complex_op_def.CopyFrom(complex_op_def)
for attr_def in complex_op_def.attr:
attr_def.ClearField("default_value")
# Loading the SavedModel via the loader must fail because the SavedModel
# does not have any attr values for the "Complex" node and the current
# op registry does not have have any default values for the "Complex" op.
sess = session.Session(graph=ops.Graph())
with self.assertRaisesRegexp(
ValueError,
"Expected one attr with name .*T(out)?.* in name: \"complex\".*"):
loader.load(sess, ["foo"], export_dir)
# Update the Op registry to change the defaults for attr "Tout"
# (complex64 -> complex128).
complex_op_def.CopyFrom(original_complex_op_def)
for attr_def in complex_op_def.attr:
if attr_def.name == "Tout":
attr_def.default_value.type = types_pb2.DT_COMPLEX128
# Loading the SavedModel via the loader must set "Tout" attr_value for the
# "Complex" node according to the latest defaults (complex128). This is
# expected to fail the model import as there is no OpKernel registered to
# handle attrs "T" (float32) and "Tout" (complex128).
sess = session.Session(graph=ops.Graph())
with self.assertRaisesRegexp(
errors.InvalidArgumentError,
".*No OpKernel was registered to support Op \'Complex\' with these "
"attrs..*"):
loader.load(sess, ["foo"], export_dir)
def freeze_saved_model(saved_model_dir, input_arrays, input_shapes,
output_arrays, tag_set, signature_key):
"""Converts a SavedModel to a frozen graph.
Args:
saved_model_dir: SavedModel directory to convert.
input_arrays: List of input tensors to freeze graph with. Uses input arrays
from SignatureDef when none are provided.
input_shapes: Dict of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" : None}).
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided.
tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
analyze. All tags in the tag set must be present.
signature_key: Key identifying SignatureDef containing inputs and outputs.
Returns:
frozen_graph_def: Frozen GraphDef.
in_tensors: List of input tensors for the graph.
out_tensors: List of output tensors for the graph.
Raises:
ValueError:
SavedModel doesn't contain a MetaGraphDef identified by tag_set.
signature_key is not in the MetaGraphDef.
assets/ directory is in the MetaGraphDef.
input_shapes does not match the length of input_arrays.
input_arrays or output_arrays are not valid.
"""
# Read SignatureDef.
meta_graph = _get_meta_graph_def(saved_model_dir, tag_set)
signature_def = _get_signature_def(meta_graph, signature_key)
inputs, outputs = _get_inputs_outputs(signature_def)
# Check SavedModel for assets directory.
collection_def = meta_graph.collection_def
if constants.ASSETS_KEY in collection_def:
raise ValueError("SavedModels with assets/ directory are not supported.")
graph = ops.Graph()
with session.Session(graph=graph) as sess:
loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir)
# Gets input and output tensors.
# TODO(zhixianyan): Use TFLite supported Op list to filter outputs.
in_tensors = _get_tensors(graph, inputs, input_arrays)
out_tensors = _get_tensors(graph, outputs, output_arrays)
set_tensor_shapes(in_tensors, input_shapes)
output_names = [node.split(":")[0] for node in outputs]
frozen_graph_def = tf_graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), output_names)
return frozen_graph_def, in_tensors, out_tensors
def __init__(self,
export_dir,
signature_def_key=None,
signature_def=None,
input_names=None,
output_names=None,
tags=None,
graph=None):
"""Initialize a `CoreEstimatorPredictor`.
Args:
export_dir: a path to a directory containing a `SavedModel`.
signature_def_key: Optional string specifying the signature to use. If
`None`, then `DEFAULT_SERVING_SIGNATURE_DEF_KEY` is used. Only one of
`signature_def_key` and `signature_def` should be specified.
signature_def: A `SignatureDef` proto specifying the inputs and outputs
for prediction. Only one of `signature_def_key` and `signature_def`
should be specified.
input_names: A dictionary mapping strings to `Tensor`s in the `SavedModel`
that represent the input. The keys can be any string of the user's
choosing.
output_names: A dictionary mapping strings to `Tensor`s in the
`SavedModel` that represent the output. The keys can be any string of
the user's choosing.
tags: Optional. Tags that will be used to retrieve the correct
`SignatureDef`. Defaults to `DEFAULT_TAGS`.
graph: Optional. The Tensorflow `graph` in which prediction should be
done.
Raises:
ValueError: If more than one of signature_def_key OR signature_def OR
(input_names AND output_names) is specified.
"""
_check_signature_arguments(
signature_def_key, signature_def, input_names, output_names)
tags = tags or DEFAULT_TAGS
self._graph = graph or ops.Graph()
with self._graph.as_default():
self._session = session.Session()
loader.load(self._session, tags.split(','), export_dir)
if input_names is None:
if signature_def is None:
signature_def = _get_signature_def(signature_def_key, export_dir, tags)
input_names = {k: v.name for k, v in signature_def.inputs.items()}
output_names = {k: v.name for k, v in signature_def.outputs.items()}
self._feed_tensors = {k: self._graph.get_tensor_by_name(v)
for k, v in input_names.items()}
self._fetch_tensors = {k: self._graph.get_tensor_by_name(v)
for k, v in output_names.items()}
def testVariables(self):
export_dir = self._get_export_dir("test_variables")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with two variables. SavedModel invoked to:
# - add with weights.
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v1", 1)
self._init_and_validate_variable(sess, "v2", 2)
builder.add_meta_graph_and_variables(sess, ["foo"])
# Graph with a single variable (subset of the variables from the previous
# graph whose weights were saved). SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v2", 3)
builder.add_meta_graph(["bar"])
# Graph with a single variable (disjoint set of variables from the previous
# graph whose weights were saved). SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v3", 4)
builder.add_meta_graph(["baz"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo", whose variables were saved.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
self.assertEqual(len(collection_vars), 2)
self.assertEqual(1, collection_vars[0].eval())
self.assertEqual(2, collection_vars[1].eval())
# Restore the graph with tag "bar", whose variables were not saved. Only the
# subset of the variables added to the graph will be restored with the
# checkpointed value.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
collection_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
self.assertEqual(len(collection_vars), 1)
self.assertEqual(2, collection_vars[0].eval())
# Try restoring the graph with tag "baz", whose variables were not saved.
# Since this graph has a disjoint set of variables from the set that was
# saved, this should raise an error.
with self.test_session(graph=ops.Graph()) as sess:
self.assertRaises(errors.NotFoundError, loader.load, sess, ["baz"],
export_dir)
def _TestStaticOp(self, use_function_backup):
if not is_tensorrt_enabled():
return
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir3")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir3")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = self._ConvertGraph(
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
maximum_cached_engines=2, # This is noop, added just for testing.
use_function_backup=use_function_backup)
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(
sess,
1,
use_function_backup=use_function_backup,
expect_engine_is_run=True)
# Run with batch size 2, which exceed the max_batch_size, it should try
# to fall back to TF function.
self._TestRun(
sess,
2,
use_function_backup=use_function_backup,
expect_engine_is_run=False)
# Test the output SavedModel
with ops.Graph().as_default():
with self.session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING], output_saved_model_dir)
# Run with batch size 1, the default engine embedded in the graphdef
# will be used.
self._TestRun(
sess,
1,
use_function_backup=use_function_backup,
expect_engine_is_run=True)
# Run with batch size 2, which exceed the max_batch_size, it should try
# to fall back to TF function.
self._TestRun(
sess,
2,
use_function_backup=use_function_backup,
expect_engine_is_run=False)
def testCollections(self):
export_dir = os.path.join(
compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("collections"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable added to a collection. SavedModel invoked to:
# - add with weights.
with self.test_session(graph=tf.Graph()) as sess:
v = tf.Variable(42, name="v")
tf.add_to_collection("foo_vars", v)
sess.run(tf.initialize_all_variables())
self.assertEqual(42, v.eval())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Graph with the same single variable added to a different collection.
# SavedModel invoked to:
# - simply add the model (weights are not updated).
with self.test_session(graph=tf.Graph()) as sess:
v = tf.Variable(43, name="v")
tf.add_to_collection("bar_vars", v)
sess.run(tf.initialize_all_variables())
self.assertEqual(43, v.eval())
builder.add_meta_graph(["bar"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with tag "foo", whose variables were saved. The
# collection 'foo_vars' should contain a single element. The collection
# 'bar_vars' should not be found.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["foo"], export_dir)
collection_foo_vars = tf.get_collection("foo_vars")
self.assertEqual(len(collection_foo_vars), 1)
self.assertEqual(42, collection_foo_vars[0].eval())
self.assertEqual(len(tf.get_collection("bar_vars")), 0)
# Restore the graph with tag "bar", whose variables were not saved. The
# collection-def exported as part of the meta graph def is updated to
# reflect the new collection. The value of the variable in the
# collection-def corresponds to the saved value (from the previous graph
# with tag "foo").
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["bar"], export_dir)
collection_bar_vars = tf.get_collection("bar_vars")
self.assertEqual(len(collection_bar_vars), 1)
self.assertEqual(42, collection_bar_vars[0].eval())
self.assertEqual(len(tf.get_collection("foo_vars")), 0)
def process(self, inputs):
# Create a session for every worker only once. The session is not
# pickleable, so it can't be created at the DoFn constructor.
if not self.session:
self.graph = ops.Graph()
with self.graph.as_default():
self.session = tf.Session()
metagraph_def = loader.load(
self.session, {self.meta_tag}, self.model_dir)
signature_def = metagraph_def.signature_def[self.meta_signature]
# inputs
self.feed_tensors = {
k: self.graph.get_tensor_by_name(v.name)
for k, v in signature_def.inputs.items()
}
# outputs/predictions
self.fetch_tensors = {
k: self.graph.get_tensor_by_name(v.name)
for k, v in signature_def.outputs.items()
}
# Create a feed_dict for a single element.
feed_dict = {
tensor: [inputs[key]]
for key, tensor in self.feed_tensors.items()
if key in inputs
}
results = self.session.run(self.fetch_tensors, feed_dict)
yield {
'id': inputs[self.id_key],
'predictions': results[self.meta_predictions][0].tolist()
}
def local_predict(args):
"""Runs prediction locally."""
sess = session.Session()
_ = loader.load(sess, [tag_constants.SERVING], args.model_dir)
# get the mappings between aliases and tensor names
# for both inputs and outputs
input_alias_map = json.loads(sess.graph.get_collection('inputs')[0])
output_alias_map = json.loads(sess.graph.get_collection('outputs')[0])
aliases, tensor_names = zip(*output_alias_map.items())
for input_file in args.input:
feed_dict = collections.defaultdict(list)
for line in tf_record.tf_record_iterator(input_file):
feed_dict[input_alias_map['examples_bytes']].append(line)
if args.dry_run:
print('Feed data dict %s to graph and fetch %s' % (
feed_dict, tensor_names))
else:
result = sess.run(fetches=tensor_names, feed_dict=feed_dict)
for row in zip(*result):
print(json.dumps(
{name: (value.tolist() if getattr(value, 'tolist', None) else value)
for name, value in zip(aliases, row)}))
def testAssets(self):
export_dir = self._get_export_dir("test_assets")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Build an asset collection.
ignored_filepath = os.path.join(
compat.as_bytes(test.get_temp_dir()), compat.as_bytes("ignored.txt"))
file_io.write_string_to_file(ignored_filepath, "will be ignored")
asset_collection = self._build_asset_collection("hello42.txt",
"foo bar baz",
"asset_file_tensor")
builder.add_meta_graph_and_variables(
sess, ["foo"], assets_collection=asset_collection)
# Save the SavedModel to disk.
builder.save()
with self.test_session(graph=ops.Graph()) as sess:
foo_graph = loader.load(sess, ["foo"], export_dir)
self._validate_asset_collection(export_dir, foo_graph.collection_def,
"hello42.txt", "foo bar baz",
"asset_file_tensor:0")
ignored_asset_path = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes(constants.ASSETS_DIRECTORY),
compat.as_bytes("ignored.txt"))
self.assertFalse(file_io.file_exists(ignored_asset_path))
def testDuplicateAssets(self):
export_dir = self._get_export_dir("test_duplicate_assets")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Build an asset collection with `foo.txt` that has `foo` specific
# content.
asset_collection = self._build_asset_collection("foo.txt", "content_foo",
"asset_file_tensor")
# Add the asset collection as part of the graph with tag "foo".
builder.add_meta_graph_and_variables(
sess, ["foo"], assets_collection=asset_collection)
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Build an asset collection with `foo.txt` that has `bar` specific
# content.
asset_collection = self._build_asset_collection("foo.txt", "content_bar",
"asset_file_tensor")
# Add the asset collection as part of the graph with tag "bar".
builder.add_meta_graph(["bar"], assets_collection=asset_collection)
# Save the SavedModel to disk.
builder.save()
# Check assets restored for graph with tag "foo".
with self.test_session(graph=ops.Graph()) as sess:
foo_graph = loader.load(sess, ["foo"], export_dir)
self._validate_asset_collection(export_dir, foo_graph.collection_def,
"foo.txt", "content_foo",
"asset_file_tensor:0")
# Check assets restored for graph with tag "bar".
with self.test_session(graph=ops.Graph()) as sess:
bar_graph = loader.load(sess, ["bar"], export_dir)
# Validate the assets for `bar` graph. `foo.txt` should contain the
# original contents corresponding to `foo` graph since an asset with the
# same name across multiple graphs is only stored the first time
self._validate_asset_collection(export_dir, bar_graph.collection_def,
"foo.txt", "content_foo",
"asset_file_tensor:0")
def test_scaffold_is_used_for_local_init(self):
tmpdir = tempfile.mkdtemp()
def _model_fn_scaffold(features, labels, mode):
_, _ = features, labels
my_int = variables.Variable(1, name='my_int',
collections=[ops.GraphKeys.LOCAL_VARIABLES])
scores = constant_op.constant([3.])
with ops.control_dependencies(
[variables.local_variables_initializer(),
data_flow_ops.tables_initializer()]):
assign_op = state_ops.assign(my_int, 12345)
# local_initSop must be an Operation, not a Tensor.
custom_local_init_op = control_flow_ops.group(assign_op)
return model_fn_lib.EstimatorSpec(
mode=mode,
predictions=constant_op.constant([[1.]]),
loss=constant_op.constant(0.),
train_op=constant_op.constant(0.),
scaffold=training.Scaffold(local_init_op=custom_local_init_op),
export_outputs={'test': export_output.ClassificationOutput(scores)})
est = estimator.Estimator(model_fn=_model_fn_scaffold)
est.train(dummy_input_fn, steps=1)
feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
# Perform the export.
export_dir_base = os.path.join(
compat.as_bytes(tmpdir), compat.as_bytes('export'))
export_dir = est.export_savedmodel(export_dir_base,
serving_input_receiver_fn)
# Restore, to validate that the custom local_init_op runs.
with ops.Graph().as_default() as graph:
with session.Session(graph=graph) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
my_int = graph.get_tensor_by_name('my_int:0')
my_int_value = sess.run(my_int)
self.assertEqual(12345, my_int_value)
def test_export_savedmodel_with_saveables_proto_roundtrip(self):
tmpdir = tempfile.mkdtemp()
est = estimator.Estimator(
model_fn=_model_fn_with_saveables_for_export_tests)
est.train(input_fn=dummy_input_fn, steps=1)
feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
# Perform the export.
export_dir_base = os.path.join(
compat.as_bytes(tmpdir), compat.as_bytes('export'))
export_dir = est.export_savedmodel(
export_dir_base, serving_input_receiver_fn)
# Check that all the files are in the right places.
self.assertTrue(gfile.Exists(export_dir_base))
self.assertTrue(gfile.Exists(export_dir))
self.assertTrue(gfile.Exists(os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('saved_model.pb'))))
self.assertTrue(gfile.Exists(os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables'))))
self.assertTrue(gfile.Exists(os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.index'))))
self.assertTrue(gfile.Exists(os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.data-00000-of-00001'))))
# Restore, to validate that the export was well-formed.
with ops.Graph().as_default() as graph:
with session.Session(graph=graph) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
graph_ops = [x.name for x in graph.get_operations()]
self.assertTrue('input_example_tensor' in graph_ops)
self.assertTrue('ParseExample/ParseExample' in graph_ops)
self.assertTrue('save/LookupTableImport' in graph_ops)
# Clean up.
gfile.DeleteRecursively(tmpdir)
def testOp(self):
export_dir = os.path.join(
compat.as_bytes(tf.test.get_temp_dir()), compat.as_bytes("op"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
with tf.Session(
graph=tf.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
v1 = tf.Variable(1, name="v1")
with sess.graph.device("/cpu:1"):
v2 = tf.Variable(2, name="v2")
# v3 is an unsaved variable derived from v1 and v2. It is used to
# exercise the ability to run an init op when restoring a graph.
v3 = tf.Variable(1, name="v3", trainable=False, collections=[])
assign_v3 = tf.assign(v3, tf.add(v1, v2))
init_op = tf.group(assign_v3, name="init_op")
tf.add_to_collection("v", v1)
tf.add_to_collection("v", v2)
tf.add_to_collection("v", v3)
tf.add_to_collection("init_op", init_op)
sess.run(tf.initialize_all_variables())
self.assertEqual(1, tf.get_collection("v")[0].eval())
self.assertEqual(2, tf.get_collection("v")[1].eval())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Save the SavedModel to disk.
builder.save()
with tf.Session(
graph=tf.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
loader.load(sess, ["foo"], export_dir)
# Validate variables, run the init op and verify result.
self.assertEqual(1, tf.get_collection("v")[0].eval())
self.assertEqual(2, tf.get_collection("v")[1].eval())
tf.get_collection("init_op")[0].run()
self.assertEqual(3, tf.get_collection("v")[2].eval())
def testVerifySessionGraphUsage(self):
export_dir = self._get_export_dir("test_verify_session_graph_usage")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess, [tag_constants.TRAINING])
# Save the SavedModel to disk.
builder.save()
# Build a session and supply it to the load operation.
sess = session.Session(graph=ops.Graph())
loader.load(sess, [tag_constants.TRAINING], export_dir)
# Check the variable within the scope of the session and its graph.
with sess:
self.assertEqual(
42, ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
def testClearDevices(self):
export_dir = os.path.join(tf.test.get_temp_dir(), "test_clear_devices")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Specify a device and save a variable.
tf.reset_default_graph()
with tf.Session(target="", config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess, [tag_constants.TRAINING], clear_devices=True)
# Save the SavedModel to disk.
builder.save()
# Restore the graph with a single predefined tag whose variables were saved
# without any device information.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, [tag_constants.TRAINING], export_dir)
self.assertEqual(42, tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
def testCreateInferenceGraph_DynamicOp(self):
if not trt_convert.is_tensorrt_enabled():
return
trt_convert.enable_test_value()
tmp_dir = self.get_temp_dir()
input_saved_model_dir = os.path.join(tmp_dir, "in_dir2")
output_saved_model_dir = os.path.join(tmp_dir, "out_dir2")
self._WriteInputSavedModel(input_saved_model_dir)
output_graph_def = trt_convert.create_inference_graph(
None,
None,
max_workspace_size_bytes=TrtConvertTest._TRT_MAX_WORKSPACE_SIZE_BYTES,
is_dynamic_op=True,
maximum_cached_engines=2,
input_saved_model_dir=input_saved_model_dir,
output_saved_model_dir=output_saved_model_dir,
session_config=self._GetConfigProto())
# Test the output GraphDef.
with ops.Graph().as_default():
importer.import_graph_def(output_graph_def, name="")
with self.test_session(config=self._GetConfigProto()) as sess:
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should evict an old engine and create a new one.
self._TestRun(sess, 3, True)
# Test the output SavedModel
with ops.Graph().as_default():
with self.test_session(config=self._GetConfigProto()) as sess:
loader.load(sess, [tag_constants.SERVING], output_saved_model_dir)
# Run with batch size 1, a new engine is created and cached.
self._TestRun(sess, 1, True)
# Run with batch size 2, a new engine is created and cached.
self._TestRun(sess, 2, True)
# Run with batch size 3, since the number of cached engines has reached
# the max, it should evict an old engine and create a new one.
self._TestRun(sess, 3, True)
def testOp(self):
export_dir = self._get_export_dir("test_op")
builder = saved_model_builder.SavedModelBuilder(export_dir)
with session.Session(
graph=ops.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
with sess.graph.device("/cpu:0"):
v1 = variables.Variable(1, name="v1")
with sess.graph.device("/cpu:1"):
v2 = variables.Variable(2, name="v2")
# v3 is an unsaved variable derived from v1 and v2. It is used to
# exercise the ability to run an init op when restoring a graph.
v3 = variables.Variable(1, name="v3", trainable=False, collections=[])
assign_v3 = state_ops.assign(v3, math_ops.add(v1, v2))
init_op = control_flow_ops.group(assign_v3, name="init_op")
ops.add_to_collection("v", v1)
ops.add_to_collection("v", v2)
ops.add_to_collection("v", v3)
ops.add_to_collection("init_op", init_op)
sess.run(variables.global_variables_initializer())
self.assertEqual(1, ops.get_collection("v")[0].eval())
self.assertEqual(2, ops.get_collection("v")[1].eval())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Save the SavedModel to disk.
builder.save()
with session.Session(
graph=ops.Graph(),
config=config_pb2.ConfigProto(device_count={"CPU": 2})) as sess:
loader.load(sess, ["foo"], export_dir)
# Validate variables, run the init op and verify result.
self.assertEqual(1, ops.get_collection("v")[0].eval())
self.assertEqual(2, ops.get_collection("v")[1].eval())
ops.get_collection("init_op")[0].run()
self.assertEqual(3, ops.get_collection("v")[2].eval())
def test_export_savedmodel_with_resource(self):
tmpdir = tempfile.mkdtemp()
est, serving_input_fn = _build_estimator_for_resource_export_test()
export_dir_base = os.path.join(
compat.as_bytes(tmpdir), compat.as_bytes('export'))
export_dir = est.export_savedmodel(export_dir_base, serving_input_fn)
self.assertTrue(gfile.Exists(export_dir_base))
self.assertTrue(gfile.Exists(export_dir))
self.assertTrue(
gfile.Exists(
os.path.join(
compat.as_bytes(export_dir), compat.as_bytes(
'saved_model.pb'))))
self.assertTrue(
gfile.Exists(
os.path.join(
compat.as_bytes(export_dir), compat.as_bytes('variables'))))
self.assertTrue(
gfile.Exists(
os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.index'))))
self.assertTrue(
gfile.Exists(
os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.data-00000-of-00001'))))
# Restore, to validate that the export was well-formed.
with ops.Graph().as_default() as graph:
with session_lib.Session(graph=graph) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
graph_ops = [x.name for x in graph.get_operations()]
self.assertTrue('input_example_tensor' in graph_ops)
self.assertTrue('ParseExample/ParseExample' in graph_ops)
self.assertTrue('LookupTableModel' in graph_ops)
self.assertFalse('LookupTableTrainingState' in graph_ops)
# cleanup
gfile.DeleteRecursively(tmpdir)
def testInconsistentConsumerDefaultAttrs(self):
export_dir = self._get_export_dir(
"test_strip_default_attrs_no_consumer_defaults")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Add a graph with a single variable and a test op with a defaultless
# float32 attr, "test_attr".
with session.Session(graph=ops.Graph()) as sess:
variables.Variable(1.0, dtype=dtypes.float64, name="var")
test_ops.test_attr(T=dtypes.float32, name="test_attr")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["foo"])
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Rewrite the SavedModel to remove the T attr from "test_attr".
saved_model_file = os.path.join(
export_dir, constants.SAVED_MODEL_FILENAME_PBTXT)
with open(saved_model_file) as f:
original_saved_model = f.read()
no_attr_saved_model = original_saved_model.replace("""
attr {
key: "T"
value {
type: DT_FLOAT
}
}""", "")
with open(saved_model_file, "w") as f:
f.write(no_attr_saved_model)
# Loading the SavedModel via the loader must fail because the SavedModel
# does not have any attr values for the "TestAttr" node, and there is no
# default specified in the TestAttr OpDef.
sess = session.Session(graph=ops.Graph())
if ops._USE_C_API:
error_message = "NodeDef missing attr 'T' from Op<name=TestAttr"
else:
error_message = ("Expected one attr with name .*T(out)?.* in name: "
"\"test_attr\".*")
with self.assertRaisesRegexp(ValueError, error_message):
loader.load(sess, ["foo"], export_dir)
# Rewrite the SavedModel to change the type of the T attr in "test_attr"
bad_type_saved_model = original_saved_model.replace("""
attr {
key: "T"
value {
type: DT_FLOAT
}
}""", """
attr {
key: "T"
value {
type: DT_DOUBLE
}
}""")
with open(saved_model_file, "w") as f:
f.write(bad_type_saved_model)
# Loading the SavedModel via the loader must fail because there is no
# OpKernel registered to handle T = double.
sess = session.Session(graph=ops.Graph())
with self.assertRaisesRegexp(
errors.InvalidArgumentError,
".*No OpKernel was registered to support Op \'TestAttr\' with these "
"attrs..*"):
loader.load(sess, ["foo"], export_dir)
def convert(saved_model_dir,
output_tflite=None,
output_arrays=None,
tag_set=None,
signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
batch_size=1):
"""Convert a savedmodel to tflite flatbuffer.
Args:
saved_model_dir: Saved model directory to convert.
output_tflite: File path to write result flatbuffer.
output_arrays: List of output tensor names, the default value is None, which
means conversion keeps all output tensors. This is also used to filter
tensors that are from Op currently not supported in tflite, e.g., Argmax).
tag_set: This is the set of tags to get meta_graph_def in saved_model.
signature_key: This is the signature key to extract inputs, outputs.
batch_size: If input tensor shape has None at first dimension,
e.g. (None,224,224,3), replace None with batch_size.
Returns:
The converted data. For example if tflite was the destination, then
this will be a tflite flatbuffer in a bytes array.
Raises:
ValueError: If tag_set does not indicate any meta_graph_def in saved_model,
or signature_key is not in relevant meta_graph_def,
or input shape has None beyond 1st dimension, e.g., (1,None, None, 3),
or given output_arrays are not valid causing empty outputs.
"""
if tag_set is None:
tag_set = set([tag_constants.SERVING])
meta_graph = get_meta_graph_def(saved_model_dir, tag_set)
signature_def = get_signature_def(meta_graph, signature_key)
inputs, outputs = get_inputs_outputs(signature_def)
graph = ops.Graph()
with session.Session(graph=graph) as sess:
loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir)
in_tensors = [graph.get_tensor_by_name(input_) for input_ in inputs]
# Users can use output_arrays to filter output tensors for conversion.
# If output_arrays is None, we keep all output tensors. In future, we may
# use tflite supported Op list and check whether op is custom Op to
# automatically filter output arrays.
# TODO(zhixianyan): Use tflite supported Op list to filter outputs.
if output_arrays is not None:
output_arrays = output_arrays.split(",")
out_tensors = [
graph.get_tensor_by_name(output)
for output in outputs
if output.split(":")[0] in output_arrays
]
else:
out_tensors = [graph.get_tensor_by_name(output) for output in outputs]
output_names = [node.split(":")[0] for node in outputs]
if not out_tensors:
raise ValueError(
"No valid output tensors for '{}', possible values are '{}'".format(
output_arrays, output_names))
frozen_graph_def = tf_graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), output_names)
# Toco requires fully defined tensor shape, for input tensor with None in
# their shape, e.g., (None, 224, 224, 3), we need to replace first None with
# a given batch size. For shape with more None, e.g. (None, None, None, 3),
# still be able to replace and convert, but require further investigation.
# TODO(zhixianyan): Add supports for input tensor with more None in shape.
for i in range(len(in_tensors)):
shape = in_tensors[i].get_shape().as_list()
if shape[0] is None:
shape[0] = batch_size
if None in shape[1:]:
raise ValueError(
"Only support None shape at 1st dim as batch_size. But tensor "
"'{}' 's shape '{}' has None at other dimension. ".format(
inputs[i], shape))
in_tensors[i].set_shape(shape)
result = lite.toco_convert(frozen_graph_def, in_tensors, out_tensors)
if output_tflite is not None:
with gfile.Open(output_tflite, "wb") as f:
f.write(result)
logging.info("Successfully converted to: %s", output_tflite)
return result
def test_export_savedmodel_assets(self):
tmpdir = tempfile.mkdtemp()
est = estimator.Estimator(model_fn=_model_fn_for_export_tests)
est.train(input_fn=dummy_input_fn, steps=1)
feature_spec = {'x': parsing_ops.VarLenFeature(dtype=dtypes.int64),
'y': parsing_ops.VarLenFeature(dtype=dtypes.int64)}
serving_input_receiver_fn = export.build_parsing_serving_input_receiver_fn(
feature_spec)
# Create a fake asset.
vocab_file_name = os.path.join(
compat.as_bytes(tmpdir), compat.as_bytes('my_vocab_file'))
vocab_file = gfile.GFile(vocab_file_name, mode='w')
vocab_file.write(_VOCAB_FILE_CONTENT)
vocab_file.close()
# hack in an op that uses the asset, in order to test asset export.
# this is not actually valid, of course.
def serving_input_receiver_with_asset_fn():
features, receiver_tensor = serving_input_receiver_fn()
filename = ops.convert_to_tensor(vocab_file_name,
dtypes.string,
name='asset_filepath')
ops.add_to_collection(ops.GraphKeys.ASSET_FILEPATHS, filename)
features['bogus_filename'] = filename
return export.ServingInputReceiver(features, receiver_tensor)
# Perform the export.
export_dir_base = os.path.join(
compat.as_bytes(tmpdir), compat.as_bytes('export'))
export_dir = est.export_savedmodel(
export_dir_base, serving_input_receiver_with_asset_fn)
# Check that the asset files are in the right places.
expected_vocab_file_name = os.path.join(
compat.as_bytes(export_dir), compat.as_bytes('assets/my_vocab_file'))
self.assertTrue(gfile.Exists(os.path.join(
compat.as_bytes(export_dir), compat.as_bytes('assets'))))
self.assertTrue(gfile.Exists(expected_vocab_file_name))
self.assertEqual(
compat.as_bytes(_VOCAB_FILE_CONTENT),
compat.as_bytes(gfile.GFile(expected_vocab_file_name).read()))
# Restore, to validate that the export was well-formed.
with ops.Graph().as_default() as graph:
with session.Session(graph=graph) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
assets = [
x.eval()
for x in graph.get_collection(ops.GraphKeys.ASSET_FILEPATHS)
]
self.assertItemsEqual([vocab_file_name], assets)
graph_ops = [x.name for x in graph.get_operations()]
self.assertTrue('input_example_tensor' in graph_ops)
self.assertTrue('ParseExample/ParseExample' in graph_ops)
self.assertTrue('asset_filepath' in graph_ops)
self.assertTrue('weight' in graph_ops)
# cleanup
gfile.DeleteRecursively(tmpdir)
def freeze_graph_with_def_protos(input_graph_def,
input_saver_def,
input_checkpoint,
output_node_names,
restore_op_name,
filename_tensor_name,
output_graph,
clear_devices,
initializer_nodes,
variable_names_whitelist="",
variable_names_blacklist="",
input_meta_graph_def=None,
input_saved_model_dir=None,
saved_model_tags=None,
checkpoint_version=saver_pb2.SaverDef.V2):
"""Converts all variables in a graph and checkpoint into constants.
Args:
input_graph_def: A `GraphDef`.
input_saver_def: A `SaverDef` (optional).
input_checkpoint: The prefix of a V1 or V2 checkpoint, with V2 taking
priority. Typically the result of `Saver.save()` or that of
`tf.train.latest_checkpoint()`, regardless of sharded/non-sharded or
V1/V2.
output_node_names: The name(s) of the output nodes, comma separated.
restore_op_name: Unused.
filename_tensor_name: Unused.
output_graph: String where to write the frozen `GraphDef`.
clear_devices: A Bool whether to remove device specifications.
initializer_nodes: Comma separated string of initializer nodes to run before
freezing.
variable_names_whitelist: The set of variable names to convert (optional, by
default, all variables are converted).
variable_names_blacklist: The set of variable names to omit converting
to constants (optional).
input_meta_graph_def: A `MetaGraphDef` (optional),
input_saved_model_dir: Path to the dir with TensorFlow 'SavedModel' file
and variables (optional).
saved_model_tags: Group of comma separated tag(s) of the MetaGraphDef to
load, in string format (optional).
checkpoint_version: Tensorflow variable file format (saver_pb2.SaverDef.V1
or saver_pb2.SaverDef.V2)
Returns:
Location of the output_graph_def.
"""
del restore_op_name, filename_tensor_name # Unused by updated loading code.
# 'input_checkpoint' may be a prefix if we're using Saver V2 format
if (not input_saved_model_dir and
not checkpoint_management.checkpoint_exists(input_checkpoint)):
raise ValueError("Input checkpoint '" + input_checkpoint +
"' doesn't exist!")
if not output_node_names:
raise ValueError(
"You need to supply the name of a node to --output_node_names.")
# Remove all the explicit device specifications for this node. This helps to
# make the graph more portable.
if clear_devices:
if input_meta_graph_def:
for node in input_meta_graph_def.graph_def.node:
node.device = ""
elif input_graph_def:
for node in input_graph_def.node:
node.device = ""
if input_graph_def:
_ = importer.import_graph_def(input_graph_def, name="")
with session.Session() as sess:
if input_saver_def:
saver = saver_lib.Saver(
saver_def=input_saver_def, write_version=checkpoint_version)
saver.restore(sess, input_checkpoint)
elif input_meta_graph_def:
restorer = saver_lib.import_meta_graph(
input_meta_graph_def, clear_devices=True)
restorer.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
elif input_saved_model_dir:
if saved_model_tags is None:
saved_model_tags = []
loader.load(sess, saved_model_tags, input_saved_model_dir)
else:
var_list = {}
#reader = py_checkpoint_reader.NewCheckpointReader(input_checkpoint)
reader = NewCheckpointReader(input_checkpoint)
var_to_shape_map = reader.get_variable_to_shape_map()
# List of all partition variables. Because the condition is heuristic
# based, the list could include false positives.
all_partition_variable_names = [
tensor.name.split(":")[0]
for op in sess.graph.get_operations()
for tensor in op.values()
if re.search(r"/part_\d+/", tensor.name)
]
has_partition_var = False
for key in var_to_shape_map:
try:
tensor = sess.graph.get_tensor_by_name(key + ":0")
if any(key in name for name in all_partition_variable_names):
has_partition_var = True
except KeyError:
# This tensor doesn't exist in the graph (for example it's
# 'global_step' or a similar housekeeping element) so skip it.
continue
var_list[key] = tensor
try:
saver = saver_lib.Saver(
var_list=var_list, write_version=checkpoint_version)
except TypeError as e:
# `var_list` is required to be a map of variable names to Variable
# tensors. Partition variables are Identity tensors that cannot be
# handled by Saver.
if has_partition_var:
raise ValueError(
"Models containing partition variables cannot be converted "
"from checkpoint files. Please pass in a SavedModel using "
"the flag --input_saved_model_dir.")
# Models that have been frozen previously do not contain Variables.
elif _has_no_variables(sess):
raise ValueError(
"No variables were found in this model. It is likely the model "
"was frozen previously. You cannot freeze a graph twice.")
return 0
else:
raise e
saver.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
variable_names_whitelist = (
variable_names_whitelist.replace(" ", "").split(",")
if variable_names_whitelist else None)
variable_names_blacklist = (
variable_names_blacklist.replace(" ", "").split(",")
if variable_names_blacklist else None)
if input_meta_graph_def:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_meta_graph_def.graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
else:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
# Write GraphDef to file if output path has been given.
if output_graph:
with gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
return output_graph_def
def testTags(self):
export_dir = os.path.join(test.get_temp_dir(), "test_tags")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable. SavedModel invoked to:
# - add with weights.
# - a single tag (from predefined constants).
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 42)
builder.add_meta_graph_and_variables(sess,
[tag_constants.TRAINING])
# Graph that updates the single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
# - a single tag (from predefined constants).
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 43)
builder.add_meta_graph([tag_constants.SERVING])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple custom tags.
with self.test_session(graph=ops.Graph()) as sess:
self._init_and_validate_variable(sess, "v", 44)
builder.add_meta_graph(["foo", "bar"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with a single predefined tag whose variables were saved.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, [tag_constants.TRAINING], export_dir)
self.assertEqual(
42,
ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Restore the graph with a single predefined tag whose variables were not
# saved.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
self.assertEqual(
42,
ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Restore the graph with multiple tags. Provide duplicate tags to test set
# semantics.
with self.test_session(graph=ops.Graph()) as sess:
loader.load(sess, ["foo", "bar", "foo"], export_dir)
self.assertEqual(
42,
ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Try restoring a graph with a non-existent tag. This should yield a runtime
# error.
with self.test_session(graph=ops.Graph()) as sess:
self.assertRaises(RuntimeError, loader.load, sess, ["INVALID"],
export_dir)
# Try restoring a graph where a subset of the tags match. Since tag matching
# for meta graph defs follows "all" semantics, this should yield a runtime
# error.
with self.test_session(graph=ops.Graph()) as sess:
self.assertRaises(RuntimeError, loader.load, sess, ["foo", "baz"],
export_dir)
def _run_model(iterator, args):
"""Run single-node inferencing on a checkpoint/saved_model using input tensors obtained from a Spark partition iterator and returning output tensors."""
single_node_env(args)
logging.info("===== input_mapping: {}".format(args.input_mapping))
logging.info("===== output_mapping: {}".format(args.output_mapping))
input_tensor_names = [ tensor for col,tensor in sorted(args.input_mapping.items()) ]
output_tensor_names = [ tensor for tensor,col in sorted(args.output_mapping.items()) ]
# if using a signature_def_key, get input/output tensor info from the requested signature
if args.signature_def_key:
assert args.export_dir, "Inferencing with signature_def_key requires --export_dir argument"
logging.info("===== loading meta_graph_def for tag_set ({0}) from saved_model: {1}".format(args.tag_set, args.export_dir))
meta_graph_def = get_meta_graph_def(args.export_dir, args.tag_set)
signature = signature_def_utils.get_signature_def_by_key(meta_graph_def, args.signature_def_key)
logging.info("signature: {}".format(signature))
inputs_tensor_info = signature.inputs
logging.info("inputs_tensor_info: {0}".format(inputs_tensor_info))
outputs_tensor_info = signature.outputs
logging.info("outputs_tensor_info: {0}".format(outputs_tensor_info))
result = []
with tf.Session(graph=ops_lib.Graph()) as sess:
if args.export_dir:
assert args.tag_set, "Inferencing from a saved_model requires --tag_set"
# load graph from a saved_model
logging.info("===== restoring from saved_model: {}".format(args.export_dir))
loader.load(sess, args.tag_set.split(','), args.export_dir)
elif args.model_dir:
# load graph from a checkpoint
ckpt = tf.train.latest_checkpoint(args.model_dir)
assert ckpt, "Invalid model checkpoint path: {}".format(args.model_dir)
logging.info("===== restoring from checkpoint: {}".format(ckpt + ".meta"))
saver = tf.train.import_meta_graph(ckpt + ".meta", clear_devices=True)
saver.restore(sess, ckpt)
else:
raise Exception("Inferencing requires either --model_dir or --export_dir argument")
# get list of input/output tensors (by name)
if args.signature_def_key:
input_tensors = [inputs_tensor_info[t].name for t in input_tensor_names]
output_tensors = [outputs_tensor_info[output_tensor_names[0]].name]
else:
input_tensors = [t + ':0' for t in input_tensor_names]
output_tensors = [t + ':0' for t in output_tensor_names]
logging.info("input_tensors: {0}".format(input_tensors))
logging.info("output_tensors: {0}".format(output_tensors))
# feed data in batches and return output tensors
for tensors in yield_batch(iterator, args.batch_size, len(input_tensor_names)):
inputs_feed_dict = {}
for i in range(len(input_tensors)):
inputs_feed_dict[input_tensors[i]] = tensors[i]
outputs = sess.run(output_tensors, feed_dict=inputs_feed_dict)
lengths = [ len(output) for output in outputs ]
input_size = len(tensors[0])
assert all([ l == input_size for l in lengths ]), "Output array sizes {} must match input size: {}".format(lengths, input_size)
python_outputs = [ output.tolist() for output in outputs ] # convert from numpy to standard python types
result.extend(zip(*python_outputs)) # convert to an array of tuples of "output columns"
return result
def main():
args = parse_args()
np.random.seed(args.random_seed)
tf.set_random_seed(args.random_seed)
use_crf = args.use_crf
print('Use CRF: ', use_crf)
BATCH_SIZE = args.batch_size
print('Setting up...')
train_filenames = pd.read_csv(TRAIN_FILENAME,
dtype=object,
keep_default_na=False,
na_values=[]).as_matrix()
train0_filenames = pd.read_csv(TRAIN0_FILENAME,
dtype=object,
keep_default_na=False,
na_values=[]).as_matrix()
val_filenames = pd.read_csv(VAL_FILENAME,
dtype=object,
keep_default_na=False,
na_values=[]).as_matrix()
test_filenames = pd.read_csv(TEST_FILENAME,
dtype=object,
keep_default_na=False,
na_values=[]).as_matrix()
subsets = ['train', 'val', 'test']
image_mean = IMAGE_MEAN
image_std = IMAGE_STD
image_mean_tensor = tf.constant(image_mean, dtype=tf.float32)
image_std_tensor = tf.constant(image_std, dtype=tf.float32)
print('num_train: ', NUM_TRAIN_SAMPLES)
STEPS_ONE_EPOCH = int(NUM_TRAIN_SAMPLES / BATCH_SIZE)
max_steps = STEPS_ONE_EPOCH * 21
DECAY_STEPS = STEPS_ONE_EPOCH * 5
print('steps one epoch: ', STEPS_ONE_EPOCH)
print('decay steps: ', DECAY_STEPS)
train_filenames_tfrecords = [
f[1] + '/data2.tfrecords' for f in train_filenames
]
val_filenames_tfrecords = [
f[1] + '/data2.tfrecords' for f in val_filenames
]
train_images, train_labels, train_case_names = dataset_input_from_tfrecords(
train_filenames_tfrecords,
batch_size=BATCH_SIZE,
num_epochs=5000,
shuffle=True)
val_images, val_labels, val_case_names = dataset_input_from_tfrecords(
val_filenames_tfrecords,
batch_size=BATCH_SIZE,
num_epochs=5000,
shuffle=False)
phase_train = tf.placeholder(tf.bool, name='phase_train')
global_step = tf.Variable(0, trainable=False)
image_node = tf.placeholder(tf.float32,
shape=[None, DEPTH, HEIGHT, WIDTH, 2])
label_node = tf.placeholder(tf.int32,
shape=[None, DEPTH, HEIGHT, WIDTH, 2])
if args.norm_type == 'nonorm':
image_node_new = image_node
elif args.norm_type == 'globalnorm_mean':
image_node_new = image_node - image_mean_tensor
elif args.norm_type == 'globalnorm_meanstd':
image_node_new = image_node - image_mean_tensor
image_node_new /= image_std_tensor
elif args.norm_type == 'instancenorm_mean':
image_node_new = tf.map_fn(
lambda frame: frame - tf.reduce_mean(
frame, axis=[0, 1, 2], keep_dims=True), image_node)
elif args.norm_type == 'instancenorm_meanstd':
batch_mean, batch_var = tf.nn.moments(image_node,
axes=[1, 2, 3],
keep_dims=True)
image_node_new = (image_node - batch_mean) / tf.sqrt(batch_var + 1e-6)
if args.net_type == 'myunet3d_bn_crf':
from myunet3d_basic import myunet3d_bn_crf
net_output_ops = myunet3d_bn_crf(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
elif args.net_type == 'myunet3d_crf':
from myunet3d_basic import myunet3d_crf
net_output_ops = myunet3d_crf(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.glorot_uniform_initializer(
), #tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
elif args.net_type == 'myunet3d_isbi2018':
from myunet3d_basic import myunet3d_isbi2018
net_output_ops = myunet3d_isbi2018(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
elif args.net_type == 'myunet3d_isbi2018_2':
from myunet3d_basic import myunet3d_isbi2018_2
net_output_ops = myunet3d_isbi2018_2(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
elif args.net_type == 'myunet3d_improved':
from myunet3d_improved import myunet3d_improved
net_output_ops = myunet3d_improved(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
elif args.net_type == 'myunet3d_improved_usingresize':
from myunet3d_improved_usingresize import myunet3d_improved_usingresize
net_output_ops = myunet3d_improved_usingresize(
name='ct' if args.feat_index == 0 else 'pt',
inputs=image_node_new[..., args.feat_index][..., tf.newaxis],
num_classes=NUM_CLASSES,
phase_train=phase_train,
use_bias=True,
kernel_initializer=tf.truncated_normal_initializer(stddev=0.01),
bias_initializer=tf.constant_initializer(value=0.1),
kernel_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
bias_regularizer=tf.contrib.layers.l2_regularizer(1e-4),
use_crf=args.use_crf,
args=args)
modes = ['ct', 'pt']
image_node_shaped = tf.reshape(image_node[..., args.feat_index],
[-1, DEPTH, HEIGHT * WIDTH, 1])
label_node_shaped = tf.cast(
tf.reshape(label_node[..., args.feat_index],
[-1, DEPTH, HEIGHT * WIDTH, 1]), tf.float32)
for bi in range(BATCH_SIZE):
tf.summary.image('input_{}_{}'.format(modes[args.feat_index], str(bi)),
image_node_shaped[bi][tf.newaxis, ...], 1)
tf.summary.image('label_{}_{}'.format(modes[args.feat_index], str(bi)),
label_node_shaped[bi][tf.newaxis, ...], 1)
prob_flat = tf.unstack(net_output_ops['y_prob'], num=2,
axis=4) # [-1, 32, 32, 32]
prob_flat = [
tf.reshape(v, [-1, DEPTH, HEIGHT * WIDTH]) for v in prob_flat
] #[-1,32,32*32]
prob_flat = tf.concat(prob_flat, axis=1) #[-1, 32*16, 32*32]
prob_flat = tf.cast(
tf.reshape(prob_flat, [-1, 2 * DEPTH, HEIGHT * WIDTH, 1]), tf.float32)
pred_flat = tf.cast(
tf.reshape(net_output_ops['y_'], [-1, DEPTH, HEIGHT * WIDTH, 1]),
tf.float32)
for bi in range(BATCH_SIZE):
tf.summary.image('prob_{}_{}'.format(modes[args.feat_index], str(bi)),
prob_flat[bi][tf.newaxis, ...], 1)
tf.summary.image('pred_{}_{}'.format(modes[args.feat_index], str(bi)),
pred_flat[bi][tf.newaxis, ...], 1)
pred_op = net_output_ops['y_']
prob_op = net_output_ops['y_prob']
print('pred_op shape: ', pred_op.shape)
print('prob_op shape: ', prob_op.shape)
dice_op = dice_tf(labels=label_node[..., args.feat_index],
logits=prob_op[..., 1])
if args.use_crf:
prob_crf_op = net_output_ops['y_prob_crf']
prob_crf_flat = tf.unstack(prob_crf_op, num=2,
axis=4) # [-1, 32, 32, 32]
prob_crf_flat = [
tf.reshape(v, [-1, DEPTH, HEIGHT * WIDTH]) for v in prob_crf_flat
] #[-1,32,32*32]
prob_crf_flat = tf.concat(prob_crf_flat, axis=1) #[-1, 32*16, 32*32]
prob_crf_flat = tf.cast(
tf.reshape(prob_crf_flat, [-1, 2 * DEPTH, HEIGHT * WIDTH, 1]),
tf.float32)
pred_crf_flat = tf.cast(
tf.reshape(net_output_ops['y_crf'],
[-1, DEPTH, HEIGHT * WIDTH, 1]), tf.float32)
for bi in range(BATCH_SIZE):
tf.summary.image(
'prob_crf_{}_{}'.format(modes[args.feat_index], str(bi)),
prob_crf_flat[bi][tf.newaxis, ...], 1)
tf.summary.image(
'pred_crf_{}_{}'.format(modes[args.feat_index], str(bi)),
pred_crf_flat[bi][tf.newaxis, ...], 1)
pred_crf_op = net_output_ops['y_crf']
dice_crf_op = dice_tf(labels=label_node[..., args.feat_index],
logits=prob_crf_op[..., 1])
# 2. set up a loss function # regularization loss
reg_constant = tf.constant(args.reg_coef, dtype=tf.float32)
reg_losses = tf.get_collection(tf.GraphKeys.REGULARIZATION_LOSSES)
if len(reg_losses) != 0:
loss_reg_op = tf.add_n(reg_losses)
else:
loss_reg_op = tf.constant(0, dtype=tf.float32)
tf.summary.scalar('loss_reg_{}'.format(modes[args.feat_index]),
loss_reg_op)
# 2. set up a loss function
if args.loss_type == 'ce':
ce = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=net_output_ops['logits'],
labels=label_node[..., args.feat_index])
loss_op = tf.reduce_mean(ce)
elif args.loss_type == 'dice':
loss_op = dice_loss(labels=label_node[..., args.feat_index],
logits=net_output_ops['logits'])
elif args.loss_type == 'focal':
loss_op = focal_loss(labels=label_node[..., args.feat_index],
logits=net_output_ops['logits'])
elif args.loss_type == 'bce':
loss_op = binary_cross_entropy(labels=label_node[..., args.feat_index],
logits=net_output_ops['logits'])
tf.summary.scalar('loss_{}'.format(modes[args.feat_index]), loss_op)
total_loss_op = loss_op + tf.multiply(reg_constant, loss_reg_op)
tf.summary.scalar('total_loss_{}'.format(modes[args.feat_index]),
total_loss_op)
if args.use_crf:
if args.loss_type == 'ce':
ce_crf = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=net_output_ops['logits_crf'],
labels=label_node[..., args.feat_index])
loss_crf_op = tf.reduce_mean(ce_crf)
elif args.loss_type == 'dice':
loss_crf_op = dice_loss(labels=label_node[..., args.feat_index],
logits=net_output_ops['logits_crf'])
elif args.loss_type == 'focal':
loss_crf_op = focal_loss(labels=label_node[..., args.feat_index],
logits=net_output_ops['logits_crf'])
elif args.loss_type == 'bce':
loss_crf_op = binary_cross_entropy(
labels=label_node[..., args.feat_index],
logits=net_output_ops['logits_crf'])
tf.summary.scalar('loss_crf_{}'.format(modes[args.feat_index]),
loss_crf_op)
total_loss_crf_op = loss_crf_op + tf.multiply(reg_constant,
loss_reg_op)
tf.summary.scalar('total_loss_{}_crf'.format(modes[args.feat_index]),
total_loss_crf_op)
if args.use_crf:
crf_sw_op = tf.get_default_graph().get_tensor_by_name(
modes[args.feat_index] + '/crf/spatial_weights:0')
crf_bw_op = tf.get_default_graph().get_tensor_by_name(
modes[args.feat_index] + '/crf/bilateral_weights:0')
crf_cm_op = tf.get_default_graph().get_tensor_by_name(
modes[args.feat_index] + '/crf/compatibility_matrix:0')
noncrf_vars = []
crf_vars = []
for v in tf.trainable_variables():
if 'crf' in v.name:
crf_vars.append(v)
else:
noncrf_vars.append(v)
if args.lr_policy == 'constant':
lr_op = args.base_lr
elif args.lr_policy == 'piecewise':
#### piecewise learning rate
boundaries = [i * DECAY_STEPS for i in [1, 2, 3, 4, 5]]
staged_lr = [
args.base_lr * x for x in [1, 0.5, 0.25, 0.125, 0.0625, 0.03125]
]
lr_op = tf.train.piecewise_constant(global_step, boundaries, staged_lr)
elif args.lr_policy == 'expdecay':
learning_rate_decay_factor = 0.5
decay_steps = 1000
# Decay the learning rate exponentially based on the number of steps.
lr_op = tf.train.exponential_decay(args.base_lr,
global_step,
decay_steps,
learning_rate_decay_factor,
staircase=True)
tf.summary.scalar('learning_rate_{}'.format(modes[args.feat_index]), lr_op)
if args.opt_type == 'gd':
optimiser = tf.train.GradientDescentOptimizer(learning_rate=lr_op)
elif args.opt_type == 'gd_diff_lr':
optimiser = tf.train.GradientDescentOptimizer(learning_rate=lr_op)
optimiser_crf = tf.train.GradientDescentOptimizer(learning_rate=lr_op *
args.crf_lr_scale)
elif args.opt_type == 'adam_diff_lr':
optimiser = tf.train.AdamOptimizer(learning_rate=lr_op)
optimiser_crf = tf.train.AdamOptimizer(learning_rate=lr_op *
args.crf_lr_scale)
elif args.opt_type == 'adam_gd_diff_lr':
optimiser = tf.train.AdamOptimizer(learning_rate=lr_op)
optimiser_crf = tf.train.GradientDescentOptimizer(learning_rate=lr_op *
args.crf_lr_scale)
elif args.opt_type == 'momentum':
optimiser = tf.train.MomentumOptimizer(learning_rate=lr_op,
momentum=0.9)
elif args.opt_type == 'adam':
optimiser = tf.train.AdamOptimizer(learning_rate=lr_op)
elif args.opt_type == 'adadelta':
optimiser = tf.train.AdadeltaOptimizer(learning_rate=lr_op)
elif args.opt_type == 'adagrad':
optimiser = tf.train.AdagradOptimizer(learning_rate=lr_op)
update_ops = tf.get_collection(tf.GraphKeys.UPDATE_OPS)
if 'diff_lr' in args.opt_type:
grads = tf.gradients(total_loss_crf_op, noncrf_vars + crf_vars)
grads1 = grads[:len(noncrf_vars)]
grads2 = grads[len(noncrf_vars):]
with tf.control_dependencies(update_ops):
train_op1 = optimiser.apply_gradients(zip(grads1, noncrf_vars))
train_op2 = optimiser_crf.apply_gradients(zip(grads2, crf_vars))
train_op = tf.group(train_op1, train_op2)
else:
with tf.control_dependencies(update_ops):
if args.use_crf:
train_op = optimiser.minimize(total_loss_crf_op,
global_step=global_step)
else:
train_op = optimiser.minimize(total_loss_op,
global_step=global_step)
log_postfix = '{}_{}x{}x{}_f{}_{}_opt{}_lr{}{}_b{}_loss{}_crf{}_reg{}_rs{}'.format(
args.net_type, HEIGHT, WIDTH, DEPTH, modes[args.feat_index],
args.norm_type, args.opt_type, args.lr_policy, args.base_lr,
args.batch_size, args.loss_type, args.use_crf, args.reg_coef,
args.random_seed)
if args.use_crf:
log_postfix = '{}/{}_{}_{}_{}_{}_{}_lr_scale_{}'.format(
log_postfix, args.sw_weight, args.bw_weight, args.cm_weight,
args.theta_alpha, args.theta_beta, args.theta_gamma,
args.crf_lr_scale)
log_dir = '{}_{}'.format(args.log_dir, log_postfix)
if args.restore_ckpt == "" and args.action == 'train':
if tf.gfile.Exists(log_dir):
tf.gfile.DeleteRecursively(log_dir)
tf.gfile.MakeDirs(log_dir)
# set up log file
current_time = datetime.now().strftime("%Y-%m-%d_%H:%M:%S.%f")
log_filename = '{}/log_{}.txt'.format(log_dir, current_time)
cmd_filename = '{}/log_curve_{}.sh'.format(log_dir, current_time)
log_file_handle = open(log_filename, 'w')
with open(cmd_filename, 'w') as cmd_file_handle:
cmd_file_handle.write(
'python plot_learning_curves_single.py {} {} {}\n'.format(
log_dir, log_postfix, log_filename))
saver = tf.train.Saver(tf.global_variables(), max_to_keep=0)
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
if args.action == 'train':
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(log_dir, sess.graph)
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
if args.restore_ckpt != "" and args.action == 'train':
variables = tf.global_variables()
restore_vars = []
restore_var_names = [line.strip() for line in \
open('trainable_variables_{}_only.txt'.format(modes[args.feat_index]),'r').readlines()]
for v in variables:
found = False
for name in restore_var_names:
if name in v.name:
found = True
break
if found:
restore_vars.append(v)
for v in restore_vars:
print(v.name)
restore_saver = tf.train.Saver(restore_vars)
restore_saver.restore(sess, args.restore_ckpt)
if args.restore_ckpt != "" and args.action == 'test':
saver1 = tf.train.import_meta_graph('{}.meta'.format(
args.restore_ckpt))
saver1.restore(sess, '{}'.format(args.restore_ckpt))
if args.restore_from_saved_model != "":
pretrained_variables = {}
with tf.Graph().as_default():
# restore_from = '/media/ubuntu/working/petct_cnn/logs_dltk_ctpt_1e-4_v1(dont delete)/ct_only/1513006564'
with tf.Session() as sess1:
graph_old = loader.load(sess1, ['serve'],
args.restore_from_saved_model)
for v in tf.global_variables():
pretrained_variables[v.name] = v.eval()
count = 0
for v in tf.global_variables():
for name in pretrained_variables.keys():
if modes[args.feat_index] + '/' + name == v.name:
print('{}: {} ==> {}'.format(count, name, v.name))
sess.run(v.assign(pretrained_variables[name]))
count += 1
break
print(count, ' variables copied from previous model!')
if args.action == 'test':
if args.save_prob_dir != '' and not os.path.exists(args.save_prob_dir):
os.makedirs(args.save_prob_dir)
if args.save_dir != '' and not os.path.exists(args.save_dir):
os.makedirs(args.save_dir)
for idx, filename in enumerate(args.test_filenames.split(',')):
temp = pd.read_csv(DATA_ROOT + '/' + filename,
dtype=object,
keep_default_na=False,
na_values=[]).as_matrix()
if idx == 0:
test_filenames = temp
else:
test_filenames = np.concatenate([test_filenames, temp], axis=0)
dice_val = []
if args.use_crf:
dice_val_crf = []
alldata = {}
for f in test_filenames:
# for f in np.concatenate([val_filenames, test_filenames], axis=0):
subject_id = f[0]
img_fn = f[1]
case_name = img_fn.split('/')[-1]
# if '2141' in case_name:
# continue
# Read the image nii with sitk and keep the pointer to the sitk.Image
# of an input
ct_sitk = sitk.ReadImage(
str(os.path.join(img_fn, 'InputCT_ROI.nii.gz')))
ct = sitk.GetArrayFromImage(ct_sitk).astype((np.float32))
pt_sitk = sitk.ReadImage(
str(os.path.join(img_fn, 'InputPET_SUV_ROI.nii.gz')))
pt = sitk.GetArrayFromImage(pt_sitk).astype((np.float32))
lbl_ct = sitk.GetArrayFromImage(
sitk.ReadImage(
str(
os.path.join(
img_fn, 'GTV_Primary_ROI_CT{}.nii.gz'.format(
GT_POSTFIX))))).astype(np.uint8)
lbl_pt = sitk.GetArrayFromImage(
sitk.ReadImage(
str(
os.path.join(
img_fn, 'GTV_Primary_ROI_PET{}.nii.gz'.format(
GT_POSTFIX))))).astype(np.uint8)
ct[ct > 200.] = 200.
ct[ct < -500.] = -500.
ct = 255 * (ct + 500) / (700.)
pt[pt < 0.01] = 0.01
pt[pt > 20.] = 20.
pt = 255 * (pt - 0.01) / (19.99)
image = np.concatenate([ct[..., np.newaxis], pt[..., np.newaxis]],
axis=3)
label = np.concatenate(
[lbl_ct[..., np.newaxis], lbl_pt[..., np.newaxis]], axis=3)
# if image_mean != None:
# image -= np.reshape(image_mean, [1,1,1,2])
# if image_std != None:
# image /= np.reshape(image_std, [1,1,1,2])
if args.use_crf:
pred, pred_crf, prob, prob_crf = sess.run(
[pred_op, pred_crf_op, prob_op, prob_crf_op],
feed_dict={
image_node: image[np.newaxis, ...],
label_node: label[np.newaxis, ...],
phase_train: False
})
else:
pred, prob = sess.run(
[pred_op, prob_op],
feed_dict={
image_node: image[np.newaxis, ...],
label_node: label[np.newaxis, ...],
phase_train: False
})
dice_val_ = computeDice(label[..., args.feat_index], pred[0])
dice_val.append(dice_val_)
if args.use_crf:
dice_val_crf_ = computeDice(label[..., args.feat_index],
pred_crf[0])
dice_val_crf.append(dice_val_crf_)
if args.save_prob_dir != '':
alldata[case_name] = {}
alldata[case_name]['image'] = image
alldata[case_name]['label'] = label
alldata[case_name]['pred'] = pred[0]
alldata[case_name]['prob'] = prob[0]
#alldata[case_name]['ct_sitk'] = ct_sitk
log_file_handle.write('{} {} {}'.format(
case_name, dice_val_,
dice_val_crf_ if args.use_crf == 1 else ''))
if args.save_dir != '':
case_save_dir = '{}/{}'.format(args.save_dir, case_name)
if not os.path.exists(case_save_dir):
os.makedirs(case_save_dir)
if args.use_crf == 1:
new_sitk_ct = sitk.GetImageFromArray(pred[0].astype(
np.int32))
new_sitk_ct.CopyInformation(ct_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf1_pred_{}_before.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
new_sitk_ct = sitk.GetImageFromArray(pred_crf[0].astype(
np.int32))
new_sitk_ct.CopyInformation(ct_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf1_pred_{}_after.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
new_sitk_ct = sitk.GetImageFromArray(
prob[0][..., 1].astype(np.float32))
new_sitk_ct.CopyInformation(pt_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf1_prob_{}_before.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
new_sitk_ct = sitk.GetImageFromArray(
prob_crf[0][..., 1].astype(np.float32))
new_sitk_ct.CopyInformation(pt_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf1_prob_{}_after.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
else:
new_sitk_ct = sitk.GetImageFromArray(pred[0].astype(
np.int32))
new_sitk_ct.CopyInformation(ct_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf0_pred_{}.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
new_sitk_ct = sitk.GetImageFromArray(
prob[0][..., 1].astype(np.float32))
new_sitk_ct.CopyInformation(pt_sitk)
sitk.WriteImage(
new_sitk_ct,
str('{}/crf0_prob_{}.nii.gz'.format(
case_save_dir,
'ct' if args.feat_index == 0 else 'pt')))
if args.save_prob_dir != '':
alldata['val_filenames'] = val_filenames
alldata['test_filenames'] = test_filenames
pickle.dump(alldata, open(args.save_prob_dir + '/alldata.p', "wb"))
log_file_handle.write('Mean {} Dice(Before CRF): {}'.format(
modes[args.feat_index], np.mean(np.array(dice_val))))
if args.use_crf:
log_file_handle.write('Mean {} Dice(After CRF): {}'.format(
modes[args.feat_index], np.mean(np.array(dice_val_crf))))
return
coord = tf.train.Coordinator()
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
step = 0
if args.step > 0:
step = args.step
while not coord.should_stop():
if step % STEPS_ONE_EPOCH == 0:
# val one epoch
# for subi, filenames in enumerate([train0_filenames, val_filenames, test_filenames]):
for subi, filenames in enumerate([
train0_filenames,
np.concatenate([val_filenames, test_filenames], axis=0)
]):
dice_val = []
if args.use_crf:
dice_val_crf = []
for f in filenames:
subject_id = f[0]
img_fn = f[1]
case_name = img_fn.split('/')[-1]
# if '2141' in case_name:
# continue
# Read the image nii with sitk and keep the pointer to the sitk.Image
# of an input
ct_sitk = sitk.ReadImage(
str(os.path.join(img_fn, 'InputCT_ROI.nii.gz')))
ct = sitk.GetArrayFromImage(ct_sitk).astype((np.float32))
pt_sitk = sitk.ReadImage(
str(os.path.join(img_fn, 'InputPET_SUV_ROI.nii.gz')))
pt = sitk.GetArrayFromImage(pt_sitk).astype((np.float32))
lbl_ct = sitk.GetArrayFromImage(
sitk.ReadImage(
str(
os.path.join(
img_fn,
'GTV_Primary_ROI_CT{}.nii.gz'.format(
GT_POSTFIX))))).astype(np.uint8)
lbl_pt = sitk.GetArrayFromImage(
sitk.ReadImage(
str(
os.path.join(
img_fn,
'GTV_Primary_ROI_PET{}.nii.gz'.format(
GT_POSTFIX))))).astype(np.uint8)
ct[ct > 200.] = 200.
ct[ct < -500.] = -500.
ct = 255 * (ct + 500) / (700.)
pt[pt < 0.01] = 0.01
pt[pt > 20.] = 20.
pt = 255 * (pt - 0.01) / (19.99)
image = np.concatenate(
[ct[..., np.newaxis], pt[..., np.newaxis]], axis=3)
label = np.concatenate(
[lbl_ct[..., np.newaxis], lbl_pt[..., np.newaxis]],
axis=3)
# if image_mean != None:
# image -= np.reshape(image_mean, [1,1,1,2])
# if image_std != None:
# image /= np.reshape(image_std, [1,1,1,2])
if args.use_crf:
pred, pred_crf = sess.run(
[pred_op, pred_crf_op],
feed_dict={
image_node: image[np.newaxis, ...],
label_node: label[np.newaxis, ...],
phase_train: False
})
else:
pred = sess.run(
[pred_op],
feed_dict={
image_node: image[np.newaxis, ...],
label_node: label[np.newaxis, ...],
phase_train: False
})
dice_val_ = computeDice(label[..., args.feat_index],
pred[0])
dice_val.append(dice_val_)
if args.use_crf:
dice_val_crf_ = computeDice(
label[..., args.feat_index], pred_crf[0])
dice_val_crf.append(dice_val_crf_)
if args.use_crf:
log_file_handle.write('{} {} {}\n'.format(
case_name, dice_val_, dice_val_crf_))
else:
log_file_handle.write('{} {}\n'.format(
case_name, dice_val_))
log_file_handle.write(
'{} {} Mean {} Dice(Before CRF): {}\n'.format(
step, subsets[subi], modes[args.feat_index],
np.mean(np.array(dice_val))))
if args.use_crf:
log_file_handle.write(
'{} {} Mean {} Dice(After CRF): {}\n'.format(
step, subsets[subi], modes[args.feat_index],
np.mean(np.array(dice_val_crf))))
if args.use_crf:
crf_sw_value, crf_bw_value, crf_cm_value \
= sess.run([crf_sw_op, crf_bw_op, crf_cm_op])
print('sw:', crf_sw_value)
print('bw:', crf_bw_value)
print('cm:', crf_cm_value)
if step % STEPS_ONE_EPOCH == 0:
if args.use_crf:
checkpoint_path = os.path.join(log_dir, 'model_crf.ckpt')
else:
checkpoint_path = os.path.join(log_dir, 'model.ckpt')
saver.save(sess, checkpoint_path, global_step=step)
# train one epoch
start_time = time.time()
image_batch, label_batch, case_name_batch = sess.run(
[train_images, train_labels, train_case_names])
if args.use_crf:
summary, _, loss_reg_value, \
total_loss_crf_value, loss_crf_value, \
total_loss_value, loss_value, \
pred, pred_crf, \
dice, dice_crf \
= sess.run([summary_op, train_op, loss_reg_op,
total_loss_crf_op, loss_crf_op,
total_loss_op, loss_op,
pred_op, pred_crf_op,
dice_op, dice_crf_op],
feed_dict={image_node: image_batch,
label_node: label_batch,
phase_train: False})
else:
summary, _, loss_reg_value, \
total_loss_value, loss_value, \
pred, dice = sess.run([summary_op, train_op, loss_reg_op,
total_loss_op, loss_op,
pred_op, dice_op],
feed_dict={image_node: image_batch,
label_node: label_batch,
phase_train: True})
duration = time.time() - start_time
# if args.use_crf:
# log_file_handle.write('{}, step {:d}, {:.6f}, ' \
# '{:.6f}, {:.6f}, ' \
# '{:.6f}, {:.6f}, ' \
# '{:.3f}, {:.3f}, {}\n'.format(modes[args.feat_index],
# step, loss_reg_value,
# total_loss_crf_value, loss_crf_value,
# total_loss_value, loss_value,
# dice, dice_crf, case_name_batch[0]))
# log_file_handle.flush()
# else:
# log_file_handle.write('{}, step {:d}, {:.6f}, ' \
# '{:.6f}, {:.6f}, ' \
# '{:.3f}, {}\n'.format(modes[args.feat_index],
# step, loss_reg_value,
# total_loss_value, loss_value,
# dice, case_name_batch[0]))
# log_file_handle.flush()
if step % STEPS_ONE_EPOCH == 0:
summary_writer.add_summary(summary, global_step=step)
step += 1
if step == max_steps:
break
log_file_handle.close()
coord.request_stop()
coord.join(threads)
sess.close()
def load_tf_model(model_path, tags=(tag_constants.SERVING, ), config=None):
"""Loads the model at the specified path.
Args:
model_path: the path to either session_bundle or SavedModel
tags: the tags that determines the model to load.
config: tf.ConfigProto containing session configuration options.
Returns:
A pair of (Session, map<string, SignatureDef>) objects.
Raises:
PredictionError: if the model could not be loaded.
"""
if loader.maybe_saved_model_directory(model_path):
try:
logging.info("Importing tensorflow.contrib in load_tf_model")
# pylint: disable=redefined-outer-name,unused-variable,g-import-not-at-top
import tensorflow as tf
from tensorflow.python.framework.ops import Graph
# pylint: enable=redefined-outer-name,unused-variable,g-import-not-at-top
if tf.__version__.startswith("1.0"):
session = tf_session.Session(target="",
graph=None,
config=config)
else:
session = tf_session.Session(target="",
graph=Graph(),
config=config)
meta_graph = loader.load(session,
tags=list(tags),
export_dir=model_path)
except Exception as e: # pylint: disable=broad-except
raise PredictionError(
PredictionError.FAILED_TO_LOAD_MODEL,
"Failed to load the model due to bad model data."
" tags: %s\n%s" % (list(tags), str(e)))
else:
raise PredictionError(
PredictionError.FAILED_TO_LOAD_MODEL,
"Cloud ML only supports TF 1.0 or above and models "
"saved in SavedModel format.")
if session is None:
raise PredictionError(
PredictionError.FAILED_TO_LOAD_MODEL,
"Failed to create session when loading the model")
if not meta_graph.signature_def:
raise PredictionError(
PredictionError.FAILED_TO_LOAD_MODEL,
"MetaGraph must have at least one signature_def.")
# Remove invalid signatures from the signature map.
invalid_signatures = []
for signature_name in meta_graph.signature_def:
try:
signature = meta_graph.signature_def[signature_name]
_update_dtypes(session.graph, signature.inputs)
_update_dtypes(session.graph, signature.outputs)
except ValueError as e:
logging.warn("Error updating signature %s: %s", signature_name,
str(e))
invalid_signatures.append(signature_name)
for signature_name in invalid_signatures:
del meta_graph.signature_def[signature_name]
return session, meta_graph.signature_def
def freeze_saved_model(saved_model_dir, input_arrays, input_shapes,
output_arrays, tag_set, signature_key):
"""Converts a SavedModel to a frozen graph.
Args:
saved_model_dir: SavedModel directory to convert.
input_arrays: List of input tensors to freeze graph with. Uses input arrays
from SignatureDef when none are provided. (default None)
input_shapes: Map of strings representing input tensor names to list of
integers representing input shapes (e.g., {"foo": : [1, 16, 16, 3]}).
Automatically determined when input shapes is None (e.g., {"foo" : None}).
(default None)
output_arrays: List of output tensors to freeze graph with. Uses output
arrays from SignatureDef when none are provided. (default None)
tag_set: Set of tags identifying the MetaGraphDef within the SavedModel to
analyze. All tags in the tag set must be present. (default "serve")
signature_key: Key identifying SignatureDef containing inputs and outputs.
Returns:
frozen_graph_def: Frozen GraphDef.
in_tensors: List of input tensors for the graph.
out_tensors: List of output tensors for the graph.
Raises:
ValueError:
SavedModel doesn't contain a MetaGraphDef identified by tag_set.
signature_key is not in the MetaGraphDef.
input_shapes does not match the length of input_arrays.
input_arrays or output_arrays are not valid.
Unable to load Session.
"""
# Set default values for inputs if they are set to None.
if signature_key is None:
signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
if tag_set is None:
tag_set = set([tag_constants.SERVING])
# Read SignatureDef.
meta_graph = _get_meta_graph_def(saved_model_dir, tag_set)
signature_def = _get_signature_def(meta_graph, signature_key)
inputs, outputs = _get_inputs_outputs(signature_def)
graph = ops.Graph()
with session.Session(graph=graph) as sess:
# TODO(nupurgarg): Throw ValueError if SavedModel has assets/ directory.
loader.load(sess, meta_graph.meta_info_def.tags, saved_model_dir)
# Gets input and output tensors.
# TODO(zhixianyan): Use TFLite supported Op list to filter outputs.
in_tensors = _get_tensors(graph, inputs, input_arrays)
out_tensors = _get_tensors(graph, outputs, output_arrays)
# Gets fully defined tensor shape.
for tensor in in_tensors:
if (input_shapes and tensor.name in input_shapes
and input_shapes[tensor.name] is not None):
shape = input_shapes[tensor.name]
else:
shape = tensor.get_shape().as_list()
tensor.set_shape(shape)
output_names = [node.split(":")[0] for node in outputs]
frozen_graph_def = tf_graph_util.convert_variables_to_constants(
sess, graph.as_graph_def(), output_names)
return frozen_graph_def, in_tensors, out_tensors
raise ValueError("Unable to load Session.")
def test_savedmodel_state_override(self):
random_model = RandomStateSpaceModel(
state_dimension=5,
state_noise_dimension=4,
configuration=state_space_model.StateSpaceModelConfiguration(
exogenous_feature_columns=[layers.real_valued_column("exogenous")],
dtype=dtypes.float64, num_features=1))
estimator = estimators.StateSpaceRegressor(
model=random_model,
optimizer=gradient_descent.GradientDescentOptimizer(0.1))
combined_input_fn = input_pipeline.WholeDatasetInputFn(
input_pipeline.NumpyReader({
feature_keys.FilteringFeatures.TIMES: [1, 2, 3, 4],
feature_keys.FilteringFeatures.VALUES: [1., 2., 3., 4.],
"exogenous": [-1., -2., -3., -4.]
}))
estimator.train(combined_input_fn, steps=1)
export_location = estimator.export_savedmodel(
self.get_temp_dir(),
estimator.build_raw_serving_input_receiver_fn())
with ops.Graph().as_default() as graph:
random_model.initialize_graph()
with self.session(graph=graph) as session:
variables.global_variables_initializer().run()
evaled_start_state = session.run(random_model.get_start_state())
evaled_start_state = [
state_element[None, ...] for state_element in evaled_start_state]
with ops.Graph().as_default() as graph:
with self.session(graph=graph) as session:
signatures = loader.load(
session, [tag_constants.SERVING], export_location)
first_split_filtering = saved_model_utils.filter_continuation(
continue_from={
feature_keys.FilteringResults.STATE_TUPLE: evaled_start_state},
signatures=signatures,
session=session,
features={
feature_keys.FilteringFeatures.TIMES: [1, 2],
feature_keys.FilteringFeatures.VALUES: [1., 2.],
"exogenous": [[-1.], [-2.]]})
second_split_filtering = saved_model_utils.filter_continuation(
continue_from=first_split_filtering,
signatures=signatures,
session=session,
features={
feature_keys.FilteringFeatures.TIMES: [3, 4],
feature_keys.FilteringFeatures.VALUES: [3., 4.],
"exogenous": [[-3.], [-4.]]
})
combined_filtering = saved_model_utils.filter_continuation(
continue_from={
feature_keys.FilteringResults.STATE_TUPLE: evaled_start_state},
signatures=signatures,
session=session,
features={
feature_keys.FilteringFeatures.TIMES: [1, 2, 3, 4],
feature_keys.FilteringFeatures.VALUES: [1., 2., 3., 4.],
"exogenous": [[-1.], [-2.], [-3.], [-4.]]
})
split_predict = saved_model_utils.predict_continuation(
continue_from=second_split_filtering,
signatures=signatures,
session=session,
steps=1,
exogenous_features={
"exogenous": [[[-5.]]]})
combined_predict = saved_model_utils.predict_continuation(
continue_from=combined_filtering,
signatures=signatures,
session=session,
steps=1,
exogenous_features={
"exogenous": [[[-5.]]]})
for state_key, combined_state_value in combined_filtering.items():
if state_key == feature_keys.FilteringResults.TIMES:
continue
self.assertAllClose(
combined_state_value, second_split_filtering[state_key])
for prediction_key, combined_value in combined_predict.items():
self.assertAllClose(combined_value, split_predict[prediction_key])
def run_saved_model_with_feed_dict(saved_model_dir,
tag_set,
signature_def_key,
input_tensor_key_feed_dict,
outdir,
overwrite_flag,
tf_debug=False):
"""Runs SavedModel and fetch all outputs.
Runs the input dictionary through the MetaGraphDef within a SavedModel
specified by the given tag_set and SignatureDef. Also save the outputs to file
if outdir is not None.
Args:
saved_model_dir: Directory containing the SavedModel to execute.
tag_set: Group of tag(s) of the MetaGraphDef with the SignatureDef map, in
string format, separated by ','. For tag-set contains multiple tags, all
tags must be passed in.
signature_def_key: A SignatureDef key string.
input_tensor_key_feed_dict: A dictionary maps input keys to numpy ndarrays.
outdir: A directory to save the outputs to. If the directory doesn't exist,
it will be created.
overwrite_flag: A boolean flag to allow overwrite output file if file with
the same name exists.
tf_debug: A boolean flag to use TensorFlow Debugger (TFDBG) to observe the
intermediate Tensor values and runtime GraphDefs while running the
SavedModel.
Raises:
RuntimeError: An error when output file already exists and overwrite is not
enabled.
"""
# Get a list of output tensor names.
meta_graph_def = get_meta_graph_def(saved_model_dir, tag_set)
# Re-create feed_dict based on input tensor name instead of key as session.run
# uses tensor name.
inputs_tensor_info = _get_inputs_tensor_info_from_meta_graph_def(
meta_graph_def, signature_def_key)
inputs_feed_dict = {
inputs_tensor_info[key].name: tensor
for key, tensor in input_tensor_key_feed_dict.items()
}
# Get outputs
outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def(
meta_graph_def, signature_def_key)
# Sort to preserve order because we need to go from value to key later.
output_tensor_keys_sorted = sorted(outputs_tensor_info.keys())
output_tensor_names_sorted = [
outputs_tensor_info[tensor_key].name
for tensor_key in output_tensor_keys_sorted
]
with session.Session(graph=ops_lib.Graph()) as sess:
loader.load(sess, tag_set.split(','), saved_model_dir)
if tf_debug:
sess = local_cli_wrapper.LocalCLIDebugWrapperSession(sess)
outputs = sess.run(output_tensor_names_sorted,
feed_dict=inputs_feed_dict)
for i, output in enumerate(outputs):
output_tensor_key = output_tensor_keys_sorted[i]
print('Result for output key %s:\n%s' %
(output_tensor_key, output))
# Only save if outdir is specified.
if outdir:
# Create directory if outdir does not exist
if not os.path.isdir(outdir):
os.makedirs(outdir)
output_full_path = os.path.join(outdir,
output_tensor_key + '.npy')
# If overwrite not enabled and file already exist, error out
if not overwrite_flag and os.path.exists(output_full_path):
raise RuntimeError(
'Output file %s already exists. Add \"--overwrite\" to overwrite'
' the existing output files.' % output_full_path)
np.save(output_full_path, output)
print('Output %s is saved to %s' %
(output_tensor_key, output_full_path))
def testClearExtraneousSavers(self):
export_dir = os.path.join(test.get_temp_dir(),
"test_clear_extraneous_savers")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Create a variable and a Saver.
with ops.Graph().as_default() as graph:
with session.Session(target="",
config=config_pb2.ConfigProto(
device_count={"CPU": 2})) as sess:
self._init_and_validate_variable(sess, "v", 42)
# Add two Savers, which should be removed in
# add_meta_graph_and_variables() in favor of the locally added one.
saver1 = tf_saver.Saver()
graph.add_to_collection(ops.GraphKeys.SAVERS, saver1)
saver2 = tf_saver.Saver()
graph.add_to_collection(ops.GraphKeys.SAVERS, saver2)
# Confirm there are two SaverDefs.
savers = graph.get_collection(ops.GraphKeys.SAVERS)
self.assertEqual(2, len(savers))
# Confirm there are two Save and two Restore ops.
save_op_names = set([
x.name for x in graph.get_operations()
if x.type == "SaveV2"
])
self.assertSetEqual(set(["save/SaveV2", "save_1/SaveV2"]),
save_op_names)
restore_op_names = set([
x.name for x in graph.get_operations()
if x.type == "RestoreV2"
])
self.assertSetEqual(
set(["save/RestoreV2", "save_1/RestoreV2"]),
restore_op_names)
# The SavedModel builder adds its own Saver' for a total of three.
builder.add_meta_graph_and_variables(sess,
[tag_constants.TRAINING],
clear_devices=True)
# Save the SavedModel to disk.
builder.save()
# Restore the graph.
with ops.Graph().as_default() as graph:
with self.test_session(graph=graph) as sess:
loader.load(sess, [tag_constants.TRAINING], export_dir)
self.assertEqual(
42,
ops.get_collection(
ops.GraphKeys.GLOBAL_VARIABLES)[0].eval())
# Confirm that the reloaded graph has only one SaverDef.
savers = ops.get_collection(ops.GraphKeys.SAVERS)
self.assertEqual(1, len(savers))
# The reloaded graph should have exactly one Save and one Restore op.
save_op_names = set([
x.name for x in graph.get_operations()
if x.type == "SaveV2"
])
self.assertSetEqual(set(["save_2/SaveV2"]), save_op_names)
restore_op_names = set([
x.name for x in graph.get_operations()
if x.type == "RestoreV2"
])
self.assertSetEqual(set(["save_2/RestoreV2"]),
restore_op_names)
def test_export_savedmodel(self):
tmpdir = tempfile.mkdtemp()
est, export_input_fn = _build_estimator_for_export_tests(tmpdir)
extra_file_name = os.path.join(compat.as_bytes(tmpdir),
compat.as_bytes('my_extra_file'))
extra_file = tf.gfile.GFile(extra_file_name, mode='w')
extra_file.write(EXTRA_FILE_CONTENT)
extra_file.close()
assets_extra = {'some/sub/directory/my_extra_file': extra_file_name}
export_dir_base = os.path.join(compat.as_bytes(tmpdir),
compat.as_bytes('export'))
export_dir = est.export_savedmodel(export_dir_base, export_input_fn,
assets_extra=assets_extra)
self.assertTrue(tf.gfile.Exists(export_dir_base))
self.assertTrue(tf.gfile.Exists(export_dir))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('saved_model.pb'))))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('variables'))))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.index'))))
self.assertTrue(tf.gfile.Exists(os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('variables/variables.data-00000-of-00001'))))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir), compat.as_bytes('assets'))))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('assets/my_vocab_file'))))
self.assertEqual(
compat.as_bytes(VOCAB_FILE_CONTENT),
compat.as_bytes(tf.gfile.GFile(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('assets/my_vocab_file'))).read()))
expected_extra_path = os.path.join(
compat.as_bytes(export_dir),
compat.as_bytes('assets.extra/some/sub/directory/my_extra_file'))
self.assertTrue(tf.gfile.Exists(
os.path.join(compat.as_bytes(export_dir),
compat.as_bytes('assets.extra'))))
self.assertTrue(tf.gfile.Exists(expected_extra_path))
self.assertEqual(
compat.as_bytes(EXTRA_FILE_CONTENT),
compat.as_bytes(tf.gfile.GFile(expected_extra_path).read()))
expected_vocab_file = os.path.join(compat.as_bytes(tmpdir),
compat.as_bytes('my_vocab_file'))
# Restore, to validate that the export was well-formed.
with tf.Graph().as_default() as graph:
with tf.Session(graph=graph) as sess:
loader.load(sess, [tag_constants.SERVING], export_dir)
assets = [x.eval()
for x in graph.get_collection(ops.GraphKeys.ASSET_FILEPATHS)]
self.assertItemsEqual([expected_vocab_file], assets)
graph_ops = [x.name for x in graph.get_operations()]
self.assertTrue('input_example_tensor' in graph_ops)
self.assertTrue('ParseExample/ParseExample' in graph_ops)
self.assertTrue('linear/linear/feature/matmul' in graph_ops)
# cleanup
tf.gfile.DeleteRecursively(tmpdir)
def testStripDefaultAttrs(self):
export_dir = self._get_export_dir("test_strip_default_attrs")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Add a graph with two float32 variables and a Complex Op composing them
# with strip_default_attrs enabled.
with session.Session(graph=ops.Graph()) as sess:
real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(
sess, ["foo"], strip_default_attrs=True)
# Add a graph with the same float32 variables and a Complex Op composing
# them with strip_default_attrs disabled.
with session.Session(graph=ops.Graph()) as sess:
real_num = variables.Variable(1.0, dtype=dtypes.float32, name="real")
imag_num = variables.Variable(2.0, dtype=dtypes.float32, name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph(["bar"], strip_default_attrs=False)
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Loading graph "foo" via the loader must restore the defaults for the
# "Complex" node based on the "Complex" OpDef in the Op registry.
sess = session.Session(graph=ops.Graph())
meta_graph_def = loader.load(sess, ["foo"], export_dir)
complex_node = test_util.get_node_def_from_graph("complex",
meta_graph_def.graph_def)
self.assertIn("T", complex_node.attr)
self.assertIn("Tout", complex_node.attr)
# Load graph "foo" from disk as-is to verify default attrs are stripped.
# pylint: disable=protected-access
saved_model_pb = loader_impl._parse_saved_model(export_dir)
self.assertIsNotNone(saved_model_pb)
# pylint: enable=protected-access
meta_graph_foo_def = None
meta_graph_bar_def = None
for meta_graph_def in saved_model_pb.meta_graphs:
if set(meta_graph_def.meta_info_def.tags) == set(["foo"]):
meta_graph_foo_def = meta_graph_def
elif set(meta_graph_def.meta_info_def.tags) == set(["bar"]):
meta_graph_bar_def = meta_graph_def
self.assertIsNotNone(meta_graph_foo_def)
self.assertIsNotNone(meta_graph_bar_def)
# "Complex" Op has 2 attributes with defaults:
# o "T" : float32. (input type)
# o "Tout" : complex64. (output type)
# "Complex" Op in graph "foo" shouldn't have attributes "T" and "Tout".
# Graph "foo" was saved with strip_default_attrs set to True.
node_def = test_util.get_node_def_from_graph("complex",
meta_graph_foo_def.graph_def)
self.assertNotIn("T", node_def.attr)
self.assertNotIn("Tout", node_def.attr)
# "Complex" Op in graph "bar" must have attributes "T" and "Tout".
# Graph "bar" was saved with strip_default_attrs set to False.
node_def = test_util.get_node_def_from_graph("complex",
meta_graph_bar_def.graph_def)
self.assertIn("T", node_def.attr)
self.assertIn("Tout", node_def.attr)
def testTags(self):
export_dir = os.path.join(compat.as_bytes(tf.test.get_temp_dir()),
compat.as_bytes("tags"))
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Graph with a single variable. SavedModel invoked to:
# - add with weights.
# - a single tag (from predefined constants).
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())
builder.add_meta_graph_and_variables(sess,
[constants.TAG_TRAINING])
# Graph that updates the single variable. SavedModel invoked to:
# - simply add the model (weights are not updated).
# - a single tag (from predefined constants).
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())
builder.add_meta_graph([constants.TAG_SERVING])
# Graph that updates the single variable. SavedModel is invoked:
# - to add the model (weights are not updated).
# - multiple custom tags.
with self.test_session(graph=tf.Graph()) as sess:
v = tf.Variable(44, name="v")
sess.run(tf.initialize_all_variables())
self.assertEqual(44, v.eval())
builder.add_meta_graph(["foo", "bar"])
# Save the SavedModel to disk.
builder.save()
# Restore the graph with a single predefined tag whose variables were saved.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, [constants.TAG_TRAINING], export_dir)
tf.train.write_graph(sess.graph.as_graph_def(),
"/tmp/",
"training_graph",
as_text=True)
self.assertEqual(
42,
tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
# Restore the graph with a single predefined tag whose variables were not
# saved.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, [constants.TAG_SERVING], export_dir)
tf.train.write_graph(sess.graph.as_graph_def(),
"/tmp/",
"serving_graph",
as_text=True)
self.assertEqual(
42,
tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
# Restore the graph with multiple tags. Provide duplicate tags to test set
# semantics.
with self.test_session(graph=tf.Graph()) as sess:
loader.load(sess, ["foo", "bar", "foo"], export_dir)
self.assertEqual(
42,
tf.get_collection(tf.GraphKeys.VARIABLES)[0].eval())
# Try restoring a graph with a non-existent tag. This should yield a runtime
# error.
with self.test_session(graph=tf.Graph()) as sess:
self.assertRaises(RuntimeError, loader.load, sess, ["INVALID"],
export_dir)
# Try restoring a graph where a subset of the tags match. Since tag matching
# for meta graph defs follows "all" semantics, this should yield a runtime
# error.
with self.test_session(graph=tf.Graph()) as sess:
self.assertRaises(RuntimeError, loader.load, sess, ["foo", "baz"],
export_dir)
def _fit_restore_fit_test_template(self, estimator_fn, dtype):
"""Tests restoring previously fit models."""
model_dir = tempfile.mkdtemp(dir=self.get_temp_dir())
exogenous_feature_columns = (
feature_column.numeric_column("exogenous"),
)
first_estimator = estimator_fn(model_dir, exogenous_feature_columns)
times = numpy.arange(20, dtype=numpy.int64)
values = numpy.arange(20, dtype=dtype.as_numpy_dtype)
exogenous = numpy.arange(20, dtype=dtype.as_numpy_dtype)
features = {
feature_keys.TrainEvalFeatures.TIMES: times,
feature_keys.TrainEvalFeatures.VALUES: values,
"exogenous": exogenous
}
train_input_fn = input_pipeline.RandomWindowInputFn(
input_pipeline.NumpyReader(features), shuffle_seed=2, num_threads=1,
batch_size=16, window_size=16)
eval_input_fn = input_pipeline.RandomWindowInputFn(
input_pipeline.NumpyReader(features), shuffle_seed=3, num_threads=1,
batch_size=16, window_size=16)
first_estimator.train(input_fn=train_input_fn, steps=5)
first_loss_before_fit = first_estimator.evaluate(
input_fn=eval_input_fn, steps=1)["loss"]
first_estimator.train(input_fn=train_input_fn, steps=50)
first_loss_after_fit = first_estimator.evaluate(
input_fn=eval_input_fn, steps=1)["loss"]
self.assertLess(first_loss_after_fit, first_loss_before_fit)
second_estimator = estimator_fn(model_dir, exogenous_feature_columns)
second_estimator.train(input_fn=train_input_fn, steps=2)
whole_dataset_input_fn = input_pipeline.WholeDatasetInputFn(
input_pipeline.NumpyReader(features))
whole_dataset_evaluation = second_estimator.evaluate(
input_fn=whole_dataset_input_fn, steps=1)
exogenous_values_ten_steps = {
"exogenous": numpy.arange(
10, dtype=dtype.as_numpy_dtype)[None, :, None]
}
predict_input_fn = input_pipeline.predict_continuation_input_fn(
evaluation=whole_dataset_evaluation,
exogenous_features=exogenous_values_ten_steps,
steps=10)
# Also tests that limit_epochs in predict_continuation_input_fn prevents
# infinite iteration
(estimator_predictions,
) = list(second_estimator.predict(input_fn=predict_input_fn))
self.assertAllEqual([10, 1], estimator_predictions["mean"].shape)
input_receiver_fn = first_estimator.build_raw_serving_input_receiver_fn()
export_location = first_estimator.export_savedmodel(self.get_temp_dir(),
input_receiver_fn)
with ops.Graph().as_default():
with session.Session() as sess:
signatures = loader.load(sess, [tag_constants.SERVING], export_location)
# Test that prediction and filtering can continue from evaluation output
saved_prediction = saved_model_utils.predict_continuation(
continue_from=whole_dataset_evaluation,
steps=10,
exogenous_features=exogenous_values_ten_steps,
signatures=signatures,
session=sess)
# Saved model predictions should be the same as Estimator predictions
# starting from the same evaluation.
for prediction_key, prediction_value in estimator_predictions.items():
self.assertAllClose(prediction_value,
numpy.squeeze(
saved_prediction[prediction_key], axis=0))
first_filtering = saved_model_utils.filter_continuation(
continue_from=whole_dataset_evaluation,
features={
feature_keys.FilteringFeatures.TIMES: times[None, -1] + 2,
feature_keys.FilteringFeatures.VALUES: values[None, -1] + 2.,
"exogenous": values[None, -1, None] + 12.
},
signatures=signatures,
session=sess)
# Test that prediction and filtering can continue from filtering output
second_saved_prediction = saved_model_utils.predict_continuation(
continue_from=first_filtering,
steps=1,
exogenous_features={
"exogenous": numpy.arange(
1, dtype=dtype.as_numpy_dtype)[None, :, None]
},
signatures=signatures,
session=sess)
self.assertEqual(
times[-1] + 3,
numpy.squeeze(
second_saved_prediction[feature_keys.PredictionResults.TIMES]))
saved_model_utils.filter_continuation(
continue_from=first_filtering,
features={
feature_keys.FilteringFeatures.TIMES: times[-1] + 3,
feature_keys.FilteringFeatures.VALUES: values[-1] + 3.,
"exogenous": values[-1, None] + 13.
},
signatures=signatures,
session=sess)
# Test cold starting
six.assertCountEqual(
self,
[feature_keys.FilteringFeatures.TIMES,
feature_keys.FilteringFeatures.VALUES,
"exogenous"],
signatures.signature_def[
feature_keys.SavedModelLabels.COLD_START_FILTER].inputs.keys())
batch_numpy_times = numpy.tile(
numpy.arange(30, dtype=numpy.int64)[None, :], (10, 1))
batch_numpy_values = numpy.ones([10, 30, 1])
state = saved_model_utils.cold_start_filter(
signatures=signatures,
session=sess,
features={
feature_keys.FilteringFeatures.TIMES: batch_numpy_times,
feature_keys.FilteringFeatures.VALUES: batch_numpy_values,
"exogenous": 10. + batch_numpy_values
}
)
predict_times = numpy.tile(
numpy.arange(30, 45, dtype=numpy.int64)[None, :], (10, 1))
predictions = saved_model_utils.predict_continuation(
continue_from=state,
times=predict_times,
exogenous_features={
"exogenous": numpy.tile(numpy.arange(
15, dtype=dtype.as_numpy_dtype), (10,))[None, :, None]
},
signatures=signatures,
session=sess)
self.assertAllEqual([10, 15, 1], predictions["mean"].shape)
def test_one_shot_prediction_head_export(self):
def _new_temp_dir():
return os.path.join(test.get_temp_dir(), str(ops.uid()))
model_dir = _new_temp_dir()
categorical_column = feature_column.categorical_column_with_hash_bucket(
key="categorical_exogenous_feature", hash_bucket_size=16)
exogenous_feature_columns = [
feature_column.numeric_column("2d_exogenous_feature", shape=(2, )),
feature_column.embedding_column(
categorical_column=categorical_column, dimension=10)
]
estimator = ts_estimators.TimeSeriesRegressor(
model=ar_model.ARModel(
periodicities=10,
input_window_size=10,
output_window_size=6,
num_features=5,
exogenous_feature_columns=exogenous_feature_columns,
prediction_model_factory=functools.partial(
ar_model.LSTMPredictionModel, num_units=10)),
head_type=ts_head_lib.OneShotPredictionHead,
model_dir=model_dir)
def train_input_fn():
num_range = math_ops.range(16, dtype=dtypes.int64)
features = {
feature_keys.TrainEvalFeatures.TIMES:
array_ops.expand_dims(num_range, axis=0),
feature_keys.TrainEvalFeatures.VALUES:
array_ops.expand_dims(array_ops.tile(num_range[:, None],
[1, 5]),
axis=0),
"2d_exogenous_feature":
array_ops.ones([1, 16, 2]),
"categorical_exogenous_feature":
array_ops.expand_dims(array_ops.tile(["strkey"], [16])[:,
None],
axis=0)
}
return features
estimator.train(input_fn=train_input_fn, steps=5)
result = estimator.evaluate(input_fn=train_input_fn, steps=1)
self.assertIn("average_loss", result)
self.assertNotIn(feature_keys.State.STATE_TUPLE, result)
input_receiver_fn = estimator.build_raw_serving_input_receiver_fn()
export_location = estimator.export_saved_model(_new_temp_dir(),
input_receiver_fn)
graph = ops.Graph()
with graph.as_default():
with session_lib.Session() as session:
signatures = loader.load(session, [tag_constants.SERVING],
export_location)
self.assertEqual([feature_keys.SavedModelLabels.PREDICT],
list(signatures.signature_def.keys()))
predict_signature = signatures.signature_def[
feature_keys.SavedModelLabels.PREDICT]
six.assertCountEqual(self, [
feature_keys.FilteringFeatures.TIMES,
feature_keys.FilteringFeatures.VALUES,
"2d_exogenous_feature", "categorical_exogenous_feature"
], predict_signature.inputs.keys())
features = {
feature_keys.TrainEvalFeatures.TIMES:
numpy.tile(
numpy.arange(35, dtype=numpy.int64)[None, :], [2, 1]),
feature_keys.TrainEvalFeatures.VALUES:
numpy.tile(
numpy.arange(20, dtype=numpy.float32)[None, :, None],
[2, 1, 5]),
"2d_exogenous_feature":
numpy.ones([2, 35, 2]),
"categorical_exogenous_feature":
numpy.tile(
numpy.array(["strkey"] * 35)[None, :, None], [2, 1, 1])
}
feeds = {
graph.as_graph_element(input_value.name):
features[input_key]
for input_key, input_value in
predict_signature.inputs.items()
}
fetches = {
output_key: graph.as_graph_element(output_value.name)
for output_key, output_value in
predict_signature.outputs.items()
}
output = session.run(fetches, feed_dict=feeds)
self.assertEqual((2, 15, 5), output["mean"].shape)
# Build a parsing input function, then make a tf.Example for it to parse.
export_location = estimator.export_saved_model(
_new_temp_dir(),
estimator.build_one_shot_parsing_serving_input_receiver_fn(
filtering_length=20, prediction_length=15))
graph = ops.Graph()
with graph.as_default():
with session_lib.Session() as session:
example = example_pb2.Example()
times = example.features.feature[
feature_keys.TrainEvalFeatures.TIMES]
values = example.features.feature[
feature_keys.TrainEvalFeatures.VALUES]
times.int64_list.value.extend(range(35))
for i in range(20):
values.float_list.value.extend([
float(i) * 2. + feature_number
for feature_number in range(5)
])
real_feature = example.features.feature["2d_exogenous_feature"]
categortical_feature = example.features.feature[
"categorical_exogenous_feature"]
for i in range(35):
real_feature.float_list.value.extend([1, 1])
categortical_feature.bytes_list.value.append(b"strkey")
# Serialize the tf.Example for feeding to the Session
examples = [example.SerializeToString()] * 2
signatures = loader.load(session, [tag_constants.SERVING],
export_location)
predict_signature = signatures.signature_def[
feature_keys.SavedModelLabels.PREDICT]
((_, input_value), ) = predict_signature.inputs.items()
feeds = {graph.as_graph_element(input_value.name): examples}
fetches = {
output_key: graph.as_graph_element(output_value.name)
for output_key, output_value in
predict_signature.outputs.items()
}
output = session.run(fetches, feed_dict=feeds)
self.assertEqual((2, 15, 5), output["mean"].shape)
def testStripDefaultAttrsInconsistentConsumerDefaults(self):
export_dir = os.path.join(
test.get_temp_dir(),
"test_strip_default_attrs_no_consumer_defaults")
builder = saved_model_builder.SavedModelBuilder(export_dir)
# Add a graph with two float32 variables and a Complex Op composing them
# with strip_default_attrs enabled. This must remove the following
# defaults for the "Complex" Op:
# o "T" : float32. (input type)
# o "Tout" : complex64. (output type)
with session.Session(graph=ops.Graph()) as sess:
real_num = variables.Variable(1.0,
dtype=dtypes.float32,
name="real")
imag_num = variables.Variable(2.0,
dtype=dtypes.float32,
name="imag")
math_ops.complex(real_num, imag_num, name="complex")
sess.run(variables.global_variables_initializer())
builder.add_meta_graph_and_variables(sess, ["foo"],
strip_default_attrs=True)
# Save the SavedModel to disk in text format.
builder.save(as_text=True)
# Update the Op registry to remove defaults for all attrs("T", "Tout") from
# the "Complex" OpDef.
complex_op_def = op_def_registry.get_registered_ops()["Complex"]
original_complex_op_def = op_def_pb2.OpDef()
original_complex_op_def.CopyFrom(complex_op_def)
for attr_def in complex_op_def.attr:
attr_def.ClearField("default_value")
# Loading the SavedModel via the loader must fail because the SavedModel
# does not have any attr values for the "Complex" node and the current
# op registry does not have have any default values for the "Complex" op.
sess = session.Session(graph=ops.Graph())
with self.assertRaisesRegexp(
ValueError,
"Expected one attr with name .*T(out)?.* in name: \"complex\".*"
):
loader.load(sess, ["foo"], export_dir)
# Update the Op registry to change the defaults for attr "Tout"
# (complex64 -> complex128).
complex_op_def.CopyFrom(original_complex_op_def)
for attr_def in complex_op_def.attr:
if attr_def.name == "Tout":
attr_def.default_value.type = types_pb2.DT_COMPLEX128
# Loading the SavedModel via the loader must set "Tout" attr_value for the
# "Complex" node according to the latest defaults (complex128). This is
# expected to fail the model import as there is no OpKernel registered to
# handle attrs "T" (float32) and "Tout" (complex128).
sess = session.Session(graph=ops.Graph())
with self.assertRaisesRegexp(
errors.InvalidArgumentError,
".*No OpKernel was registered to support Op \'Complex\' with these "
"attrs..*"):
loader.load(sess, ["foo"], export_dir)
def freeze_graph_with_def_protos(input_graph_def,
input_saver_def,
input_checkpoint,
output_node_names,
restore_op_name,
filename_tensor_name,
output_graph,
clear_devices,
initializer_nodes,
variable_names_whitelist="",
variable_names_blacklist="",
input_meta_graph_def=None,
input_saved_model_dir=None,
saved_model_tags=None,
checkpoint_version=saver_pb2.SaverDef.V2):
"""Converts all variables in a graph and checkpoint into constants."""
del restore_op_name, filename_tensor_name # Unused by updated loading code.
# 'input_checkpoint' may be a prefix if we're using Saver V2 format
if (not input_saved_model_dir
and not saver_lib.checkpoint_exists(input_checkpoint)):
print("Input checkpoint '" + input_checkpoint + "' doesn't exist!")
return -1
if not output_node_names:
print("You need to supply the name of a node to --output_node_names.")
return -1
# Remove all the explicit device specifications for this node. This helps to
# make the graph more portable.
if clear_devices:
if input_meta_graph_def:
for node in input_meta_graph_def.graph_def.node:
node.device = ""
elif input_graph_def:
for node in input_graph_def.node:
node.device = ""
if input_graph_def:
_ = importer.import_graph_def(input_graph_def, name="")
with session.Session() as sess:
if input_saver_def:
saver = saver_lib.Saver(saver_def=input_saver_def,
write_version=checkpoint_version)
saver.restore(sess, input_checkpoint)
elif input_meta_graph_def:
restorer = saver_lib.import_meta_graph(input_meta_graph_def,
clear_devices=True)
restorer.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
elif input_saved_model_dir:
if saved_model_tags is None:
saved_model_tags = []
loader.load(sess, saved_model_tags, input_saved_model_dir)
else:
var_list = {}
reader = pywrap_tensorflow.NewCheckpointReader(input_checkpoint)
var_to_shape_map = reader.get_variable_to_shape_map()
for key in var_to_shape_map:
try:
tensor = sess.graph.get_tensor_by_name(key + ":0")
except KeyError:
# This tensor doesn't exist in the graph (for example it's
# 'global_step' or a similar housekeeping element) so skip it.
continue
var_list[key] = tensor
saver = saver_lib.Saver(var_list=var_list,
write_version=checkpoint_version)
saver.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
variable_names_whitelist = (variable_names_whitelist.replace(
" ", "").split(",") if variable_names_whitelist else None)
variable_names_blacklist = (variable_names_blacklist.replace(
" ", "").split(",") if variable_names_blacklist else None)
if input_meta_graph_def:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_meta_graph_def.graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
else:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
# Write GraphDef to file if output path has been given.
if output_graph:
with gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
return output_graph_def
def load(self, tags, path):
assert self.session is not None
loader.load(self.session, tags, path)
def run_saved_model_with_feed_dict(saved_model_dir,
tag_set,
signature_def_key,
input_tensor_key_feed_dict,
outdir,
overwrite_flag,
worker=None,
init_tpu=False,
use_tfrt=False,
tf_debug=False):
"""Runs SavedModel and fetch all outputs.
Runs the input dictionary through the MetaGraphDef within a SavedModel
specified by the given tag_set and SignatureDef. Also save the outputs to file
if outdir is not None.
Args:
saved_model_dir: Directory containing the SavedModel to execute.
tag_set: Group of tag(s) of the MetaGraphDef with the SignatureDef map, in
string format, separated by ','. For tag-set contains multiple tags, all
tags must be passed in.
signature_def_key: A SignatureDef key string.
input_tensor_key_feed_dict: A dictionary maps input keys to numpy ndarrays.
outdir: A directory to save the outputs to. If the directory doesn't exist,
it will be created.
overwrite_flag: A boolean flag to allow overwrite output file if file with
the same name exists.
worker: If provided, the session will be run on the worker. Valid worker
specification is a bns or gRPC path.
init_tpu: If true, the TPU system will be initialized after the session
is created.
use_tfrt: If true, TFRT session will be used.
tf_debug: A boolean flag to use TensorFlow Debugger (TFDBG) to observe the
intermediate Tensor values and runtime GraphDefs while running the
SavedModel.
Raises:
ValueError: When any of the input tensor keys is not valid.
RuntimeError: An error when output file already exists and overwrite is not
enabled.
"""
# Get a list of output tensor names.
meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir,
tag_set)
# Re-create feed_dict based on input tensor name instead of key as session.run
# uses tensor name.
inputs_tensor_info = _get_inputs_tensor_info_from_meta_graph_def(
meta_graph_def, signature_def_key)
# Check if input tensor keys are valid.
for input_key_name in input_tensor_key_feed_dict.keys():
if input_key_name not in inputs_tensor_info:
raise ValueError(
'"%s" is not a valid input key. Please choose from %s, or use '
'--show option.' %
(input_key_name, '"' + '", "'.join(inputs_tensor_info.keys()) + '"'))
inputs_feed_dict = {
inputs_tensor_info[key].name: tensor
for key, tensor in input_tensor_key_feed_dict.items()
}
# Get outputs
outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def(
meta_graph_def, signature_def_key)
# Sort to preserve order because we need to go from value to key later.
output_tensor_keys_sorted = sorted(outputs_tensor_info.keys())
output_tensor_names_sorted = [
outputs_tensor_info[tensor_key].name
for tensor_key in output_tensor_keys_sorted
]
config = None
if use_tfrt:
logging.info('Using TFRT session.')
config = config_pb2.ConfigProto(
experimental=config_pb2.ConfigProto.Experimental(use_tfrt=True))
with session.Session(worker, graph=ops_lib.Graph(), config=config) as sess:
if init_tpu:
print('Initializing TPU System ...')
# This is needed for freshly started worker, or if the job
# restarts after a preemption.
sess.run(tpu.initialize_system())
loader.load(sess, tag_set.split(','), saved_model_dir)
if tf_debug:
sess = local_cli_wrapper.LocalCLIDebugWrapperSession(sess)
outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
for i, output in enumerate(outputs):
output_tensor_key = output_tensor_keys_sorted[i]
print('Result for output key %s:\n%s' % (output_tensor_key, output))
# Only save if outdir is specified.
if outdir:
# Create directory if outdir does not exist
if not os.path.isdir(outdir):
os.makedirs(outdir)
output_full_path = os.path.join(outdir, output_tensor_key + '.npy')
# If overwrite not enabled and file already exist, error out
if not overwrite_flag and os.path.exists(output_full_path):
raise RuntimeError(
'Output file %s already exists. Add \"--overwrite\" to overwrite'
' the existing output files.' % output_full_path)
np.save(output_full_path, output)
print('Output %s is saved to %s' % (output_tensor_key,
output_full_path))
def freeze_graph_with_def_protos(input_graph_def,
input_saver_def,
input_checkpoint,
output_node_names,
restore_op_name,
filename_tensor_name,
output_graph,
clear_devices,
initializer_nodes,
variable_names_whitelist="",
variable_names_blacklist="",
input_meta_graph_def=None,
input_saved_model_dir=None,
saved_model_tags=None,
checkpoint_version=saver_pb2.SaverDef.V2):
"""Converts all variables in a graph and checkpoint into constants."""
del restore_op_name, filename_tensor_name # Unused by updated loading code.
if not output_node_names:
print("You need to supply the name of a node to --output_node_names.")
return -1
# Remove all the explicit device specifications for this node. This helps to
# make the graph more portable.
if clear_devices:
if input_meta_graph_def:
for node in input_meta_graph_def.graph_def.node:
node.device = ""
elif input_graph_def:
for node in input_graph_def.node:
node.device = ""
if input_graph_def:
_ = importer.import_graph_def(input_graph_def, name="")
with session.Session() as sess:
if input_saver_def:
saver = saver_lib.Saver(saver_def=input_saver_def,
write_version=checkpoint_version)
saver.restore(sess, input_checkpoint)
elif input_meta_graph_def:
restorer = saver_lib.import_meta_graph(input_meta_graph_def,
clear_devices=True)
restorer.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
elif input_saved_model_dir:
if saved_model_tags is None:
saved_model_tags = []
loader.load(sess, saved_model_tags, input_saved_model_dir)
else:
var_list = {}
reader = pywrap_tensorflow.NewCheckpointReader(input_checkpoint)
var_to_shape_map = reader.get_variable_to_shape_map()
# List of all partition variables. Because the condition is heuristic
# based, the list could include false positives.
all_parition_variable_names = [
tensor.name.split(":")[0]
for op in sess.graph.get_operations()
for tensor in op.values()
if re.search(r"/part_\d+/", tensor.name)
]
has_partition_var = False
for key in var_to_shape_map:
try:
tensor = sess.graph.get_tensor_by_name(key + ":0")
if any(key in name
for name in all_parition_variable_names):
has_partition_var = True
except KeyError:
# This tensor doesn't exist in the graph (for example it's
# 'global_step' or a similar housekeeping element) so skip it.
continue
var_list[key] = tensor
try:
saver = saver_lib.Saver(var_list=var_list,
write_version=checkpoint_version)
except TypeError as e:
# `var_list` is required to be a map of variable names to Variable
# tensors. Partition variables are Identity tensors that cannot be
# handled by Saver.
if has_partition_var:
print(
"Models containing partition variables cannot be converted "
"from checkpoint files. Please pass in a SavedModel using "
"the flag --input_saved_model_dir.")
return -1
else:
raise e
saver.restore(sess, input_checkpoint)
if initializer_nodes:
sess.run(initializer_nodes.replace(" ", "").split(","))
variable_names_whitelist = (variable_names_whitelist.replace(
" ", "").split(",") if variable_names_whitelist else None)
variable_names_blacklist = (variable_names_blacklist.replace(
" ", "").split(",") if variable_names_blacklist else None)
if input_meta_graph_def:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_meta_graph_def.graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
else:
output_graph_def = graph_util.convert_variables_to_constants(
sess,
input_graph_def,
output_node_names.replace(" ", "").split(","),
variable_names_whitelist=variable_names_whitelist,
variable_names_blacklist=variable_names_blacklist)
# Write GraphDef to file if output path has been given.
if output_graph:
with gfile.GFile(output_graph, "wb") as f:
f.write(output_graph_def.SerializeToString())
return output_graph_def
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 _run_model(iterator, args, tf_args):
"""mapPartitions function to run single-node inferencing from a checkpoint/saved_model, using the model's input/output mappings.
Args:
:iterator: input RDD partition iterator.
:args: arguments for TFModel, in argparse format
:tf_args: arguments for TensorFlow inferencing code, in argparse or ARGV format.
Returns:
An iterator of result data.
"""
single_node_env(tf_args)
logging.info("===== input_mapping: {}".format(args.input_mapping))
logging.info("===== output_mapping: {}".format(args.output_mapping))
input_tensor_names = [
tensor for col, tensor in sorted(args.input_mapping.items())
]
output_tensor_names = [
tensor for tensor, col in sorted(args.output_mapping.items())
]
# if using a signature_def_key, get input/output tensor info from the requested signature
if args.signature_def_key:
assert args.export_dir, "Inferencing with signature_def_key requires --export_dir argument"
logging.info(
"===== loading meta_graph_def for tag_set ({0}) from saved_model: {1}"
.format(args.tag_set, args.export_dir))
meta_graph_def = get_meta_graph_def(args.export_dir, args.tag_set)
signature = signature_def_utils.get_signature_def_by_key(
meta_graph_def, args.signature_def_key)
logging.debug("signature: {}".format(signature))
inputs_tensor_info = signature.inputs
logging.debug("inputs_tensor_info: {0}".format(inputs_tensor_info))
outputs_tensor_info = signature.outputs
logging.debug("outputs_tensor_info: {0}".format(outputs_tensor_info))
result = []
global global_sess, global_args
if global_sess and global_args == args:
# if graph/session already loaded/started (and using same args), just reuse it
sess = global_sess
else:
# otherwise, create new session and load graph from disk
tf.reset_default_graph()
sess = tf.Session(graph=tf.get_default_graph())
if args.export_dir:
assert args.tag_set, "Inferencing from a saved_model requires --tag_set"
# load graph from a saved_model
logging.info("===== restoring from saved_model: {}".format(
args.export_dir))
loader.load(sess, args.tag_set.split(','), args.export_dir)
elif args.model_dir:
# load graph from a checkpoint
ckpt = tf.train.latest_checkpoint(args.model_dir)
assert ckpt, "Invalid model checkpoint path: {}".format(
args.model_dir)
logging.info("===== restoring from checkpoint: {}".format(ckpt +
".meta"))
saver = tf.train.import_meta_graph(ckpt + ".meta",
clear_devices=True)
saver.restore(sess, ckpt)
else:
raise Exception(
"Inferencing requires either --model_dir or --export_dir argument"
)
global_sess = sess
global_args = args
# get list of input/output tensors (by name)
if args.signature_def_key:
input_tensors = [
inputs_tensor_info[t].name for t in input_tensor_names
]
output_tensors = [outputs_tensor_info[output_tensor_names[0]].name]
else:
input_tensors = [t + ':0' for t in input_tensor_names]
output_tensors = [t + ':0' for t in output_tensor_names]
logging.info("input_tensors: {0}".format(input_tensors))
logging.info("output_tensors: {0}".format(output_tensors))
# feed data in batches and return output tensors
for tensors in yield_batch(iterator, args.batch_size,
len(input_tensor_names)):
inputs_feed_dict = {}
for i in range(len(input_tensors)):
inputs_feed_dict[input_tensors[i]] = tensors[i]
outputs = sess.run(output_tensors, feed_dict=inputs_feed_dict)
lengths = [len(output) for output in outputs]
input_size = len(tensors[0])
assert all([
length == input_size for length in lengths
]), "Output array sizes {} must match input size: {}".format(
lengths, input_size)
python_outputs = [output.tolist() for output in outputs
] # convert from numpy to standard python types
result.extend(zip(*python_outputs)
) # convert to an array of tuples of "output columns"
return result
def __init__(self,
export_dir,
signature_def_key=None,
signature_def=None,
input_names=None,
output_names=None,
tags=None,
graph=None,
config=None):
"""Initialize a `CoreEstimatorPredictor`.
Args:
export_dir: a path to a directory containing a `SavedModel`.
signature_def_key: Optional string specifying the signature to use. If
`None`, then `DEFAULT_SERVING_SIGNATURE_DEF_KEY` is used. Only one of
`signature_def_key` and `signature_def` should be specified.
signature_def: A `SignatureDef` proto specifying the inputs and outputs
for prediction. Only one of `signature_def_key` and `signature_def`
should be specified.
input_names: A dictionary mapping strings to `Tensor`s in the `SavedModel`
that represent the input. The keys can be any string of the user's
choosing.
output_names: A dictionary mapping strings to `Tensor`s in the
`SavedModel` that represent the output. The keys can be any string of
the user's choosing.
tags: Optional. Comma separated list of tags that will be used to retrieve
the correct `SignatureDef`. Defaults to `DEFAULT_TAGS`.
graph: Optional. The Tensorflow `graph` in which prediction should be
done.
config: `ConfigProto` proto used to configure the session.
Raises:
ValueError: If more than one of signature_def_key OR signature_def OR
(input_names AND output_names) is specified.
"""
_check_signature_arguments(signature_def_key, signature_def,
input_names, output_names)
tags = tags or DEFAULT_TAGS
self._graph = graph or ops.Graph()
with self._graph.as_default():
self._session = session.Session(config=config)
loader.load(self._session, tags.split(','), export_dir)
if input_names is None:
if signature_def is None:
signature_def = _get_signature_def(signature_def_key,
export_dir, tags)
input_names = {k: v.name for k, v in signature_def.inputs.items()}
output_names = {
k: v.name
for k, v in signature_def.outputs.items()
}
self._feed_tensors = {
k: self._graph.get_tensor_by_name(v)
for k, v in input_names.items()
}
self._fetch_tensors = {
k: self._graph.get_tensor_by_name(v)
for k, v in output_names.items()
}
def test_one_shot_prediction_head_export(self):
model_dir = self.get_temp_dir()
categorical_column = feature_column.categorical_column_with_hash_bucket(
key="categorical_exogenous_feature", hash_bucket_size=16)
exogenous_feature_columns = [
feature_column.numeric_column("2d_exogenous_feature", shape=(2, )),
feature_column.embedding_column(
categorical_column=categorical_column, dimension=10)
]
estimator = ts_estimators.TimeSeriesRegressor(
model=lstm_example._LSTMModel(
num_features=5,
num_units=128,
exogenous_feature_columns=exogenous_feature_columns),
optimizer=adam.AdamOptimizer(0.001),
config=estimator_lib.RunConfig(tf_random_seed=4),
state_manager=state_management.ChainingStateManager(),
head_type=ts_head_lib.OneShotPredictionHead,
model_dir=model_dir)
train_features = {
feature_keys.TrainEvalFeatures.TIMES:
numpy.arange(20, dtype=numpy.int64),
feature_keys.TrainEvalFeatures.VALUES:
numpy.tile(numpy.arange(20, dtype=numpy.float32)[:, None], [1, 5]),
"2d_exogenous_feature":
numpy.ones([20, 2]),
"categorical_exogenous_feature":
numpy.array(["strkey"] * 20)[:, None]
}
train_input_fn = input_pipeline.RandomWindowInputFn(
input_pipeline.NumpyReader(train_features),
shuffle_seed=2,
num_threads=1,
batch_size=16,
window_size=16)
estimator.train(input_fn=train_input_fn, steps=5)
input_receiver_fn = estimator.build_raw_serving_input_receiver_fn()
export_location = estimator.export_savedmodel(self.get_temp_dir(),
input_receiver_fn)
graph = ops.Graph()
with graph.as_default():
with session_lib.Session() as session:
signatures = loader.load(session, [tag_constants.SERVING],
export_location)
self.assertEqual([feature_keys.SavedModelLabels.PREDICT],
list(signatures.signature_def.keys()))
predict_signature = signatures.signature_def[
feature_keys.SavedModelLabels.PREDICT]
six.assertCountEqual(self, [
feature_keys.FilteringFeatures.TIMES,
feature_keys.FilteringFeatures.VALUES,
"2d_exogenous_feature", "categorical_exogenous_feature"
], predict_signature.inputs.keys())
features = {
feature_keys.TrainEvalFeatures.TIMES:
numpy.tile(
numpy.arange(35, dtype=numpy.int64)[None, :], [2, 1]),
feature_keys.TrainEvalFeatures.VALUES:
numpy.tile(
numpy.arange(20, dtype=numpy.float32)[None, :, None],
[2, 1, 5]),
"2d_exogenous_feature":
numpy.ones([2, 35, 2]),
"categorical_exogenous_feature":
numpy.tile(
numpy.array(["strkey"] * 35)[None, :, None], [2, 1, 1])
}
feeds = {
graph.as_graph_element(input_value.name):
features[input_key]
for input_key, input_value in
predict_signature.inputs.items()
}
fetches = {
output_key: graph.as_graph_element(output_value.name)
for output_key, output_value in
predict_signature.outputs.items()
}
output = session.run(fetches, feed_dict=feeds)
self.assertAllEqual((2, 15, 5), output["mean"].shape)
# Name of the output placeholder
model_outputs = "map_1/TensorArrayStack/TensorArrayGatherV3:0"
# Run our session
with tf.Session() as sess:
# Print inputs and outputs (not required to run the script)
graph = tf.Graph()
with graph.as_default():
metagraph = tf.saved_model.loader.load(sess, [tag_constants.SERVING],
model_dir)
inputs_mapping = dict(metagraph.signature_def['serving_default'].inputs)
outputs_mapping = dict(metagraph.signature_def['serving_default'].outputs)
#print (inputs_mapping)
#print (outputs_mapping)
'''
Prepare the image to pass into the tf session. This will need to be improved
to allow for all images in a given directory to be read as bytes, appended to
a numpy array, then pass into the session.
'''
with open("StaM_6805_nf_duo_tripod_20171112_21_58_33p.JPG", "rb") as image:
f = image.read()
b = bytes(f)
#print (b)
# Run the inference
loader.load(sess, [tag_constants.SERVING], model_dir).signature_def
a = sess.run(model_outputs,
feed_dict={model_input: np.array(b).reshape(-1)})
print(a)
