Remove v1 test decorators from saved_model:loader_test.
The test exercises v1 APIs that nevertheless operate correctly under v2 behavior - the only difference for the tests is that the operations to run have different names in the graph. PiperOrigin-RevId: 324089971 Change-Id: I9dbb9ac23417b47fe55011c92292f0e89d21dc18
This commit is contained in:
Cesar Crusius
TensorFlower Gardener
parent
481957da0f
commit
e321e00f43
@ -26,9 +26,9 @@ from absl.testing import parameterized
|
||||
from tensorflow.python.client import session
|
||||
from tensorflow.python.framework import errors
|
||||
from tensorflow.python.framework import ops
|
||||
from tensorflow.python.framework import test_util
|
||||
from tensorflow.python.ops import control_flow_ops
|
||||
from tensorflow.python.ops import state_ops
|
||||
from tensorflow.python.ops import variable_scope
|
||||
from tensorflow.python.ops import variables
|
||||
from tensorflow.python.platform import test
|
||||
from tensorflow.python.saved_model import builder as saved_model_builder
|
||||
@ -41,6 +41,13 @@ from tensorflow.python.training import saver as tf_saver
|
||||
def _get_export_dir(label):
|
||||
return os.path.join(test.get_temp_dir(), label)
|
||||
|
||||
|
||||
def _tensor_name(name):
|
||||
if variable_scope.resource_variables_enabled():
|
||||
return name + "/Read/ReadVariableOp:0"
|
||||
return name + ":0"
|
||||
|
||||
|
||||
SIMPLE_ADD_SAVED_MODEL = _get_export_dir("simple_add_saved_model")
|
||||
SAVED_MODEL_WITH_MAIN_OP = _get_export_dir("saved_model_with_main_op")
|
||||
|
||||
@ -94,105 +101,116 @@ class SavedModelLoaderTest(test.TestCase, parameterized.TestCase):
|
||||
super(SavedModelLoaderTest, self).tearDown()
|
||||
shutil.rmtree(test.get_temp_dir(), ignore_errors=True)
|
||||
|
||||
@test_util.run_v1_only("b/120545219")
|
||||
def test_load_function(self, builder_cls):
|
||||
self.export_simple_graph(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
|
||||
self.assertEqual(11, sess.graph.get_tensor_by_name("y:0").eval())
|
||||
# Force test to run in graph mode.
|
||||
# The SaveModelLoader.load method is a v1-only API that requires a session
|
||||
# to work.
|
||||
with ops.Graph().as_default():
|
||||
self.export_simple_graph(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.run(_tensor_name("x")))
|
||||
self.assertEqual(11, sess.run(_tensor_name("y")))
|
||||
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader2 = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader2.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
|
||||
self.assertEqual(7, sess.graph.get_tensor_by_name("y:0").eval())
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader2 = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader2.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.run(_tensor_name("x")))
|
||||
self.assertEqual(7, sess.run(_tensor_name("y")))
|
||||
|
||||
@test_util.run_v1_only("b/120545219")
|
||||
def test_load_graph(self, builder_cls):
|
||||
self.export_simple_graph(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SIMPLE_ADD_SAVED_MODEL)
|
||||
graph = ops.Graph()
|
||||
loader.load_graph(graph, ["foo_graph"])
|
||||
|
||||
x = graph.get_tensor_by_name("x:0")
|
||||
y = graph.get_tensor_by_name("y:0")
|
||||
x = graph.get_tensor_by_name(_tensor_name("x"))
|
||||
y = graph.get_tensor_by_name(_tensor_name("y"))
|
||||
|
||||
with self.assertRaises(KeyError):
|
||||
graph.get_tensor_by_name("z:0")
|
||||
graph.get_tensor_by_name(_tensor_name("z"))
|
||||
|
||||
with self.session(graph=graph):
|
||||
with graph.as_default(), self.session():
|
||||
# Check that x and y are not initialized
|
||||
with self.assertRaises(errors.FailedPreconditionError):
|
||||
self.evaluate(x)
|
||||
with self.assertRaises(errors.FailedPreconditionError):
|
||||
self.evaluate(y)
|
||||
|
||||
@test_util.run_v1_only("b/120545219")
|
||||
def test_load_with_import_scope(self, builder_cls):
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
saver, _ = loader.load_graph(
|
||||
sess.graph, ["foo_graph"], import_scope="baz")
|
||||
# Force test to run in graph mode.
|
||||
# The SaveModelLoader.restore_variables ahd SaveModelLoader.run_init_ops
|
||||
# methods are v1-only APIs that require a session to work.
|
||||
with ops.Graph().as_default():
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
saver, _ = loader.load_graph(
|
||||
sess.graph, ["foo_graph"], import_scope="baz")
|
||||
|
||||
# The default saver should not work when the import scope is set.
|
||||
with self.assertRaises(errors.NotFoundError):
|
||||
loader.restore_variables(sess, tf_saver.Saver())
|
||||
|
||||
loader.restore_variables(sess, saver)
|
||||
|
||||
if builder_cls == saved_model_builder._SavedModelBuilder:
|
||||
# The default saver should not work when the import scope is set.
|
||||
with self.assertRaises(errors.NotFoundError):
|
||||
loader.restore_variables(sess, tf_saver.Saver())
|
||||
|
||||
loader.restore_variables(sess, saver)
|
||||
|
||||
if builder_cls == saved_model_builder._SavedModelBuilder:
|
||||
with self.assertRaises(errors.NotFoundError):
|
||||
loader.run_init_ops(sess, ["foo_graph"])
|
||||
loader.run_init_ops(sess, ["foo_graph"], import_scope="baz")
|
||||
else:
|
||||
loader.run_init_ops(sess, ["foo_graph"])
|
||||
loader.run_init_ops(sess, ["foo_graph"], import_scope="baz")
|
||||
else:
|
||||
loader.run_init_ops(sess, ["foo_graph"])
|
||||
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("baz/x:0").eval())
|
||||
self.assertEqual(7, sess.graph.get_tensor_by_name("baz/y:0").eval())
|
||||
self.assertEqual(5, sess.run(_tensor_name("baz/x")))
|
||||
self.assertEqual(7, sess.run(_tensor_name("baz/y")))
|
||||
|
||||
# Test combined load function.
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"], import_scope="baa")
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("baa/x:0").eval())
|
||||
self.assertEqual(7, sess.graph.get_tensor_by_name("baa/y:0").eval())
|
||||
# Test combined load function.
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"], import_scope="baa")
|
||||
self.assertEqual(5, sess.run(_tensor_name("baa/x")))
|
||||
self.assertEqual(7, sess.run(_tensor_name("baa/y")))
|
||||
|
||||
@test_util.run_deprecated_v1
|
||||
def test_restore_variables(self, builder_cls):
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
x = variables.VariableV1(0, name="x")
|
||||
y = variables.VariableV1(0, name="y")
|
||||
z = x * y
|
||||
# Force test to run in graph mode.
|
||||
# The SaveModelLoader.restore_variables method is a v1-only API requiring a
|
||||
# session to work.
|
||||
with ops.Graph().as_default():
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
with self.session() as sess:
|
||||
x = variables.VariableV1(0, name="x")
|
||||
y = variables.VariableV1(0, name="y")
|
||||
z = x * y
|
||||
|
||||
self.evaluate(variables.global_variables_initializer())
|
||||
self.evaluate(variables.global_variables_initializer())
|
||||
|
||||
# There are variables to restore, so a saver must be created.
|
||||
with self.assertRaises(ValueError):
|
||||
loader.restore_variables(sess, None)
|
||||
# There are variables to restore, so a saver must be created.
|
||||
with self.assertRaises(ValueError):
|
||||
loader.restore_variables(sess, None)
|
||||
|
||||
loader.restore_variables(sess, tf_saver.Saver())
|
||||
self.assertEqual(55, self.evaluate(z))
|
||||
loader.restore_variables(sess, tf_saver.Saver())
|
||||
self.assertEqual(55, self.evaluate(z))
|
||||
|
||||
@test_util.run_v1_only("b/120545219")
|
||||
def test_run_init_op(self, builder_cls):
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
graph = ops.Graph()
|
||||
saver, _ = loader.load_graph(graph, ["foo_graph"])
|
||||
with self.session(graph=graph) as sess:
|
||||
loader.restore_variables(sess, saver)
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
|
||||
self.assertEqual(11, sess.graph.get_tensor_by_name("y:0").eval())
|
||||
# Force test to run in graph mode.
|
||||
# The SaveModelLoader.restore_variables ahd SaveModelLoader.run_init_ops
|
||||
# methods are v1-only APIs that require a session to work.
|
||||
with ops.Graph().as_default():
|
||||
self.export_graph_with_main_op(builder_cls)
|
||||
loader = loader_impl.SavedModelLoader(SAVED_MODEL_WITH_MAIN_OP)
|
||||
graph = ops.Graph()
|
||||
saver, _ = loader.load_graph(graph, ["foo_graph"])
|
||||
with self.session(graph=graph) as sess:
|
||||
loader.restore_variables(sess, saver)
|
||||
self.assertEqual(5, sess.run(_tensor_name("x")))
|
||||
self.assertEqual(11, sess.run(_tensor_name("y")))
|
||||
|
||||
loader.run_init_ops(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
|
||||
self.assertEqual(7, sess.graph.get_tensor_by_name("y:0").eval())
|
||||
loader.run_init_ops(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.run(_tensor_name("x")))
|
||||
self.assertEqual(7, sess.run(_tensor_name("y")))
|
||||
|
||||
def test_parse_saved_model(self, builder_cls):
|
||||
self.export_simple_graph(builder_cls)
|
||||
@ -212,7 +230,6 @@ class SavedModelLoaderTest(test.TestCase, parameterized.TestCase):
|
||||
with self.assertRaises(RuntimeError):
|
||||
loader.get_meta_graph_def_from_tags(["not_a_graph"])
|
||||
|
||||
@test_util.run_v1_only("b/120545219")
|
||||
def test_load_saved_model_with_no_variables(self, builder_cls):
|
||||
"""Test that SavedModel runs saver when there appear to be no variables.
|
||||
|
||||
@ -224,36 +241,40 @@ class SavedModelLoaderTest(test.TestCase, parameterized.TestCase):
|
||||
Args:
|
||||
builder_cls: SavedModelBuilder or _SavedModelBuilder class
|
||||
"""
|
||||
path = _get_export_dir("no_variable_saved_model")
|
||||
with session.Session(graph=ops.Graph()) as sess:
|
||||
x = variables.VariableV1(
|
||||
5, name="x", collections=["not_global_variable"])
|
||||
y = variables.VariableV1(
|
||||
11, name="y", collections=["not_global_variable"])
|
||||
self.assertFalse(variables._all_saveable_objects())
|
||||
z = x + y
|
||||
self.evaluate(variables.variables_initializer([x, y]))
|
||||
# Force test to run in graph mode.
|
||||
# The SaveModelBuilder.add_meta_graph_and_variables and SaveModelLoader.load
|
||||
# methods are v1-only APIs that require a session to work.
|
||||
with ops.Graph().as_default():
|
||||
path = _get_export_dir("no_variable_saved_model")
|
||||
with session.Session(graph=ops.Graph()) as sess:
|
||||
x = variables.VariableV1(
|
||||
5, name="x", collections=["not_global_variable"])
|
||||
y = variables.VariableV1(
|
||||
11, name="y", collections=["not_global_variable"])
|
||||
self.assertFalse(variables._all_saveable_objects())
|
||||
z = x + y
|
||||
self.evaluate(variables.variables_initializer([x, y]))
|
||||
|
||||
foo_sig_def = signature_def_utils.build_signature_def(
|
||||
{"foo_input": utils.build_tensor_info(x)},
|
||||
{"foo_output": utils.build_tensor_info(z)})
|
||||
foo_sig_def = signature_def_utils.build_signature_def(
|
||||
{"foo_input": utils.build_tensor_info(x)},
|
||||
{"foo_output": utils.build_tensor_info(z)})
|
||||
|
||||
builder = saved_model_builder.SavedModelBuilder(path)
|
||||
builder.add_meta_graph_and_variables(
|
||||
sess, ["foo_graph"], {"foo": foo_sig_def},
|
||||
saver=tf_saver.Saver([x, y]))
|
||||
builder.save()
|
||||
builder = saved_model_builder.SavedModelBuilder(path)
|
||||
builder.add_meta_graph_and_variables(
|
||||
sess, ["foo_graph"], {"foo": foo_sig_def},
|
||||
saver=tf_saver.Saver([x, y]))
|
||||
builder.save()
|
||||
|
||||
loader = loader_impl.SavedModelLoader(path)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
saver, _ = loader.load_graph(sess.graph, ["foo_graph"])
|
||||
self.assertFalse(variables._all_saveable_objects())
|
||||
self.assertIsNotNone(saver)
|
||||
loader = loader_impl.SavedModelLoader(path)
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
saver, _ = loader.load_graph(sess.graph, ["foo_graph"])
|
||||
self.assertFalse(variables._all_saveable_objects())
|
||||
self.assertIsNotNone(saver)
|
||||
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.graph.get_tensor_by_name("x:0").eval())
|
||||
self.assertEqual(11, sess.graph.get_tensor_by_name("y:0").eval())
|
||||
with self.session(graph=ops.Graph()) as sess:
|
||||
loader.load(sess, ["foo_graph"])
|
||||
self.assertEqual(5, sess.run(_tensor_name("x")))
|
||||
self.assertEqual(11, sess.run(_tensor_name("y")))
|
||||
|
||||
def test_load_saved_model_graph_with_return_elements(self, builder_cls):
|
||||
"""Ensure that the correct elements are returned."""
|
||||
|
||||
Loading…
Reference in New Issue
Block a user