wehub-resource-sync
1581 lines
58 KiB
Python
1581 lines
58 KiB
Python
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
# ==============================================================================
|
|
"""Tests for trackable object SavedModel save."""
|
|
|
|
import os
|
|
|
|
from absl.testing import parameterized
|
|
|
|
from google.protobuf import text_format
|
|
from tensorflow.core.config import flags
|
|
from tensorflow.core.framework import attr_value_pb2
|
|
from tensorflow.core.framework import function_pb2
|
|
from tensorflow.core.framework import graph_debug_info_pb2
|
|
from tensorflow.core.framework import graph_pb2
|
|
from tensorflow.core.framework import node_def_pb2
|
|
from tensorflow.core.protobuf import meta_graph_pb2
|
|
from tensorflow.python.checkpoint import checkpoint
|
|
from tensorflow.python.checkpoint.sharding import sharding_policies
|
|
from tensorflow.python.client import session as session_lib
|
|
from tensorflow.python.data.ops import dataset_ops
|
|
from tensorflow.python.distribute import mirrored_strategy
|
|
from tensorflow.python.eager import backprop
|
|
from tensorflow.python.eager import context
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import remote
|
|
from tensorflow.python.eager import test
|
|
from tensorflow.python.framework import constant_op
|
|
from tensorflow.python.framework import dtypes
|
|
from tensorflow.python.framework import meta_graph
|
|
from tensorflow.python.framework import ops
|
|
from tensorflow.python.framework import tensor_spec
|
|
from tensorflow.python.framework import test_util
|
|
from tensorflow.python.framework import versions
|
|
from tensorflow.python.lib.io import file_io
|
|
from tensorflow.python.module import module
|
|
from tensorflow.python.ops import array_ops
|
|
from tensorflow.python.ops import control_flow_switch_case
|
|
from tensorflow.python.ops import io_ops
|
|
from tensorflow.python.ops import lookup_ops
|
|
from tensorflow.python.ops import math_ops
|
|
from tensorflow.python.ops import resource_variable_ops
|
|
from tensorflow.python.ops import variables
|
|
from tensorflow.python.ops.ragged import ragged_factory_ops
|
|
from tensorflow.python.ops.ragged import ragged_tensor
|
|
from tensorflow.python.platform import gfile
|
|
from tensorflow.python.saved_model import load
|
|
from tensorflow.python.saved_model import loader
|
|
from tensorflow.python.saved_model import loader_impl
|
|
from tensorflow.python.saved_model import save
|
|
from tensorflow.python.saved_model import save_options
|
|
from tensorflow.python.saved_model import signature_constants
|
|
from tensorflow.python.saved_model import tag_constants
|
|
from tensorflow.python.trackable import asset
|
|
from tensorflow.python.trackable import autotrackable
|
|
from tensorflow.python.training import saver
|
|
from tensorflow.python.training import server_lib
|
|
from tensorflow.python.util import compat
|
|
|
|
|
|
def _run_signature(
|
|
session,
|
|
meta_graph_def,
|
|
inputs,
|
|
signature_key,
|
|
disable_check_for_input_signature_size_match=False,
|
|
):
|
|
signature = meta_graph_def.signature_def[signature_key]
|
|
if not disable_check_for_input_signature_size_match:
|
|
assert set(inputs.keys()) == set(signature.inputs.keys())
|
|
feed_dict = {}
|
|
for arg_name in inputs.keys():
|
|
input_tensor = session.graph.get_tensor_by_name(
|
|
signature.inputs[arg_name].name)
|
|
feed_dict[input_tensor] = inputs[arg_name]
|
|
output_dict = {}
|
|
for output_name, output_tensor_info in signature.outputs.items():
|
|
output_dict[output_name] = session.graph.get_tensor_by_name(
|
|
output_tensor_info.name)
|
|
return session.run(output_dict, feed_dict=feed_dict)
|
|
|
|
|
|
def _import_and_infer(
|
|
save_dir,
|
|
inputs,
|
|
signature_key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY,
|
|
disable_check_for_input_signature_size_match=False,
|
|
):
|
|
"""Import a SavedModel into a TF 1.x-style graph and run `signature_key`."""
|
|
graph = ops.Graph()
|
|
with graph.as_default(), session_lib.Session() as session:
|
|
model = loader.load(session, [tag_constants.SERVING], save_dir)
|
|
return _run_signature(
|
|
session,
|
|
model,
|
|
inputs,
|
|
signature_key,
|
|
disable_check_for_input_signature_size_match,
|
|
)
|
|
|
|
|
|
class SaveTest(test.TestCase, parameterized.TestCase):
|
|
|
|
def test_method_save_signature(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: 2. * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
root.f(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, root.f)
|
|
self.assertEqual({"output_0": 2.}, _import_and_infer(save_dir, {"x": 1.}))
|
|
|
|
def test_method_save_list_func(self):
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@def_function.function
|
|
def case_fn(x):
|
|
branch_index = constant_op.constant(1)
|
|
branches = [lambda: x, lambda: x + 1]
|
|
case_out = control_flow_switch_case.switch_case(branch_index, branches)
|
|
return case_out
|
|
|
|
root.f = def_function.function(
|
|
lambda x: 2. * case_fn(x),
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
root.f(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, root.f)
|
|
self.assertEqual({"output_0": 4.}, _import_and_infer(save_dir, {"x": 1.}))
|
|
|
|
def test_method_save_concrete(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(lambda z: {"out": 2. * z})
|
|
root.f(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root, save_dir, {
|
|
"non_default_key":
|
|
root.f.get_concrete_function(
|
|
tensor_spec.TensorSpec(None, dtypes.float32))
|
|
})
|
|
self.assertEqual({"out": 2.},
|
|
_import_and_infer(
|
|
save_dir, {"z": 1.}, signature_key="non_default_key"))
|
|
|
|
def test_method_save_annotated_function(self):
|
|
# This test is only meaningful with Python 3 because Python 2's
|
|
# inspect.getargspec doesn't save annotations.
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
class UnknownType(object): # pylint: disable=unused-variable
|
|
pass
|
|
|
|
def annotated_function(z):
|
|
return {"out": 2. * z}
|
|
|
|
# Same effect as annotating function like the following.
|
|
# def annotated_function("z": UnknownType) -> UnknownType:
|
|
# This is a workaround since Python 2 does not support annotations and
|
|
# our presubmit linter catches it.
|
|
annotated_function.__annotations__ = {
|
|
"z": UnknownType,
|
|
"return": UnknownType
|
|
}
|
|
|
|
root.f = def_function.function(annotated_function)
|
|
root.f(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root, save_dir, {
|
|
"non_default_key":
|
|
root.f.get_concrete_function(
|
|
tensor_spec.TensorSpec(None, dtypes.float32))
|
|
})
|
|
self.assertEqual({"out": 2.},
|
|
_import_and_infer(
|
|
save_dir, {"z": 1.}, signature_key="non_default_key"))
|
|
|
|
def test_method_save_defaults(self):
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec([], dtypes.float32)]
|
|
)
|
|
def f(x, y=constant_op.constant(5.0)):
|
|
return x + y
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec([], dtypes.float32)]
|
|
)
|
|
def g(x=constant_op.constant(10.0), y=constant_op.constant(20.0)):
|
|
return x + y
|
|
|
|
root = module.Module()
|
|
root.f = f
|
|
root.g = g
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, {"f": root.f, "g": root.g})
|
|
|
|
self.assertEqual(
|
|
{"output_0": 7.0},
|
|
_import_and_infer(
|
|
save_dir,
|
|
inputs={"x": 2.0},
|
|
signature_key="f",
|
|
disable_check_for_input_signature_size_match=True,
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
{"output_0": 30.0},
|
|
_import_and_infer(
|
|
save_dir,
|
|
inputs={},
|
|
signature_key="g",
|
|
disable_check_for_input_signature_size_match=True,
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
{"output_0": 15.0},
|
|
_import_and_infer(
|
|
save_dir,
|
|
inputs={"y": 5.0},
|
|
signature_key="g",
|
|
disable_check_for_input_signature_size_match=True,
|
|
),
|
|
)
|
|
|
|
def test_save_defaults_dict(self):
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@def_function.function(
|
|
input_signature=[{
|
|
"temperature": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.float32, name="d"
|
|
),
|
|
"per_example_max_decode_steps": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32, name="c"
|
|
),
|
|
"per_example_top_k": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32, name="b"
|
|
),
|
|
"gumbel_prng_key": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32, name="a"
|
|
),
|
|
}]
|
|
)
|
|
def f(
|
|
x={
|
|
"temperature": constant_op.constant(0.5),
|
|
"per_example_max_decode_steps": constant_op.constant(1024),
|
|
"per_example_top_k": constant_op.constant(40),
|
|
"gumbel_prng_key": constant_op.constant(0),
|
|
}
|
|
): # pylint: disable=dangerous-default-value
|
|
return x
|
|
|
|
root.f = f
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, root.f)
|
|
|
|
self.assertEqual(
|
|
{
|
|
"gumbel_prng_key": 0,
|
|
"per_example_max_decode_steps": 1024,
|
|
"temperature": 0.5,
|
|
"per_example_top_k": 40,
|
|
},
|
|
_import_and_infer(
|
|
save_dir, {}, disable_check_for_input_signature_size_match=True
|
|
),
|
|
)
|
|
|
|
def test_save_defaults_nested_structure(self):
|
|
root = autotrackable.AutoTrackable()
|
|
|
|
@def_function.function(
|
|
input_signature=[
|
|
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name="m"),
|
|
[
|
|
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
|
|
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
|
|
{
|
|
"temperature": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.float32
|
|
),
|
|
"per_example_max_decode_steps": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32
|
|
),
|
|
"per_example_top_k": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32
|
|
),
|
|
"gumbel_prng_key": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.int32
|
|
),
|
|
"dict_entry": {
|
|
"a": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.float32
|
|
),
|
|
"b": tensor_spec.TensorSpec(
|
|
shape=(), dtype=dtypes.float32
|
|
),
|
|
},
|
|
},
|
|
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32),
|
|
],
|
|
tensor_spec.TensorSpec(shape=(), dtype=dtypes.float32, name="y"),
|
|
]
|
|
)
|
|
def f(
|
|
m,
|
|
x=[
|
|
constant_op.constant(5.0),
|
|
constant_op.constant(1.0),
|
|
{
|
|
"tempurature": constant_op.constant(0.5),
|
|
"per_example_max_decode_steps": constant_op.constant(1024),
|
|
"per_example_top_k": constant_op.constant(40),
|
|
"gumbel_prng_key": constant_op.constant(0),
|
|
"dict_entry": {
|
|
"a": constant_op.constant(1.0),
|
|
"b": constant_op.constant(2.0),
|
|
},
|
|
},
|
|
constant_op.constant(3.0),
|
|
],
|
|
y=constant_op.constant(2.0),
|
|
): # pylint: disable=dangerous-default-value
|
|
return m + x[2]["dict_entry"]["a"] + x[3] + y
|
|
|
|
root.f = f
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, {"f": root.f})
|
|
|
|
self.assertEqual(
|
|
{"output_0": 8.0},
|
|
_import_and_infer(
|
|
save_dir,
|
|
inputs={"m": 2.0},
|
|
signature_key="f",
|
|
disable_check_for_input_signature_size_match=True,
|
|
),
|
|
)
|
|
|
|
self.assertEqual(
|
|
{"output_0": 10.0},
|
|
_import_and_infer(
|
|
save_dir,
|
|
inputs={"m": 2.0, "y": 4.0},
|
|
signature_key="f",
|
|
disable_check_for_input_signature_size_match=True,
|
|
),
|
|
)
|
|
|
|
def test_unsaveable_func_graph(self):
|
|
root = module.Module()
|
|
|
|
@def_function.function(input_signature=[])
|
|
def nested_f():
|
|
ops.get_default_graph().mark_as_unsaveable("ERROR MSG")
|
|
return 1
|
|
|
|
@def_function.function(input_signature=[])
|
|
def f():
|
|
return nested_f()
|
|
|
|
root.f = f
|
|
with self.assertRaisesRegex(ValueError, "ERROR MSG"):
|
|
save.save(root, os.path.join(self.get_temp_dir(), "saved_model"))
|
|
|
|
def test_untracked_variable_useful_message(self):
|
|
root = module.Module()
|
|
v = variables.Variable(1., name="some_unique_name")
|
|
|
|
@def_function.function(input_signature=[])
|
|
def f():
|
|
return v.read_value()
|
|
|
|
root.f = f
|
|
with self.assertRaisesRegex(
|
|
AssertionError, "Trackable referencing this tensor.*some_unique_name"):
|
|
save.save(root, os.path.join(self.get_temp_dir(), "saved_model"))
|
|
|
|
def test_version_information_included(self):
|
|
root = autotrackable.AutoTrackable()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir)
|
|
saved_model_proto = loader_impl.parse_saved_model(save_dir)
|
|
self.assertEqual(
|
|
versions.__version__,
|
|
saved_model_proto.meta_graphs[0].meta_info_def.tensorflow_version)
|
|
self.assertEqual(
|
|
versions.__git_version__,
|
|
saved_model_proto.meta_graphs[0].meta_info_def.tensorflow_git_version)
|
|
|
|
def test_non_concrete_error(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(lambda x: 2. * x)
|
|
root.f(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError, "Expected a TensorFlow function"):
|
|
save.save(root, save_dir, root.f)
|
|
|
|
def test_captures_unreachable_variable(self):
|
|
root = autotrackable.AutoTrackable()
|
|
unreachable_variable = variables.Variable([5.0, 2.0])
|
|
root.reachable_variable = variables.Variable([1.0, 3.0])
|
|
|
|
@def_function.function
|
|
def increase_variable(x):
|
|
return 2 * unreachable_variable * x + root.reachable_variable
|
|
|
|
root.f = increase_variable
|
|
|
|
self.assertAllEqual([101.0, 83.0],
|
|
root.f(constant_op.constant([10.0, 20.0])).numpy())
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
|
|
with self.assertRaisesRegex(KeyError, "not reachable from root"):
|
|
save.save(root, save_dir)
|
|
|
|
def test_nested_inputs(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: 2. * x[0],
|
|
input_signature=([
|
|
tensor_spec.TensorSpec(None, dtypes.float32),
|
|
tensor_spec.TensorSpec(None, dtypes.float32)
|
|
],))
|
|
root.f([constant_op.constant(1.), constant_op.constant(1.)])
|
|
|
|
def test_nested_outputs(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(lambda x: (2. * x, (3. * x, 4. * x)))
|
|
root.f(constant_op.constant(1.))
|
|
to_save = root.f.get_concrete_function(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError, "non-Tensor value"):
|
|
save.save(root, save_dir, to_save)
|
|
|
|
def test_nested_dict_outputs(self):
|
|
root = checkpoint.Checkpoint(
|
|
f=def_function.function(lambda x: { # pylint: disable=g-long-lambda
|
|
"a": 2. * x,
|
|
"b": (3. * x, 4. * x)
|
|
}))
|
|
root.f(constant_op.constant(1.))
|
|
to_save = root.f.get_concrete_function(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError, "non-Tensor value"):
|
|
save.save(root, save_dir, to_save)
|
|
|
|
def test_variable(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.v1 = variables.Variable(3.)
|
|
root.v2 = variables.Variable(2.)
|
|
root.f = def_function.function(lambda x: root.v1 * root.v2 * x)
|
|
root.f(constant_op.constant(1.))
|
|
to_save = root.f.get_concrete_function(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir, to_save)
|
|
self.assertAllEqual({"output_0": 12.},
|
|
_import_and_infer(save_dir, {"x": 2.}))
|
|
|
|
def test_single_function_default_signature(self):
|
|
model = autotrackable.AutoTrackable()
|
|
model.f = def_function.function(lambda: 3., input_signature=())
|
|
model.f()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(model, save_dir)
|
|
self.assertAllClose({"output_0": 3.}, _import_and_infer(save_dir, {}))
|
|
|
|
def test_single_function_no_signature(self):
|
|
model = autotrackable.AutoTrackable()
|
|
model.f = def_function.function(lambda: 3.)
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(model, save_dir)
|
|
|
|
def test_save_function_no_trace(self):
|
|
|
|
class ObjWithFunction(module.Module):
|
|
|
|
@def_function.function
|
|
def foo(self, a):
|
|
return a
|
|
|
|
@def_function.function
|
|
def bar(self, a):
|
|
return a + 1
|
|
|
|
root = ObjWithFunction()
|
|
root.bar(1)
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertLogs(level="INFO") as logs:
|
|
save.save(root, save_dir)
|
|
|
|
expected_message = (
|
|
"INFO:absl:Found untraced functions such as foo while saving "
|
|
"(showing 1 of 1). These functions will not be directly callable after "
|
|
"loading.")
|
|
self.assertIn(expected_message, logs.output)
|
|
|
|
def test_find_default_save_function(self):
|
|
|
|
class ObjWithDefaultSignature(checkpoint.Checkpoint):
|
|
|
|
@def_function.function(input_signature=[
|
|
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)
|
|
])
|
|
def _default_save_signature(self, x):
|
|
return x + x + 1
|
|
|
|
obj = ObjWithDefaultSignature()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(obj, save_dir)
|
|
self.assertAllClose({"output_0": 7.},
|
|
_import_and_infer(save_dir, {"x": 3.}))
|
|
|
|
def test_docstring(self):
|
|
|
|
class Adder(module.Module):
|
|
|
|
@def_function.function(input_signature=[
|
|
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)
|
|
])
|
|
def add(self, x):
|
|
return x + x + 1.
|
|
|
|
to_save = Adder()
|
|
to_save.add(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(to_save, save_dir)
|
|
self.assertAllClose({"output_0": 7.},
|
|
_import_and_infer(save_dir, {"x": 3.}))
|
|
|
|
def test_datastructures(self):
|
|
|
|
class HasDatastructures(checkpoint.Checkpoint):
|
|
|
|
def __init__(self):
|
|
self.a = [1.]
|
|
self.a.append(variables.Variable(2.))
|
|
self.b = {"a": variables.Variable(3.)}
|
|
|
|
@def_function.function(input_signature=[
|
|
tensor_spec.TensorSpec(shape=None, dtype=dtypes.float32)
|
|
])
|
|
def add(self, x):
|
|
return x + math_ops.add_n(self.a) + self.b["a"]
|
|
|
|
to_save = HasDatastructures()
|
|
to_save.add(constant_op.constant(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(to_save, save_dir)
|
|
self.assertAllClose({"output_0": 10.},
|
|
_import_and_infer(save_dir, {"x": 4.}))
|
|
|
|
def test_default_attr_stripping(self):
|
|
|
|
class Complex(checkpoint.Checkpoint):
|
|
|
|
@def_function.function(input_signature=[])
|
|
def __call__(self):
|
|
return math_ops.complex(
|
|
constant_op.constant(1.), constant_op.constant(2.), name="complex")
|
|
|
|
to_save = Complex()
|
|
to_save()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(to_save, save_dir)
|
|
graph = ops.Graph()
|
|
with graph.as_default(), self.session(graph) as session:
|
|
loader.load(session, [tag_constants.SERVING], save_dir)
|
|
func, = [f for name, f in graph._functions.items() if "call" in name]
|
|
complex_node, = [
|
|
node for node in func.definition.node_def if node.op == "Complex"
|
|
]
|
|
self.assertNotIn("T", complex_node.attr)
|
|
self.assertNotIn("Tout", complex_node.attr)
|
|
|
|
def test_signature_attribute_reserved(self):
|
|
root = checkpoint.Checkpoint(signatures=variables.Variable(1.))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError, "del obj.signatures"):
|
|
save.save(root, save_dir)
|
|
del root.signatures
|
|
save.save(root, save_dir)
|
|
|
|
def test_function_with_captured_dataset(self):
|
|
if test_util.is_gpu_available():
|
|
self.skipTest("Currently broken when a GPU is available.")
|
|
|
|
class HasDataset(module.Module):
|
|
|
|
def __init__(self):
|
|
super(HasDataset, self).__init__()
|
|
self.dataset = (dataset_ops.Dataset.range(5).map(lambda x: x**2))
|
|
|
|
@def_function.function
|
|
def __call__(self, x):
|
|
current_sum = array_ops.zeros([], dtype=dtypes.int64)
|
|
for element in self.dataset:
|
|
current_sum += x * element
|
|
return current_sum
|
|
|
|
root = HasDataset()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root,
|
|
save_dir,
|
|
signatures=root.__call__.get_concrete_function(
|
|
tensor_spec.TensorSpec(None, dtypes.int64)))
|
|
self.assertAllClose({"output_0": 3 * (1 + 4 + 9 + 16)},
|
|
_import_and_infer(save_dir, {"x": 3}))
|
|
|
|
def test_variable_args_cannot_be_used_as_signature(self):
|
|
|
|
with self.assertRaises(TypeError):
|
|
@def_function.function(input_signature=[
|
|
resource_variable_ops.VariableSpec(shape=[], dtype=dtypes.int32)
|
|
])
|
|
def f(unused_v):
|
|
return 1
|
|
|
|
def test_export_correct_output_shapes(self):
|
|
"""Asserts that nodes are exported with the correct number of output shapes.
|
|
|
|
After backpropagation rewrite, functions are rewritten with additional
|
|
outputs. When exporting to SavedModel, the shapes of the additional outputs
|
|
were incorrectly added to the FunctionDef proto (b/133666530).
|
|
"""
|
|
obj = autotrackable.AutoTrackable()
|
|
obj.v = variables.Variable(2.)
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
def f(x):
|
|
return (math_ops.multiply(obj.v, x), math_ops.multiply(obj.v,
|
|
(x + 1)), None)
|
|
|
|
obj.f = f
|
|
|
|
@def_function.function(
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
def g(x):
|
|
return obj.f(x)[1]
|
|
|
|
obj.g = g
|
|
|
|
# After the following lines, the concrete functions of obj.g and obj.f are
|
|
# rewritten with many extra outputs.
|
|
with backprop.GradientTape():
|
|
obj.g(constant_op.constant(3.0))
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(obj, save_dir, signatures={"g": obj.g})
|
|
graph_def = loader_impl.parse_saved_model(save_dir).meta_graphs[0].graph_def
|
|
|
|
def assert_correct_number_of_output_shapes(node):
|
|
if node.op == "StatefulPartitionedCall":
|
|
fn_name = node.attr["f"].func.name
|
|
if fn_name.startswith("__inference_f"):
|
|
self.assertLen(node.attr["_output_shapes"].list.shape, 2)
|
|
if fn_name.startswith("__inference_g"):
|
|
self.assertLen(node.attr["_output_shapes"].list.shape, 1)
|
|
|
|
for f in graph_def.library.function:
|
|
if (f.signature.name.startswith("__inference_f") or
|
|
f.signature.name.startswith("__inference_g")):
|
|
for node in f.node_def:
|
|
assert_correct_number_of_output_shapes(node)
|
|
|
|
def test_save_cached_variable(self):
|
|
with ops.Graph().as_default(), session_lib.Session() as session:
|
|
obj = autotrackable.AutoTrackable()
|
|
obj.v = variables.Variable(2., caching_device=lambda op: op.device)
|
|
obj.w = variables.Variable(3.)
|
|
session.run([obj.v.initializer, obj.w.initializer])
|
|
|
|
@def_function.function(input_signature=[])
|
|
def f():
|
|
return obj.v + obj.w
|
|
|
|
obj.f = f
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(obj, save_dir, signatures=obj.f)
|
|
self.assertAllClose({"output_0": 5}, _import_and_infer(save_dir, {}))
|
|
|
|
@parameterized.named_parameters(
|
|
("_SaveDevices_ExportMetaGraph",
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES, True),
|
|
("_DiscardDevices_ExportMetaGraph", save_options.VariablePolicy.NONE,
|
|
True), ("_SaveDevices_Save",
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES, False),
|
|
("_DiscardDevices_Save", save_options.VariablePolicy.NONE, False))
|
|
def test_save_variable_devices(self, save_devices, meta_graph_only):
|
|
context._reset_context()
|
|
cpus = context.context().list_physical_devices("CPU")
|
|
if len(cpus) == 1:
|
|
context.context().set_logical_device_configuration(
|
|
cpus[0], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
context.ensure_initialized()
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
with ops.device("CPU:0"):
|
|
root.v0 = variables.Variable(1., name="v0")
|
|
with ops.device("CPU:1"):
|
|
root.v1 = variables.Variable(1., name="v1")
|
|
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy=save_devices)
|
|
file_name = os.path.join(self.get_temp_dir(), "saved_model")
|
|
if meta_graph_only:
|
|
save.export_meta_graph(obj=root, filename=file_name, options=options)
|
|
else:
|
|
save.save(obj=root, export_dir=file_name, options=options)
|
|
|
|
meta = None
|
|
if meta_graph_only:
|
|
meta = meta_graph.read_meta_graph_file(file_name)
|
|
else:
|
|
meta = loader_impl.parse_saved_model(file_name).meta_graphs[0]
|
|
|
|
# Check devices in meta graph nodes.
|
|
graph_def = meta.graph_def
|
|
v0 = next((n for n in graph_def.node if n.name == "v0"), None)
|
|
v1 = next((n for n in graph_def.node if n.name == "v1"), None)
|
|
self.assertIsNotNone(v0)
|
|
self.assertIsNotNone(v1)
|
|
if save_devices == save_options.VariablePolicy.SAVE_VARIABLE_DEVICES:
|
|
self.assertIn("CPU:0", v0.device)
|
|
self.assertIn("CPU:1", v1.device)
|
|
else:
|
|
self.assertEmpty(v0.device)
|
|
self.assertEmpty(v1.device)
|
|
|
|
# Check devices in object graph nodes.
|
|
object_graph_def = meta.object_graph_def
|
|
v0 = next((n.variable
|
|
for n in object_graph_def.nodes
|
|
if n.HasField("variable") and n.variable.name == "v0"), None)
|
|
v1 = next((n.variable
|
|
for n in object_graph_def.nodes
|
|
if n.HasField("variable") and n.variable.name == "v1"), None)
|
|
self.assertIsNotNone(v0)
|
|
self.assertIsNotNone(v1)
|
|
if save_devices == save_options.VariablePolicy.SAVE_VARIABLE_DEVICES:
|
|
self.assertIn("CPU:0", v0.device)
|
|
self.assertIn("CPU:1", v1.device)
|
|
else:
|
|
self.assertEmpty(v0.device)
|
|
self.assertEmpty(v1.device)
|
|
|
|
@parameterized.named_parameters(
|
|
("_ExpandDistributedVariablesWithPolicy",
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES, True),
|
|
("_ExpandDistributedVariablesWithoutPolicy",
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES, False),
|
|
("_DiscardDistributedVariablesWithPolicy",
|
|
save_options.VariablePolicy.NONE, True),
|
|
("_DiscardDistributedVariablesWithoutPolicy",
|
|
save_options.VariablePolicy.NONE, False))
|
|
def test_expand_distributed_variables(self, expand_strategy, policy):
|
|
# 1. Create a context with both CPU:0 and CPU:1.
|
|
context._reset_context()
|
|
cpus = context.context().list_physical_devices("CPU")
|
|
if len(cpus) == 1:
|
|
context.context().set_logical_device_configuration(
|
|
cpus[0], [
|
|
context.LogicalDeviceConfiguration(),
|
|
context.LogicalDeviceConfiguration()
|
|
])
|
|
context.ensure_initialized()
|
|
|
|
# 2. Create and save a model under a mirrored strategy.
|
|
file_name = os.path.join(self.get_temp_dir(), "saved_model.pb")
|
|
strategy = mirrored_strategy.MirroredStrategy(["CPU:0", "CPU:1"])
|
|
strategy.extended._use_var_policy = policy
|
|
with strategy.scope():
|
|
root = autotrackable.AutoTrackable()
|
|
root.v = variables.Variable([1., 1.], name="v")
|
|
|
|
@def_function.function(input_signature=[])
|
|
def f():
|
|
root.v.assign([2., 2.])
|
|
|
|
root.f = f
|
|
|
|
save.export_meta_graph(
|
|
obj=root,
|
|
filename=file_name,
|
|
options=save_options.SaveOptions(
|
|
experimental_variable_policy=expand_strategy))
|
|
|
|
# 3. Read the output file and test behavior.
|
|
meta_graph_def = meta_graph.read_meta_graph_file(file_name)
|
|
object_graph = meta_graph_def.object_graph_def
|
|
graph_def = meta_graph_def.graph_def
|
|
v = next((n.variable
|
|
for n in object_graph.nodes
|
|
if n.HasField("variable") and n.variable.name == "v"), None)
|
|
saved_function = next((f for f in graph_def.library.function
|
|
if "inference_f_" in f.signature.name), None)
|
|
self.assertIsNotNone(saved_function)
|
|
if (expand_strategy ==
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES):
|
|
# experimental_save_variable_devices should have been automatically set.
|
|
self.assertIn("CPU:0", v.device)
|
|
components = v.experimental_distributed_variable_components
|
|
self.assertLen(components, 2)
|
|
v0 = next((x for x in components if x.name == "v"), None)
|
|
v1 = next((x for x in components if x.name == "v/replica_1"), None)
|
|
self.assertIsNotNone(v0)
|
|
self.assertIsNotNone(v1)
|
|
self.assertIn("CPU:0", v0.device)
|
|
self.assertIn("CPU:1", v1.device)
|
|
self.assertLen(saved_function.signature.input_arg, 2)
|
|
else:
|
|
self.assertEmpty(v.device)
|
|
self.assertEmpty(v.experimental_distributed_variable_components)
|
|
self.assertLen(saved_function.signature.input_arg, 1)
|
|
|
|
def test_save_uninitialized_variable(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.uninitialized_variable = resource_variable_ops.UninitializedVariable(
|
|
name="uninitialized_variable", dtype=dtypes.float32)
|
|
root.initialized_variable = variables.Variable(
|
|
1.0, name="initialized_variable")
|
|
|
|
# TODO(b/149594077): Python loading does not work now partly because it
|
|
# shouldn't, as the public API and semantics of uninitialized variables
|
|
# are not properly defined, and officially supporting loading would end up
|
|
# defining semantics "by usage." We should only allow loading once the API
|
|
# is made official.
|
|
export_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, export_dir)
|
|
with self.assertRaisesRegex(FileNotFoundError,
|
|
"Key uninitialized_variable"):
|
|
load.load(export_dir)
|
|
with ops.Graph().as_default(), session_lib.Session() as session:
|
|
# The final ValueError here (with "no variables to save") is confusing,
|
|
# but errors upstream give the user the correct information (a
|
|
# NotFoundError stating that the uninitalized_variable was not found in
|
|
# the checkpoint).
|
|
with self.assertRaises(ValueError):
|
|
loader.load(session, [tag_constants.SERVING], export_dir)
|
|
|
|
def test_concrete_function_with_set_shape(self,):
|
|
# Serialized concrete function should retain the shape from the TensorSpec,
|
|
# instead of using the shape of the inputs (which are changed by set_shape).
|
|
@def_function.function
|
|
def f(x):
|
|
x.set_shape((5, 1))
|
|
return x
|
|
|
|
root = autotrackable.AutoTrackable()
|
|
path = os.path.join(self.get_temp_dir(), "saved_model")
|
|
concrete = f.get_concrete_function(
|
|
tensor_spec.TensorSpec((None, 1), name="name"))
|
|
save.save(root, path, signatures={"key": concrete})
|
|
imported = load.load(path)
|
|
self.assertEqual(imported.signatures["key"].structured_input_signature[1],
|
|
{"name": tensor_spec.TensorSpec((None, 1), name="name")})
|
|
|
|
def test_save_composite_tensor_signature(self):
|
|
@def_function.function(
|
|
input_signature=[ragged_tensor.RaggedTensorSpec(ragged_rank=2)])
|
|
def f(x):
|
|
return {"output_key": x}
|
|
root = autotrackable.AutoTrackable()
|
|
path = os.path.join(self.get_temp_dir(), "saved_model")
|
|
inp = ragged_factory_ops.constant([[[1.0, 2.0], [3.0]], [[5.]]])
|
|
flat_inp = {
|
|
"x": constant_op.constant([1., 2., 3., 5]),
|
|
"x_1": constant_op.constant([0, 2, 3], dtype=dtypes.int64),
|
|
"x_2": constant_op.constant([0, 2, 3, 4], dtype=dtypes.int64)
|
|
}
|
|
save.save(root, path, signatures={"key": f.get_concrete_function()})
|
|
|
|
# Test that the ragged signature can be loaded back into Python with V2 APIs
|
|
imported = load.load(path)
|
|
self.assertAllEqual(inp,
|
|
imported.signatures["key"](**flat_inp)["output_key"])
|
|
graph = ops.Graph()
|
|
|
|
# Try running the signature with V1 APIs.
|
|
with graph.as_default(), session_lib.Session() as session:
|
|
meta_graph_def = loader.load(session, [tag_constants.SERVING], path)
|
|
signature = meta_graph_def.signature_def["key"]
|
|
|
|
feed_dict = {}
|
|
for arg_name in flat_inp:
|
|
input_tensor = session.graph.get_tensor_by_name(
|
|
signature.inputs[arg_name].name)
|
|
feed_dict[input_tensor] = flat_inp[arg_name].numpy()
|
|
|
|
# Get composite tensor components
|
|
output_components = (
|
|
signature.outputs["output_key"].composite_tensor.components)
|
|
fetches = {}
|
|
components_keys = ["x", "x_1", "x_2"]
|
|
for k, output_tensor_info in zip(components_keys, output_components):
|
|
fetches[k] = session.graph.get_tensor_by_name(output_tensor_info.name)
|
|
|
|
outputs = session.run(fetches, feed_dict)
|
|
|
|
self.assertAllClose(flat_inp, outputs)
|
|
|
|
def test_save_uses_sanitized_signature_name(self):
|
|
|
|
@def_function.function(
|
|
input_signature=[ragged_tensor.RaggedTensorSpec(ragged_rank=2)])
|
|
def f(x):
|
|
return {"output_key": x}
|
|
|
|
# Colons are not usable as name scopes.
|
|
unsanitized_name = "foo:bar"
|
|
root = autotrackable.AutoTrackable()
|
|
path = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root, path, signatures={unsanitized_name: f.get_concrete_function()})
|
|
graph = ops.Graph()
|
|
with graph.as_default(), session_lib.Session() as session:
|
|
meta_graph_def = loader.load(session, [tag_constants.SERVING], path)
|
|
signature = meta_graph_def.signature_def[unsanitized_name]
|
|
tensor_names = [
|
|
session.graph.get_tensor_by_name(signature.inputs[key].name).name
|
|
for key in signature.inputs
|
|
]
|
|
# The placeholder names will have the sanitized version.
|
|
self.assertCountEqual(tensor_names,
|
|
["foo_bar_x:0", "foo_bar_x_1:0", "foo_bar_x_2:0"])
|
|
|
|
def test_save_returns_none(self):
|
|
# Test that `tf.saved_model.save` API returns None to user.
|
|
root = autotrackable.AutoTrackable()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
result = save.save(root, save_dir)
|
|
self.assertIsNone(result)
|
|
|
|
def test_sharding_callback_saveoption(self):
|
|
servers = [server_lib.Server.create_local_server() for _ in range(3)]
|
|
cluster_spec = server_lib.ClusterSpec({
|
|
"worker": [s.target[len("grpc://"):] for s in servers]})
|
|
remote.connect_to_cluster(cluster_spec)
|
|
root = module.Module()
|
|
with ops.device("/job:worker/task:0/cpu:0"):
|
|
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
|
|
self.evaluate(v0.initializer)
|
|
with ops.device("/job:worker/task:1/cpu:0"):
|
|
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
|
|
v2 = resource_variable_ops.ResourceVariable([2.0, 3.0], name="v2")
|
|
self.evaluate(v1.initializer)
|
|
self.evaluate(v2.initializer)
|
|
root.v0 = v0
|
|
root.v1 = v1
|
|
root.v2 = v2
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), "shard_by_task")
|
|
save.save(
|
|
root, save_dir, options=save_options.SaveOptions(
|
|
experimental_sharding_callback=(
|
|
sharding_policies.ShardByTaskPolicy())))
|
|
self.assertLen(gfile.Glob(save_dir + "/variables/variables.data*"), 3)
|
|
loaded_root = load.load(save_dir)
|
|
self.assertEqual(loaded_root.v0.numpy(), root.v0.numpy())
|
|
self.assertEqual(loaded_root.v1.numpy(), root.v1.numpy())
|
|
self.assertEqual(loaded_root.v2.numpy()[0], root.v2.numpy()[0])
|
|
self.assertEqual(loaded_root.v2.numpy()[1], root.v2.numpy()[1])
|
|
|
|
save_dir = os.path.join(self.get_temp_dir(), "max_shard_size")
|
|
save.save(
|
|
root, save_dir, options=save_options.SaveOptions(
|
|
experimental_sharding_callback=(
|
|
sharding_policies.MaxShardSizePolicy(max_shard_size=(4)))))
|
|
self.assertLen(gfile.Glob(save_dir + "/variables/variables.data*"), 5)
|
|
loaded_root = load.load(save_dir)
|
|
self.assertEqual(loaded_root.v0.numpy(), root.v0.numpy())
|
|
self.assertEqual(loaded_root.v1.numpy(), root.v1.numpy())
|
|
self.assertEqual(loaded_root.v2.numpy()[0], root.v2.numpy()[0])
|
|
self.assertEqual(loaded_root.v2.numpy()[1], root.v2.numpy()[1])
|
|
|
|
|
|
class DependencyTest(test.TestCase):
|
|
"""Tests for deserialization dependencies (saving-related only)."""
|
|
|
|
def test_validate_dependencies(self):
|
|
|
|
class Valid(autotrackable.AutoTrackable):
|
|
|
|
def _deserialization_dependencies(self, children):
|
|
return children
|
|
|
|
root = Valid()
|
|
root.f = variables.Variable(1.0)
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir)
|
|
|
|
def test_validate_dependencies_error_untracked(self):
|
|
untracked = variables.Variable(1.0)
|
|
|
|
class Invalid(autotrackable.AutoTrackable):
|
|
|
|
def _deserialization_dependencies(self, children):
|
|
del children # Unused.
|
|
return {"untracked": untracked}
|
|
invalid_deps = Invalid()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError, "Found an untracked dependency"):
|
|
save.save(invalid_deps, save_dir)
|
|
|
|
def test_validate_dependencies_error_cyclic(self):
|
|
|
|
class Invalid(autotrackable.AutoTrackable):
|
|
|
|
def __init__(self):
|
|
self.cycle_ref = None
|
|
|
|
def _deserialization_dependencies(self, children):
|
|
del children # Unused.
|
|
return {"cycle_ref": self.cycle_ref}
|
|
cycle1 = Invalid()
|
|
cycle2 = Invalid()
|
|
cycle1.cycle_ref = cycle2
|
|
cycle2.cycle_ref = cycle1
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(ValueError,
|
|
"dependency cycle in the saved Trackable"):
|
|
save.save(cycle1, save_dir)
|
|
|
|
|
|
class VariablePolicyEnumTest(test.TestCase):
|
|
|
|
def testFromObj(self):
|
|
self.assertEqual(save_options.VariablePolicy.NONE,
|
|
save_options.VariablePolicy.from_obj(None))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES,
|
|
save_options.VariablePolicy.from_obj(
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES,
|
|
save_options.VariablePolicy.from_obj(
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES,
|
|
save_options.VariablePolicy.from_obj("save_variable_devices"))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.SAVE_VARIABLE_DEVICES,
|
|
save_options.VariablePolicy.from_obj("SaVe_VaRiAbLe_DeViCeS"))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES,
|
|
save_options.VariablePolicy.from_obj("expand_distributed_variables"))
|
|
self.assertEqual(
|
|
save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES,
|
|
save_options.VariablePolicy.from_obj("eXpAnD_dIsTrIbUtEd_VaRiAbLeS"))
|
|
for invalid in ["not_a_valid_value", 2.0, []]:
|
|
with self.assertRaisesRegex(ValueError, "invalid VariablePolicy value"):
|
|
save_options.VariablePolicy.from_obj(invalid)
|
|
|
|
def testNamingConvention(self):
|
|
"""Enforces names are uppercase versions of values."""
|
|
for policy in save_options.VariablePolicy:
|
|
if policy == save_options.VariablePolicy.NONE:
|
|
self.assertIsNone(policy.value)
|
|
else:
|
|
self.assertEqual(policy.name, policy.name.upper())
|
|
self.assertEqual(policy.value, policy.value.lower())
|
|
self.assertEqual(policy.name, policy.value.upper())
|
|
|
|
|
|
class SavingOptionsTest(test.TestCase):
|
|
|
|
def testOpNameSpace(self):
|
|
# TODO(kathywu): Add test that saves out SavedModel with a custom op when
|
|
# the ">" character is allowed in op names.
|
|
graph_def = graph_pb2.GraphDef()
|
|
text_format.Parse("node { name: 'A' op: 'Test>CustomOp' }", graph_def)
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Attempted to save ops from non-whitelisted namespaces"):
|
|
save._verify_ops(graph_def, [])
|
|
save._verify_ops(graph_def, ["Test"])
|
|
|
|
# Test with multiple carrots in op name.
|
|
text_format.Parse("node { name: 'A' op: 'Test>>A>CustomOp' }", graph_def)
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Attempted to save ops from non-whitelisted namespaces"):
|
|
save._verify_ops(graph_def, [])
|
|
save._verify_ops(graph_def, ["Test"])
|
|
|
|
def test_save_custom_op_with_no_whitelist_specified(self):
|
|
# Test that we are able to save a model that contains a custom op with a
|
|
# custom namespace when the user has not explicitly specified a namespace
|
|
# whitelist (i.e. that we default to allowing all custom ops when saving
|
|
# and no whitelist is specified, rather than throwing an exception).
|
|
graph_def = graph_pb2.GraphDef()
|
|
text_format.Parse("node { name: 'A' op: 'Test>CustomOp' }", graph_def)
|
|
save._verify_ops(graph_def, namespace_whitelist=None)
|
|
|
|
# If the user passes an empty list for the namespace whitelist rather than
|
|
# nothing, we should then throw an exception if a custom op is used.
|
|
with self.assertRaisesRegex(
|
|
ValueError, "Attempted to save ops from non-whitelisted namespaces"
|
|
):
|
|
save._verify_ops(graph_def, [])
|
|
|
|
def test_strip_debug_nodes(self):
|
|
# Test that we are able to strip debug nodes from a meta_graph correctly.
|
|
test_node_defs = [
|
|
node_def_pb2.NodeDef(
|
|
name="AssertNode",
|
|
op="Assert",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
"^ControlInput:output:0",
|
|
],
|
|
attr={
|
|
"regular_node_attr": attr_value_pb2.AttrValue(i=1),
|
|
"_non_regular_node_attr": attr_value_pb2.AttrValue(i=2),
|
|
}
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="ConstNode",
|
|
op="Const",
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="CheckNumericsNode",
|
|
op="CheckNumerics",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
"NonControlInputTwo:output:0",
|
|
"^ControlInput:output:0",
|
|
],
|
|
attr={
|
|
"T": attr_value_pb2.AttrValue(i=4),
|
|
"NotT": attr_value_pb2.AttrValue(i=5),
|
|
}
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="CheckNumericsNodeTwo",
|
|
op="CheckNumerics",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
"NonControlInputTwo:output:0",
|
|
"^ControlInput:output:0",
|
|
],
|
|
attr={
|
|
"OnlyNotT": attr_value_pb2.AttrValue(i=6),
|
|
},
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="PrintNode",
|
|
op="Print",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
],
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="PrintV2Node",
|
|
op="PrintV2",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
],
|
|
),
|
|
]
|
|
|
|
expected_node_defs = [
|
|
node_def_pb2.NodeDef(
|
|
name="AssertNode",
|
|
op="NoOp",
|
|
input=[
|
|
"^NonControlInput",
|
|
"^ControlInput:output:0",
|
|
],
|
|
attr={
|
|
"_non_regular_node_attr": attr_value_pb2.AttrValue(i=2),
|
|
}
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="ConstNode",
|
|
op="Const",
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="CheckNumericsNode",
|
|
op="Identity",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
"^NonControlInputTwo",
|
|
"^ControlInput:output:0",
|
|
],
|
|
attr={
|
|
"T": attr_value_pb2.AttrValue(i=4),
|
|
}
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="CheckNumericsNodeTwo",
|
|
op="Identity",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
"^NonControlInputTwo",
|
|
"^ControlInput:output:0",
|
|
],
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="PrintNode",
|
|
op="Identity",
|
|
input=[
|
|
"NonControlInput:output:0",
|
|
],
|
|
),
|
|
node_def_pb2.NodeDef(
|
|
name="PrintV2Node",
|
|
op="NoOp",
|
|
input=[
|
|
"^NonControlInput",
|
|
],
|
|
),
|
|
]
|
|
|
|
meta_graph_def = meta_graph_pb2.MetaGraphDef(
|
|
graph_def=graph_pb2.GraphDef(
|
|
node=test_node_defs,
|
|
library=function_pb2.FunctionDefLibrary(
|
|
function=[function_pb2.FunctionDef(node_def=test_node_defs)]
|
|
),
|
|
),
|
|
)
|
|
|
|
expected = meta_graph_pb2.MetaGraphDef(
|
|
graph_def=graph_pb2.GraphDef(
|
|
node=expected_node_defs,
|
|
library=function_pb2.FunctionDefLibrary(
|
|
function=[function_pb2.FunctionDef(node_def=expected_node_defs)]
|
|
),
|
|
),
|
|
)
|
|
|
|
save._strip_debug_nodes(meta_graph_def)
|
|
self.assertEqual(expected, meta_graph_def)
|
|
|
|
def test_save_debug_info_enabled(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: math_ops.mul(2.0, x, name="DEBUG_INFO_OP"),
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)],
|
|
)
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root,
|
|
save_dir,
|
|
root.f,
|
|
options=save_options.SaveOptions(save_debug_info=True),
|
|
)
|
|
debug_info_file_name = os.path.join(
|
|
save_dir, "debug", "saved_model_debug_info.pb"
|
|
)
|
|
self.assertTrue(os.path.exists(debug_info_file_name))
|
|
debug_info = graph_debug_info_pb2.GraphDebugInfo()
|
|
with open(debug_info_file_name, "rb") as f:
|
|
debug_info.ParseFromString(f.read())
|
|
|
|
# Verify that there is a trace for DEBUG_INFO_OP just to ensure that
|
|
# function debug info tracing is nominally functioning.
|
|
found_op = False
|
|
for key in debug_info.name_to_trace_id.keys():
|
|
if key.startswith("DEBUG_INFO_OP@"):
|
|
found_op = True
|
|
break
|
|
self.assertTrue(
|
|
found_op, "Did not find DEBUG_INFO_OP in trace: %s" % debug_info
|
|
)
|
|
|
|
def test_save_debug_info_disabled(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: math_ops.mul(2., x, name="DEBUG_INFO_OP"),
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(
|
|
root,
|
|
save_dir,
|
|
root.f,
|
|
options=save_options.SaveOptions(save_debug_info=False))
|
|
debug_info_file_name = os.path.join(save_dir, "debug",
|
|
"saved_model_debug_info.pb")
|
|
self.assertFalse(os.path.exists(debug_info_file_name))
|
|
|
|
def test_function_aliases(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: 2. * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
options = save_options.SaveOptions(
|
|
function_aliases={
|
|
"my_func": root.f,
|
|
}
|
|
)
|
|
save.save(root, save_dir, root.f, options=options)
|
|
function_cache = root.f._variable_creation_config.function_cache.values()
|
|
function_aliases = loader_impl.parse_saved_model(
|
|
save_dir).meta_graphs[0].meta_info_def.function_aliases
|
|
self.assertLen(function_cache, 1)
|
|
self.assertEqual(function_cache[0].name.decode("utf-8"),
|
|
list(function_aliases.keys())[0])
|
|
|
|
def test_concrete_function_aliases(self):
|
|
root = autotrackable.AutoTrackable()
|
|
f = def_function.function(
|
|
lambda x: 2.0 * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)],
|
|
).get_concrete_function()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
options = save_options.SaveOptions(
|
|
function_aliases={
|
|
"my_func": f,
|
|
}
|
|
)
|
|
save.save(root, save_dir, f, options=options)
|
|
function_aliases = loader_impl.parse_saved_model(
|
|
save_dir).meta_graphs[0].meta_info_def.function_aliases
|
|
self.assertEqual(f.name.decode("utf-8"),
|
|
list(function_aliases.keys())[0])
|
|
|
|
def test_concrete_function_list_aliases(self):
|
|
root = autotrackable.AutoTrackable()
|
|
f = def_function.function(lambda z: {"out": z * z})
|
|
f1 = f.get_concrete_function(tensor_spec.TensorSpec(None, dtypes.float32))
|
|
f2 = f.get_concrete_function(tensor_spec.TensorSpec(None, dtypes.int32))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
options = save_options.SaveOptions(
|
|
function_aliases={
|
|
"my_func": [f1, f2],
|
|
}
|
|
)
|
|
save.save(root, save_dir, f1, options=options)
|
|
function_aliases = (
|
|
loader_impl.parse_saved_model(save_dir)
|
|
.meta_graphs[0]
|
|
.meta_info_def.function_aliases
|
|
)
|
|
self.assertSameElements(
|
|
[f1.name.decode("utf-8"), f2.name.decode("utf-8")],
|
|
list(function_aliases.keys()),
|
|
)
|
|
|
|
def test_function_aliases_incorrect_type(self):
|
|
root = autotrackable.AutoTrackable()
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
f = lambda x: 2.0 * x
|
|
root.f = def_function.function(
|
|
f, input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)]
|
|
)
|
|
options = save_options.SaveOptions(
|
|
function_aliases={
|
|
"my_func": f,
|
|
}
|
|
)
|
|
with self.assertRaisesRegex(TypeError, "Unsupported type"):
|
|
save.save(root, save_dir, root.f, options=options)
|
|
|
|
def test_accepts_io_device(self):
|
|
options = save_options.SaveOptions()
|
|
self.assertIsNone(options.experimental_io_device)
|
|
options = save_options.SaveOptions(experimental_io_device="/job:localhost")
|
|
self.assertEqual("/job:localhost", options.experimental_io_device)
|
|
|
|
def test_accepts_variable_policy(self):
|
|
options = save_options.SaveOptions()
|
|
self.assertEqual(save_options.VariablePolicy.NONE,
|
|
options.experimental_variable_policy)
|
|
# VariablePolicy instances.
|
|
options = save_options.SaveOptions(experimental_variable_policy=save_options
|
|
.VariablePolicy.SAVE_VARIABLE_DEVICES)
|
|
self.assertEqual(save_options.VariablePolicy.SAVE_VARIABLE_DEVICES,
|
|
options.experimental_variable_policy)
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy=save_options.VariablePolicy
|
|
.EXPAND_DISTRIBUTED_VARIABLES)
|
|
self.assertEqual(save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES,
|
|
options.experimental_variable_policy)
|
|
# String conversions.
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy="save_variable_devices")
|
|
self.assertEqual(save_options.VariablePolicy.SAVE_VARIABLE_DEVICES,
|
|
options.experimental_variable_policy)
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy="expand_distributed_variables")
|
|
self.assertEqual(save_options.VariablePolicy.EXPAND_DISTRIBUTED_VARIABLES,
|
|
options.experimental_variable_policy)
|
|
with self.assertRaisesRegex(ValueError, "invalid VariablePolicy value"):
|
|
options = save_options.SaveOptions(
|
|
experimental_variable_policy="not_a_valid_value")
|
|
|
|
|
|
class AssetTests(test.TestCase):
|
|
|
|
def setUp(self):
|
|
super(AssetTests, self).setUp()
|
|
self._vocab_path = os.path.join(self.get_temp_dir(), "vocab.txt")
|
|
with open(self._vocab_path, "w") as f:
|
|
f.write("alpha\nbeta\ngamma\n")
|
|
|
|
def test_asset_path_returned(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.path = asset.Asset(self._vocab_path)
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
root.get_asset = def_function.function(lambda: root.path.asset_path)
|
|
save.save(root, save_dir, signatures=root.get_asset.get_concrete_function())
|
|
second_dir = os.path.join(self.get_temp_dir(), "second_dir")
|
|
file_io.rename(save_dir, second_dir)
|
|
imported_path = _import_and_infer(second_dir, {})["output_0"]
|
|
self.assertIn(
|
|
compat.as_str_any(second_dir), compat.as_str_any(imported_path))
|
|
|
|
def test_table(self):
|
|
initializer = lookup_ops.TextFileInitializer(
|
|
self._vocab_path,
|
|
key_dtype=dtypes.string,
|
|
key_index=lookup_ops.TextFileIndex.WHOLE_LINE,
|
|
value_dtype=dtypes.int64,
|
|
value_index=lookup_ops.TextFileIndex.LINE_NUMBER)
|
|
root = checkpoint.Checkpoint(
|
|
table=lookup_ops.HashTable(initializer, default_value=-1))
|
|
root.table_user = def_function.function(
|
|
root.table.lookup,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.string)])
|
|
self.assertEqual(
|
|
2, self.evaluate(root.table_user(constant_op.constant("gamma"))))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
save.save(root, save_dir)
|
|
file_io.delete_file(self._vocab_path)
|
|
self.assertAllClose({"output_0": [2, 0]},
|
|
_import_and_infer(save_dir,
|
|
{"keys": ["gamma", "alpha"]}))
|
|
second_dir = os.path.join(self.get_temp_dir(), "second_dir")
|
|
# Asset paths should track the location the SavedModel is loaded from.
|
|
file_io.rename(save_dir, second_dir)
|
|
self.assertAllClose({"output_0": [2, 1]},
|
|
_import_and_infer(second_dir,
|
|
{"keys": ["gamma", "beta"]}))
|
|
|
|
def test_untracked_table_useful_message(self):
|
|
root = module.Module()
|
|
initializer = lookup_ops.TextFileInitializer(
|
|
self._vocab_path,
|
|
key_dtype=dtypes.string,
|
|
key_index=lookup_ops.TextFileIndex.WHOLE_LINE,
|
|
value_dtype=dtypes.int64,
|
|
value_index=lookup_ops.TextFileIndex.LINE_NUMBER)
|
|
table = lookup_ops.HashTable(initializer, default_value=-1)
|
|
root.table_user = def_function.function(
|
|
table.lookup,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.string)])
|
|
root.table_user(constant_op.constant("gamma"))
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
with self.assertRaisesRegex(AssertionError, "HashTable"):
|
|
save.save(root, save_dir)
|
|
|
|
def test_unused_asset(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.f = def_function.function(
|
|
lambda x: 2. * x,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
root.asset = asset.Asset(self._vocab_path)
|
|
|
|
export_dir = os.path.join(self.get_temp_dir(), "save_dir")
|
|
save.save(root, export_dir)
|
|
self.assertAllClose({"output_0": [0.2]},
|
|
_import_and_infer(export_dir, {"x": [0.1]}))
|
|
|
|
def test_sensible_function_building_exception(self):
|
|
root = checkpoint.Checkpoint(v=variables.Variable(2.))
|
|
root.f = def_function.function(
|
|
lambda x: 2. * root.v,
|
|
input_signature=[tensor_spec.TensorSpec(None, dtypes.float32)])
|
|
export_dir = os.path.join(self.get_temp_dir(), "save_dir")
|
|
|
|
@def_function.function
|
|
def _calls_save():
|
|
save.save(root, export_dir)
|
|
|
|
with self.assertRaisesRegex(AssertionError, "tf.function"):
|
|
_calls_save()
|
|
|
|
def test_rewrite_asset_to_same_destination(self):
|
|
save_dir = os.path.join(self.get_temp_dir(), "saved_model")
|
|
asset_path = os.path.join(self.get_temp_dir(), "asset")
|
|
|
|
def save_and_load(label):
|
|
with open(asset_path, "w") as f:
|
|
f.write(label)
|
|
|
|
model = autotrackable.AutoTrackable()
|
|
model.asset = asset.Asset(asset_path)
|
|
model.fn = def_function.function(lambda: io_ops.read_file(model.asset))
|
|
self.assertEqual(label, model.fn().numpy().decode("utf-8"))
|
|
|
|
save.save(model, save_dir)
|
|
imported = load.load(save_dir)
|
|
self.assertEqual(label, imported.fn().numpy().decode("utf-8"))
|
|
|
|
save_and_load("first")
|
|
save_and_load("second")
|
|
|
|
|
|
class ExportMetaGraphTests(test.TestCase):
|
|
|
|
def test_export_meta_graph(self):
|
|
root = autotrackable.AutoTrackable()
|
|
root.variable = resource_variable_ops.UninitializedVariable(
|
|
name="some_variable", dtype=dtypes.float32)
|
|
|
|
@def_function.function(input_signature=[tensor_spec.TensorSpec(None)])
|
|
def multiply_var(x):
|
|
return root.variable * x
|
|
|
|
@def_function.function(input_signature=[tensor_spec.TensorSpec([])])
|
|
def update(y):
|
|
root.variable.assign_add(y)
|
|
# TODO(b/150393409): All functions exported as signatures must have at
|
|
# least one output.
|
|
return 0
|
|
|
|
@def_function.function(input_signature=[])
|
|
def initialize():
|
|
root.variable.assign(1.0)
|
|
# TODO(b/150393409): All functions exported as signatures must have at
|
|
# least one output.
|
|
return 0
|
|
|
|
save_path = os.path.join(self.get_temp_dir(), "meta_graph.pb")
|
|
save.export_meta_graph(
|
|
root,
|
|
save_path,
|
|
signatures={
|
|
"multiply_var": multiply_var,
|
|
"initialize": initialize,
|
|
"update": update
|
|
})
|
|
|
|
with ops.Graph().as_default(), session_lib.Session() as session:
|
|
saver.import_meta_graph(save_path)
|
|
meta_graph_def = meta_graph.read_meta_graph_file(save_path)
|
|
|
|
# Initialize variable to 1
|
|
_run_signature(session, meta_graph_def, {}, "initialize")
|
|
out = _run_signature(session, meta_graph_def, {"x": 3}, "multiply_var")
|
|
self.assertAllEqual(out, {"output_0": 3})
|
|
|
|
# Adds 2 to the variable. Variable is now 3
|
|
_run_signature(session, meta_graph_def, {"y": 2}, "update")
|
|
out = _run_signature(session, meta_graph_def, {"x": 4}, "multiply_var")
|
|
self.assertAllEqual(out, {"output_0": 12})
|
|
|
|
|
|
class FingerprintingTests(test.TestCase):
|
|
|
|
def test_toggle_flag(self):
|
|
self.assertTrue(flags.config().saved_model_fingerprinting.value())
|
|
flags.config().saved_model_fingerprinting.reset(False)
|
|
self.assertFalse(flags.config().saved_model_fingerprinting.value())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test.main()
|