wehub-resource-sync
153 lines
6.0 KiB
Python
153 lines
6.0 KiB
Python
# Copyright 2017 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.
|
|
# ==============================================================================
|
|
"""Dependency tracking for trackable objects."""
|
|
|
|
import warnings
|
|
|
|
from absl import logging
|
|
|
|
from tensorflow.python.eager import def_function
|
|
from tensorflow.python.eager import function as defun
|
|
from tensorflow.python.trackable import base
|
|
from tensorflow.python.trackable import data_structures
|
|
from tensorflow.python.types import core as core_types
|
|
from tensorflow.python.util.tf_export import tf_export
|
|
|
|
|
|
@tf_export("__internal__.tracking.AutoTrackable", v1=[])
|
|
class AutoTrackable(base.Trackable):
|
|
"""Manages dependencies on other objects.
|
|
|
|
`Trackable` objects may have dependencies: other `Trackable` objects
|
|
which should be saved if the object declaring the dependency is saved. A
|
|
correctly saveable program has a dependency graph such that if changing a
|
|
global variable affects an object (e.g. changes the behavior of any of its
|
|
methods) then there is a chain of dependencies from the influenced object to
|
|
the variable.
|
|
|
|
Dependency edges have names, and are created implicitly when a
|
|
`Trackable` object is assigned to an attribute of another
|
|
`Trackable` object. For example:
|
|
|
|
```
|
|
obj = Trackable()
|
|
obj.v = ResourceVariable(0.)
|
|
```
|
|
|
|
The `Trackable` object `obj` now has a dependency named "v" on a
|
|
variable.
|
|
|
|
`Trackable` objects may specify `Tensor`s to be saved and restored
|
|
directly (e.g. a `Variable` indicating how to save itself) rather than through
|
|
dependencies on other objects. See
|
|
`Trackable._gather_saveables_for_checkpoint` for details.
|
|
"""
|
|
|
|
def __setattr__(self, name, value):
|
|
"""Support self.foo = trackable syntax."""
|
|
try:
|
|
if getattr(self, name) is value:
|
|
# Short circuit for `self.$x = self.$x`.
|
|
return
|
|
except AttributeError:
|
|
pass
|
|
|
|
if getattr(self, "_self_setattr_tracking", True):
|
|
value = data_structures.sticky_attribute_assignment(
|
|
trackable=self, value=value, name=name)
|
|
super(AutoTrackable, self).__setattr__(name, value)
|
|
|
|
def __delattr__(self, name):
|
|
self._delete_tracking(name)
|
|
super(AutoTrackable, self).__delattr__(name)
|
|
|
|
def _no_dependency(self, value):
|
|
"""Override to allow TrackableBase to disable dependency tracking."""
|
|
return data_structures.NoDependency(value)
|
|
|
|
def _trackable_children(self, save_type=base.SaveType.CHECKPOINT, **kwargs):
|
|
"""Returns all children of a trackable, including functions."""
|
|
if save_type != base.SaveType.SAVEDMODEL:
|
|
return super(AutoTrackable, self)._trackable_children(
|
|
save_type, **kwargs)
|
|
|
|
functions = {}
|
|
try:
|
|
# We get the attributes, suppressing warnings and exceptions.
|
|
logging_verbosity = logging.get_verbosity()
|
|
logging.set_verbosity(logging.FATAL)
|
|
for attribute_name in dir(self):
|
|
try:
|
|
with warnings.catch_warnings():
|
|
warnings.simplefilter("ignore")
|
|
attribute_value = getattr(self, attribute_name, None)
|
|
except Exception: # pylint: disable=broad-except
|
|
# NOTE: If we make the exception catching here less broad, we might
|
|
# need to revisit `finally` block below.
|
|
# We really don't want to throw an exception just because some
|
|
# object's attribute accessor is broken.
|
|
attribute_value = None
|
|
if isinstance(attribute_value, (def_function.Function,
|
|
defun.ConcreteFunction)):
|
|
functions[attribute_name] = attribute_value
|
|
finally:
|
|
logging.set_verbosity(logging_verbosity)
|
|
|
|
# Trace concrete functions to force side-effects:
|
|
# 1. populate the cache for functions that have an input_signature
|
|
# and have not been called
|
|
# 2. force side effects of creation of concrete functions, e.g. create
|
|
# variables on first run.
|
|
for fn in functions.values():
|
|
if isinstance(fn, def_function.Function):
|
|
fn._list_all_concrete_functions_for_serialization() # pylint: disable=protected-access
|
|
|
|
# Additional dependencies may have been generated during function tracing
|
|
# (e.g. captured variables). Make sure we return those too.
|
|
children = {}
|
|
for name, child in self._checkpoint_dependencies:
|
|
if isinstance(child, (core_types.PolymorphicFunction,
|
|
core_types.ConcreteFunction)):
|
|
# Skip "tracked" functions for now since there may be objects that
|
|
# automatically track functions that should not be saved.
|
|
# TODO(kathywu): remove once `_list_functions_for_serialization` has
|
|
# been fully deprecated.
|
|
continue
|
|
|
|
if name in functions and child is not functions[name]:
|
|
raise ValueError(
|
|
"Can't save object because it has multiple children with the same "
|
|
f"name. Object: {self}, attribute name: {name}, child 1: "
|
|
f"{child}, child 2: {functions[name]}")
|
|
|
|
children[name] = child
|
|
|
|
children.update(functions)
|
|
return children
|
|
|
|
def _delete_tracking(self, name):
|
|
"""Removes the tracking of name."""
|
|
self._maybe_initialize_trackable()
|
|
if name in self._unconditional_dependency_names:
|
|
del self._unconditional_dependency_names[name]
|
|
for index, (dep_name, _) in enumerate(
|
|
self._unconditional_checkpoint_dependencies):
|
|
if dep_name == name:
|
|
del self._unconditional_checkpoint_dependencies[index]
|
|
break
|
|
|
|
def _add_trackable_child(self, name, value):
|
|
self.__setattr__(name, value)
|