0e1f3de50a
`KerasTensor`s are an alternative representation for Keras `Inputs`
and for intermediate outputs of layers during Functional API construction of
models. They are a lightweight data structure comprised of only the
`tf.TypeSpec` of the Tensor that will be consumed/produced in the
corresponding position of the model.
They implement just small subset of `tf.Tensor`'s attributes and
methods, and also overload
the same operators as `tf.Tensor` and automatically turn them into
Keras layers in the model.
`KerasTensor`s are still internal-only and are a work in progress, but they
have several advantages over using a graph `tf.Tensor` to represent
symbolic values in functional models.
- Unlike symbolic tensors, they do not need to refer to a graph. This means
Keras does not need to maintain a never-deleted global background graph
containing all layers ever called during functional model construction when
constructing Functional Models with KerasTensors. These memory savings
can be significant.
- Triggering Keras functional model construction is simpler
when it just has to check whether something is a KerasTensor, rather
than trying to infer if a tensor was meant to be a symbolic keras
representation or just a value produced during function tracing. This means we can add support for cases where values in nest.flatten(*args, **kwargs) are a completely arbitrary mix of KerasTensors and objects that are not KerasTensors, as long as any value is a KerasTensor.
- Autolambda layers (converting tf ops on symbolic Keras tensors to lambda
Keras layers in the model) use TF's internal dispatching mechanism, instead
of trying to manually walk a graph and extract nodes from it.
The dispatching mechanism is simpler, works more reliably, and is less
likely to run into issues with composite tensors or strange tf ops/nodes.
(And when it fails, it's by design: because dispatch is explicitly not
supported on the op & it's more obvious that dispatch doesn't support the
setting).
- Because they support arbitrary typespecs, models/layers that use
KerasTensors are generally more friendly to composite tensors of different
types than using symbolic graph tensors (which must have a TensorSpec and
can't have arbitrary typespecs)
To experiment with using KerasTensors instead of symbolic graph `tf.Tensors`,
import keras_tensor directly and call `keras_tensor.enable_keras_tensors()`
PiperOrigin-RevId: 315009281
Change-Id: I6765f3a44da43f965ec261b6b193df26598cffae
202 lines
5.4 KiB
Python
202 lines
5.4 KiB
Python
# Description:
|
|
# Contains the Keras save model API (internal TensorFlow version).
|
|
|
|
load("//tensorflow:tensorflow.bzl", "tf_py_test")
|
|
|
|
package(
|
|
# TODO(scottzhu): Remove non-keras deps from TF.
|
|
default_visibility = [
|
|
"//tensorflow/python/distribute:__pkg__",
|
|
"//tensorflow/python/keras:__subpackages__",
|
|
],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
exports_files(["LICENSE"])
|
|
|
|
py_library(
|
|
name = "saving",
|
|
srcs = [
|
|
"__init__.py",
|
|
"hdf5_format.py",
|
|
"model_config.py",
|
|
"save.py",
|
|
"saved_model/base_serialization.py",
|
|
"saved_model/constants.py",
|
|
"saved_model/json_utils.py",
|
|
"saved_model/layer_serialization.py",
|
|
"saved_model/load.py",
|
|
"saved_model/metric_serialization.py",
|
|
"saved_model/model_serialization.py",
|
|
"saved_model/network_serialization.py",
|
|
"saved_model/save.py",
|
|
"saved_model/save_impl.py",
|
|
"saved_model/serialized_attributes.py",
|
|
"saved_model/utils.py",
|
|
"saved_model_experimental.py",
|
|
"saving_utils.py",
|
|
],
|
|
srcs_version = "PY2AND3",
|
|
deps = [
|
|
"//tensorflow/python:lib",
|
|
"//tensorflow/python:math_ops",
|
|
"//tensorflow/python:saver",
|
|
"//tensorflow/python:tensor_spec",
|
|
"//tensorflow/python/eager:def_function",
|
|
"//tensorflow/python/keras:backend",
|
|
"//tensorflow/python/keras:optimizers",
|
|
"//tensorflow/python/keras:regularizers",
|
|
"//tensorflow/python/keras/engine:input_spec",
|
|
"//tensorflow/python/keras/utils:engine_utils",
|
|
"//tensorflow/python/keras/utils:mode_keys",
|
|
"//tensorflow/python/saved_model",
|
|
"//tensorflow/python/saved_model/model_utils",
|
|
"//tensorflow/python/training/tracking",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "metrics_serialization_test",
|
|
size = "medium",
|
|
srcs = ["metrics_serialization_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 8,
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "losses_serialization_test",
|
|
size = "medium",
|
|
srcs = ["losses_serialization_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 4,
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "hdf5_format_test",
|
|
size = "medium",
|
|
srcs = ["hdf5_format_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 4,
|
|
tags = [
|
|
"no_oss_py35", # b/147011479
|
|
"no_windows",
|
|
],
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//tensorflow/python/keras:combinations",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "save_test",
|
|
size = "medium",
|
|
srcs = ["save_test.py"],
|
|
python_version = "PY3",
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/feature_column:feature_column_v2",
|
|
"//tensorflow/python/keras",
|
|
"//tensorflow/python/keras:combinations",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "saved_model_experimental_test",
|
|
size = "medium",
|
|
srcs = ["saved_model_experimental_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 4,
|
|
tags = [
|
|
"no_oss", # TODO(b/119349471): Re-enable
|
|
"no_windows",
|
|
],
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "saved_model_test",
|
|
size = "medium",
|
|
srcs = ["saved_model/saved_model_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 4,
|
|
tags = [
|
|
"no_rocm",
|
|
"no_windows",
|
|
],
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/compat:v2_compat",
|
|
"//tensorflow/python/distribute:mirrored_strategy",
|
|
"//tensorflow/python/keras",
|
|
"//tensorflow/python/keras:combinations",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "saving_utils_test",
|
|
size = "medium",
|
|
srcs = ["saving_utils_test.py"],
|
|
python_version = "PY3",
|
|
tags = ["notsan"],
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//tensorflow/python/keras:combinations",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "revive_test",
|
|
size = "medium",
|
|
srcs = ["saved_model/revive_test.py"],
|
|
python_version = "PY3",
|
|
shard_count = 4,
|
|
tags = [
|
|
"no_windows", # b/158005583
|
|
],
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//tensorflow/python/keras",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|
|
|
|
tf_py_test(
|
|
name = "json_utils_test",
|
|
size = "small",
|
|
srcs = ["saved_model/json_utils_test.py"],
|
|
python_version = "PY3",
|
|
deps = [
|
|
"//tensorflow/python:client_testlib",
|
|
"//third_party/py/numpy",
|
|
"@absl_py//absl/testing:parameterized",
|
|
],
|
|
)
|