290 lines
9.7 KiB
Python
290 lines
9.7 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.
|
||
|
# ==============================================================================
|
||
|
"""Utility functions shared between SavedModel saving/loading
|
||
|
implementations."""
|
||
|
|
||
|
import copy
|
||
|
import itertools
|
||
|
import threading
|
||
|
import types
|
||
|
|
||
|
import tensorflow.compat.v2 as tf
|
||
|
|
||
|
from keras import backend
|
||
|
from keras.engine import base_layer_utils
|
||
|
from keras.utils import control_flow_util
|
||
|
from keras.utils import tf_contextlib
|
||
|
from keras.utils.generic_utils import LazyLoader
|
||
|
from keras.utils.layer_utils import CallFunctionSpec
|
||
|
|
||
|
training_lib = LazyLoader("training_lib", globals(), "keras.engine.training")
|
||
|
|
||
|
|
||
|
def use_wrapped_call(
|
||
|
layer, call_fn, call_spec, default_training_value=None, return_method=False
|
||
|
):
|
||
|
"""Creates fn that adds losses returned by call_fn & returns the outputs.
|
||
|
|
||
|
Args:
|
||
|
layer: A Keras layer object
|
||
|
call_fn: tf.function that takes layer inputs (and possibly a training
|
||
|
arg), and returns a tuple of (outputs, list of losses).
|
||
|
call_spec: The `CallFunctionSpec` for the layer's call function.
|
||
|
default_training_value: Default value of the training kwarg. If `None`,
|
||
|
the default is `tf.keras.backend.learning_phase()`.
|
||
|
return_method: Whether to return a method bound to the layer.
|
||
|
|
||
|
Returns:
|
||
|
function that calls call_fn and returns the outputs. Losses returned by
|
||
|
call_fn are added to the layer losses.
|
||
|
"""
|
||
|
expects_training_arg = layer_uses_training_bool(layer)
|
||
|
|
||
|
fn, arg_spec = maybe_add_training_arg(
|
||
|
call_spec, call_fn, expects_training_arg, default_training_value
|
||
|
)
|
||
|
|
||
|
def return_outputs_and_add_losses(*args, **kwargs):
|
||
|
"""Returns the outputs from the layer call function, and adds the
|
||
|
losses."""
|
||
|
if return_method:
|
||
|
args = args[1:]
|
||
|
|
||
|
outputs, losses = fn(*args, **kwargs)
|
||
|
layer.add_loss(losses)
|
||
|
|
||
|
# TODO(kathywu): This is a temporary hack. When a network of layers is
|
||
|
# revived from SavedModel, only the top-level layer will have losses.
|
||
|
# This causes issues in eager mode because the child layers may have
|
||
|
# graph losses (thus model.losses returns a mix of Eager and graph
|
||
|
# tensors). To fix this, whenever eager losses are added to one layer,
|
||
|
# add eager losses to all child layers. This causes `.losses` to only
|
||
|
# return eager losses.
|
||
|
|
||
|
if tf.executing_eagerly():
|
||
|
for i in layer._flatten_layers():
|
||
|
if i is not layer:
|
||
|
i._eager_losses = [
|
||
|
base_layer_utils.REVIVED_LOSS_PLACEHOLDER
|
||
|
]
|
||
|
|
||
|
return outputs
|
||
|
|
||
|
decorated = tf.__internal__.decorator.make_decorator(
|
||
|
target=call_fn,
|
||
|
decorator_func=return_outputs_and_add_losses,
|
||
|
decorator_argspec=arg_spec,
|
||
|
)
|
||
|
|
||
|
if return_method:
|
||
|
return types.MethodType(decorated, layer)
|
||
|
else:
|
||
|
return decorated
|
||
|
|
||
|
|
||
|
def layer_uses_training_bool(layer):
|
||
|
"""Returns whether this layer or any of its children uses the training
|
||
|
arg."""
|
||
|
if layer._expects_training_arg:
|
||
|
return True
|
||
|
visited = {layer}
|
||
|
to_visit = list_all_layers(layer)
|
||
|
while to_visit:
|
||
|
layer = to_visit.pop()
|
||
|
if layer in visited:
|
||
|
continue
|
||
|
if getattr(layer, "_expects_training_arg", True):
|
||
|
return True
|
||
|
visited.add(layer)
|
||
|
to_visit.extend(list_all_layers(layer))
|
||
|
return False
|
||
|
|
||
|
|
||
|
def list_all_layers(obj):
|
||
|
if isinstance(obj, training_lib.Model):
|
||
|
# Handle special case of Sequential, which doesn't return
|
||
|
# the `Input` layer.
|
||
|
return obj.layers
|
||
|
else:
|
||
|
return list(obj._flatten_layers(include_self=False, recursive=False))
|
||
|
|
||
|
|
||
|
def list_all_layers_and_sublayers(obj):
|
||
|
s = set([obj])
|
||
|
s.update(
|
||
|
itertools.chain.from_iterable(
|
||
|
list_all_layers_and_sublayers(layer)
|
||
|
for layer in list_all_layers(obj)
|
||
|
)
|
||
|
)
|
||
|
return s
|
||
|
|
||
|
|
||
|
def maybe_add_training_arg(
|
||
|
call_spec, wrapped_call, expects_training_arg, default_training_value
|
||
|
):
|
||
|
"""Decorate call and optionally adds training argument.
|
||
|
|
||
|
If a layer expects a training argument, this function ensures that
|
||
|
'training' is present in the layer args or kwonly args, with the default
|
||
|
training value.
|
||
|
|
||
|
Args:
|
||
|
call_spec: CallFunctionSpec of the layer.
|
||
|
wrapped_call: Wrapped call function.
|
||
|
expects_training_arg: Whether to include 'training' argument.
|
||
|
default_training_value: Default value of the training kwarg to include in
|
||
|
the arg spec. If `None`, the default is
|
||
|
`tf.keras.backend.learning_phase()`.
|
||
|
|
||
|
Returns:
|
||
|
Tuple of (
|
||
|
function that calls `wrapped_call` and sets the training arg,
|
||
|
Argspec of returned function or `None` if the argspec is unchanged)
|
||
|
"""
|
||
|
if not expects_training_arg:
|
||
|
return wrapped_call, None
|
||
|
|
||
|
arg_spec = set_training_arg_spec(
|
||
|
call_spec.full_argspec, default_training_value
|
||
|
)
|
||
|
call_spec = CallFunctionSpec(arg_spec)
|
||
|
|
||
|
def wrap_with_training_arg(*args, **kwargs):
|
||
|
"""Wrap the `wrapped_call` function, and set training argument."""
|
||
|
try:
|
||
|
training = call_spec.get_arg_value(
|
||
|
"training", args, kwargs, inputs_in_args=True
|
||
|
)
|
||
|
except KeyError:
|
||
|
training = None
|
||
|
|
||
|
if training is None:
|
||
|
training = (
|
||
|
default_training_value
|
||
|
or base_layer_utils.call_context().training
|
||
|
or backend.learning_phase()
|
||
|
)
|
||
|
|
||
|
args = list(args)
|
||
|
kwargs = kwargs.copy()
|
||
|
|
||
|
def replace_training_and_call(training):
|
||
|
new_args, new_kwargs = call_spec.set_arg_value(
|
||
|
"training", training, args, kwargs, inputs_in_args=True
|
||
|
)
|
||
|
return wrapped_call(*new_args, **new_kwargs)
|
||
|
|
||
|
return control_flow_util.smart_cond(
|
||
|
training,
|
||
|
lambda: replace_training_and_call(True),
|
||
|
lambda: replace_training_and_call(False),
|
||
|
)
|
||
|
|
||
|
return wrap_with_training_arg, arg_spec
|
||
|
|
||
|
|
||
|
def set_training_arg_spec(arg_spec, default_training_value):
|
||
|
"""Set `training=DEFAULT` argument in an ArgSpec."""
|
||
|
if "training" in arg_spec.args:
|
||
|
# If `training` is already in the args list, try to set the default
|
||
|
# value.
|
||
|
index = arg_spec.args.index("training")
|
||
|
training_default_index = len(arg_spec.args) - index
|
||
|
defaults = (
|
||
|
list(arg_spec.defaults) if arg_spec.defaults is not None else []
|
||
|
)
|
||
|
if (
|
||
|
arg_spec.defaults
|
||
|
and len(arg_spec.defaults) >= training_default_index
|
||
|
and defaults[-training_default_index] is None
|
||
|
):
|
||
|
defaults[-training_default_index] = default_training_value
|
||
|
return arg_spec._replace(defaults=defaults)
|
||
|
elif "training" not in arg_spec.kwonlyargs:
|
||
|
kwonlyargs = arg_spec.kwonlyargs + ["training"]
|
||
|
kwonlydefaults = copy.copy(arg_spec.kwonlydefaults) or {}
|
||
|
kwonlydefaults["training"] = default_training_value
|
||
|
return arg_spec._replace(
|
||
|
kwonlyargs=kwonlyargs, kwonlydefaults=kwonlydefaults
|
||
|
)
|
||
|
|
||
|
return arg_spec
|
||
|
|
||
|
|
||
|
class SaveOptionsContext(threading.local):
|
||
|
def __init__(self):
|
||
|
super().__init__()
|
||
|
self.save_traces = True
|
||
|
self.in_tf_saved_model_scope = False
|
||
|
|
||
|
|
||
|
_save_options_context = SaveOptionsContext()
|
||
|
|
||
|
|
||
|
@tf_contextlib.contextmanager
|
||
|
def keras_option_scope(save_traces, in_tf_saved_model_scope=True):
|
||
|
save_traces_previous_value = _save_options_context.save_traces
|
||
|
in_scope_previous_value = _save_options_context.in_tf_saved_model_scope
|
||
|
try:
|
||
|
_save_options_context.save_traces = save_traces
|
||
|
_save_options_context.in_tf_saved_model_scope = in_tf_saved_model_scope
|
||
|
yield
|
||
|
finally:
|
||
|
_save_options_context.save_traces = save_traces_previous_value
|
||
|
_save_options_context.in_tf_saved_model_scope = in_scope_previous_value
|
||
|
|
||
|
|
||
|
def should_save_traces():
|
||
|
"""Whether to trace layer functions-can be disabled in the save_traces
|
||
|
arg."""
|
||
|
return _save_options_context.save_traces
|
||
|
|
||
|
|
||
|
def in_tf_saved_model_scope():
|
||
|
return _save_options_context.in_tf_saved_model_scope
|
||
|
|
||
|
|
||
|
@tf_contextlib.contextmanager
|
||
|
def no_automatic_dependency_tracking_scope(obj):
|
||
|
"""Context that disables automatic dependency tracking when assigning attrs.
|
||
|
|
||
|
Objects that inherit from Autotrackable automatically creates dependencies
|
||
|
to trackable objects through attribute assignments, and wraps data
|
||
|
structures (lists or dicts) with trackable classes. This scope may be used
|
||
|
to temporarily disable this behavior. This works similar to the decorator
|
||
|
`no_automatic_dependency_tracking`.
|
||
|
|
||
|
Example usage:
|
||
|
```
|
||
|
model = tf.keras.Model()
|
||
|
model.arr1 = [] # Creates a ListWrapper object
|
||
|
with no_automatic_dependency_tracking_scope(model):
|
||
|
model.arr2 = [] # Creates a regular, untracked python list
|
||
|
```
|
||
|
|
||
|
Args:
|
||
|
obj: A trackable object.
|
||
|
|
||
|
Yields:
|
||
|
a scope in which the object doesn't track dependencies.
|
||
|
"""
|
||
|
previous_value = getattr(obj, "_setattr_tracking", True)
|
||
|
obj._setattr_tracking = False
|
||
|
try:
|
||
|
yield
|
||
|
finally:
|
||
|
obj._setattr_tracking = previous_value
|