wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/python/checkpoint/BUILD
+606
View File
@@ -0,0 +1,606 @@
# Description:
# Utilities for reading and writing object-based checkpoints.
load("@xla//third_party/rules_python/python:py_binary.bzl", "py_binary")
load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
load("//tensorflow:tensorflow.default.bzl", "cuda_py_strict_test", "tf_py_strict_test")
load(
"//tensorflow/tools/test:performance.bzl",
"tf_py_benchmark_test",
"tf_py_logged_benchmark",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
py_library(
name = "checkpoint_lib",
strict_deps = True,
visibility = [
"//tensorflow:internal",
"//third_party/py/tf_slim:__subpackages__",
],
deps = [
":checkpoint",
":checkpoint_management",
":checkpoint_options",
":functional_saver",
":graph_view",
":saveable_compat",
":util",
],
)
py_library(
name = "checkpoint_adapter",
srcs = ["checkpoint_adapter.py"],
strict_deps = True,
deps = [
"//tensorflow/python/framework:tensor",
"//tensorflow/python/trackable:base",
"@absl_py//absl/logging",
],
)
py_library(
name = "async_checkpoint_helper",
srcs = ["async_checkpoint_helper.py"],
strict_deps = True,
deps = [
":checkpoint_context",
":trackable_view",
"//tensorflow/python/distribute:device_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:executor",
"//tensorflow/python/framework:ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/saved_model:pywrap_saved_model",
"//tensorflow/python/trackable:base",
"//tensorflow/python/util:object_identity",
"@absl_py//absl/logging",
],
)
py_library(
name = "checkpoint",
srcs = ["checkpoint.py"],
strict_deps = True,
deps = [
":async_checkpoint_helper",
":checkpoint_context",
":checkpoint_management",
":checkpoint_options",
":functional_saver",
":graph_view",
":restore",
":save_util",
":save_util_v1",
":util",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:monitoring",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:io_ops_gen",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model:path_helpers",
"//tensorflow/python/saved_model:pywrap_saved_model",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:data_structures",
"//tensorflow/python/training:py_checkpoint_reader",
"//tensorflow/python/training:saver",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:object_identity",
"//tensorflow/python/util:tf_contextlib",
"//tensorflow/python/util:tf_export",
"//tensorflow/python/util:tf_inspect",
],
)
tf_py_strict_test(
name = "checkpoint_test",
srcs = ["checkpoint_test.py"],
tags = [
"no_pip", # TODO(b/250108043)
"no_windows", # TODO(b/201457117)
],
deps = [
":async_checkpoint_helper",
":checkpoint",
":checkpoint_management",
":checkpoint_options",
":graph_view",
":save_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:stack",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:template",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/ops:variable_v1",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training:adam",
"//tensorflow/python/training:checkpoint_utils",
"//tensorflow/python/training:saver",
"@absl_py//absl/testing:parameterized",
],
)
py_library(
name = "checkpoint_context",
srcs = [
"checkpoint_context.py",
],
strict_deps = True,
)
tf_py_strict_test(
name = "checkpoint_with_v1_optimizers_test",
srcs = ["checkpoint_with_v1_optimizers_test.py"],
deps = [
":checkpoint",
":checkpoint_options",
"//tensorflow/python/client:session",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/ops:init_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:state_ops",
"//tensorflow/python/ops:template",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/training:adam",
"//tensorflow/python/training:checkpoint_utils",
],
)
tf_py_strict_test(
name = "checkpoint_metrics_test",
srcs = ["checkpoint_metrics_test.py"],
deps = [
":checkpoint",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/saved_model:pywrap_saved_model",
],
)
py_library(
name = "checkpoint_view",
srcs = ["checkpoint_view.py"],
strict_deps = True,
tags = ["no_pip"],
deps = [
":trackable_view",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:errors",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training:py_checkpoint_reader",
"//tensorflow/python/util:object_identity",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "checkpoint_view_test",
srcs = ["checkpoint_view_test.py"],
tags = ["no_pip"],
deps = [
":checkpoint",
":checkpoint_view",
"//tensorflow/python/eager:test",
"//tensorflow/python/trackable:autotrackable",
],
)
py_library(
name = "graph_view",
srcs = ["graph_view.py"],
strict_deps = True,
deps = [
":save_util_v1",
":trackable_view",
"//tensorflow/python/trackable:base",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "save_util",
srcs = ["save_util.py"],
strict_deps = True,
deps = [
":graph_view",
":save_util_v1",
":saveable_compat",
":util",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:python_state",
"//tensorflow/python/trackable:trackable_utils",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/types:core",
"//tensorflow/python/util:object_identity",
],
)
py_library(
name = "save_util_v1",
srcs = ["save_util_v1.py"],
strict_deps = True,
deps = [
":saveable_compat",
":util",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:python_state",
"//tensorflow/python/trackable:trackable_utils",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/util:object_identity",
],
)
tf_py_strict_test(
name = "save_util_v1_test",
srcs = ["save_util_v1_test.py"],
deps = [
":graph_view",
":save_util_v1",
"//tensorflow/python/eager:test",
"//tensorflow/python/ops:variables",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/util:object_identity",
],
)
py_library(
name = "trackable_view",
srcs = ["trackable_view.py"],
strict_deps = True,
tags = ["no_pip"],
deps = [
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:converter",
"//tensorflow/python/util:object_identity",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "trackable_view_test",
srcs = ["trackable_view_test.py"],
deps = [
":trackable_view",
"//tensorflow/python/eager:test",
"//tensorflow/python/trackable:base",
],
)
py_library(
name = "util",
srcs = ["util.py"],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:trackable_utils",
"//tensorflow/python/util:object_identity",
],
)
py_library(
name = "restore",
srcs = ["restore.py"],
strict_deps = True,
deps = [
":checkpoint_adapter",
":checkpoint_view",
":functional_saver",
":save_util_v1",
":saveable_compat",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/ops:io_ops_gen",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:constants",
"//tensorflow/python/trackable:python_state",
"//tensorflow/python/trackable:trackable_utils",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/util:object_identity",
],
)
tf_py_strict_test(
name = "restore_test",
srcs = ["restore_test.py"],
deps = [
":checkpoint",
":restore",
"//tensorflow/python/eager:test",
"//tensorflow/python/module",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:autotrackable",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training/saving:saveable_object",
],
)
tf_py_benchmark_test(
name = "benchmarks_test",
srcs = ["benchmarks_test.py"],
deps = [
":checkpoint",
"//tensorflow/python/framework:ops",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:io_ops_gen",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training:py_checkpoint_reader",
],
)
tf_py_logged_benchmark(
name = "benchmarks",
target = "//tensorflow/python/checkpoint:benchmarks_test",
)
py_library(
name = "checkpoint_options",
srcs = ["checkpoint_options.py"],
strict_deps = True,
deps = [
"//tensorflow/python/checkpoint/sharding:sharding_util",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
py_library(
name = "functional_saver",
srcs = ["functional_saver.py"],
strict_deps = True,
deps = [
":checkpoint_options",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/checkpoint/sharding:sharding_policies",
"//tensorflow/python/checkpoint/sharding:sharding_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/framework:tensor_util",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:io_ops",
"//tensorflow/python/ops:io_ops_gen",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/saved_model:pywrap_saved_model",
"//tensorflow/python/saved_model/registration",
"//tensorflow/python/trackable:base",
"//tensorflow/python/trackable:trackable_utils",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"//tensorflow/python/types:core",
"//tensorflow/python/util:nest",
"//tensorflow/python/util:object_identity",
"@absl_py//absl/logging",
],
)
cuda_py_strict_test(
name = "functional_saver_test",
size = "medium",
srcs = [
"functional_saver_test.py",
],
deps = [
":checkpoint",
":checkpoint_options",
":functional_saver",
":graph_view",
"//tensorflow/python/eager:context",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/eager:wrap_function",
"//tensorflow/python/framework:config",
"//tensorflow/python/framework:constant_op",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/module",
"//tensorflow/python/ops:io_ops_gen",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/training/saving:saveable_object_util",
],
)
py_library(
name = "tensor_callable",
srcs = ["tensor_callable.py"],
strict_deps = True,
deps = [
"//tensorflow/python/training/saving:saveable_object",
],
)
tf_py_strict_test(
name = "tensor_callable_test",
srcs = ["tensor_callable_test.py"],
deps = [
":checkpoint",
":tensor_callable",
"//tensorflow/python/eager:test",
"//tensorflow/python/ops:variables",
"//tensorflow/python/saved_model:save",
"//tensorflow/python/trackable:base",
],
)
py_library(
name = "checkpoint_management",
srcs = ["checkpoint_management.py"],
strict_deps = True,
deps = [
"//tensorflow/core:protos_all_py",
"//tensorflow/python/checkpoint:checkpoint_options",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:errors",
"//tensorflow/python/framework:ops",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:variable_scope",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/training:checkpoint_state_py",
"//tensorflow/python/training:training_util",
"//tensorflow/python/util:compat",
"//tensorflow/python/util:deprecation",
"//tensorflow/python/util:tf_export",
],
)
cuda_py_strict_test(
name = "checkpoint_management_test",
size = "small",
srcs = [
"checkpoint_management_test.py",
],
deps = [
":checkpoint",
":checkpoint_management",
"//tensorflow/core:protos_all_py",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/lib/io:file_io",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:client_testlib",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/platform:tf_logging",
"//tensorflow/python/training:checkpoint_state_py",
"//tensorflow/python/training:saver",
],
)
py_library(
name = "saveable_compat",
srcs = [
"saveable_compat.py",
],
strict_deps = True,
)
tf_py_strict_test(
name = "saveable_compat_test",
srcs = [
"saveable_compat_test.py",
],
data = [
"testdata/table_legacy_saveable_object.data-00000-of-00001",
"testdata/table_legacy_saveable_object.index",
],
tags = ["no_pip"],
deps = [
":checkpoint",
":generate_checkpoint_lib",
":saveable_compat",
"//tensorflow/python/eager:test",
"//tensorflow/python/ops:control_flow_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training:checkpoint_utils",
"//tensorflow/python/training/saving:saveable_object",
],
)
py_binary(
name = "generate_checkpoint",
srcs = ["testdata/generate_checkpoint.py"],
strict_deps = True,
deps = [
"//tensorflow/python/checkpoint",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/module",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:variables",
"@absl_py//absl:app",
],
)
py_library(
name = "generate_checkpoint_lib",
srcs = ["testdata/generate_checkpoint.py"],
strict_deps = True,
deps = [
"//tensorflow/python/checkpoint",
"//tensorflow/python/compat:v2_compat",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/module",
"//tensorflow/python/ops:lookup_ops",
"//tensorflow/python/ops:variables",
"@absl_py//absl:app",
],
)
tensorflow/python/checkpoint/async_checkpoint_helper.py
+628
View File
@@ -0,0 +1,628 @@
# Copyright 2022 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.
# ==============================================================================
"""Utilities for saving/loading Trackable objects asynchronously."""
import atexit
import copy
import queue
import threading
import time
import weakref
from absl import logging
from tensorflow.python.checkpoint import checkpoint_context
from tensorflow.python.checkpoint import trackable_view
from tensorflow.python.distribute import device_util
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.eager import executor
from tensorflow.python.framework import ops
from tensorflow.python.ops import variables
from tensorflow.python.saved_model.pywrap_saved_model import metrics
from tensorflow.python.trackable import base
from tensorflow.python.util import object_identity
# Captures the timestamp of the first Checkpoint instantiation or end of a write
# operation. Can be accessed by multiple Checkpoint instances.
_END_TIME_OF_LAST_ASYNC_WRITE = None
_END_TIME_OF_LAST_ASYNC_WRITE_LOCK = threading.Lock()
# API label for cell names used in async checkpoint metrics.
_ASYNC_CHECKPOINT = "async_checkpoint"
# Name of TPUEmbedding attribute. This is a temporary workaround
# to identify TPUEmbedding while avoiding import cycles.
_TPU_EMBEDDING_ATTR = "_create_copy_for_async_checkpoint"
def _get_duration_microseconds(start_time_seconds, end_time_seconds):
"""Calculate the duration between start and end time.
Args:
start_time_seconds: The start time in seconds.
end_time_seconds: The end time in seconds.
Returns:
The duration between the start and the end time. Return 0 if
end_time_seconds < start_time_seconds.
"""
if end_time_seconds < start_time_seconds:
# Avoid returning negative value in case of clock skew.
return 0
return round((end_time_seconds - start_time_seconds) * 1000000)
def _get_all_trackables(root, exclude_set):
"""Return the list of checkpointable trackables dependent on `root`.
Args:
root: The root trackable from where we get all its dependent trackables.
exclude_set: An ObjectIdentitySet of Trackables to exclude before returning.
Each element in `exclude_set` is a specific instance of a `Trackable`
and appears precisely once in `TrackableView(root).descendants()`.
Returns:
saveable_trackables: All trackables that are saveable in `all_trackables`
(see definition of "saveable" in `_trackable_needs_to_be_saved()`). A
subset of `all_trackables`.
all_trackables: All trackables returned by `TrackableView`'s `descendants()`
after excluding `exclude_set`. A superset of `saveable_trackables`.
"""
all_trackables = trackable_view.TrackableView(root=root).descendants()
# Kick out the trackable we want to exclude.
# The goal of writing such loop is to only scan the list once and stop
# scanning as early as possible (unlike filtering with list comprehension).
trackable_index = 0
while trackable_index < len(all_trackables) and exclude_set:
# While we have not excluded all items, or gone through all trackables.
if all_trackables[trackable_index] in exclude_set:
# If want to exclude this trackable, we pop it and do not update ptr
exclude_set.discard(all_trackables[trackable_index])
all_trackables.pop(trackable_index)
else:
# Otherwise update ptr
trackable_index += 1
# Kick out trackables that do not need to be saved (e.g. ListWrapper, etc.)
# We define any trackable that does not implement `_serialize_to_tensor` or
# `_gather_saveables` as "no need to be saved". If the trackable has one or
# both of the methods defined, it should have `_copy_trackable_to_cpu`
# defined; if not, we will raise warning in `_copy_to_cpu()`. In case of
# special case, we also check whether a trackable (who has neither of the
# other two methods defined) defines `_copy_trackable_to_cpu` only; we still
# define such cases as "needs to be saved".
def _trackable_needs_to_be_saved(obj):
"""Returns whether a trackable needs to be saved.
Returns a bool to indicate whether obj's class has `_serialize_to_tensors`,
`gather_saveables_for_checkpoint`, or `_copy_trackable_to_cpu` defined.
Args:
obj: A Trackable object.
"""
if hasattr(obj, "__dict__"):
# Data structure proxy wrappers don't have __dict__.
if ("_serialize_to_tensors" in obj.__dict__
or "_gather_saveables_for_checkpoint" in obj.__dict__
or "_copy_trackable_to_cpu" in obj.__dict__):
return True
# Use MRO so that if a parent class has one of the three methods, we still
# consider `t` as needed to be saved.
for t in type(obj).mro():
if t is base.Trackable:
# Base class always has them implemented, but would raise error.
continue
elif ("_serialize_to_tensors" in t.__dict__
or "_gather_saveables_for_checkpoint" in t.__dict__
or "_copy_trackable_to_cpu" in t.__dict__):
return True
return False
saveable_trackables = [x for x in all_trackables if
_trackable_needs_to_be_saved(x)]
return saveable_trackables, all_trackables
class AsyncCheckpointHelper:
"""Helper class for async checkpoint."""
def __init__(self, checkpointer_impl, root=None, **kwargs):
"""Initialize AsyncCheckpoint.
Args:
checkpointer_impl: The Checkpoint class to power the AsyncCheckpoint.
root: The root object to checkpoint. `root` may be a trackable object or
`WeakRef` of a trackable object.
**kwargs: The keyword arguments representing the checkpointed variables.
Raises:
AttributeError: when checkpointer_impl is None.
"""
# TODO(chienchunh): Make sure the processing for the root object is
# consistent when integrating with the public API, e.g., adding all kwarg
# items as the child of the root object.
if root:
trackable_root = root() if isinstance(root, weakref.ref) else root
kwargs["root"] = trackable_root
trackable_root._maybe_initialize_trackable()
# The underlying Checkpoint instance and its items.
if checkpointer_impl is None:
raise AttributeError(
"checkpointer_impl cannot be None for AsyncCheckpointHelper."
)
self._checkpointer_impl = checkpointer_impl
self._checkpoint_items = kwargs
self._checkpoint = None
self.checkpointer()
self._checkpoint_options = None
# Indicate whether async checkpoint has finished traversing the variable
# list and created the object map between the original and copied variables.
self._initialized = False
# The list of all nodes from the original checkpoint items.
# TODO(chienchunh): Consider changing this to local variable.
self._original_nodes = None
# The mapping between the original and the copied resource variables.
# The copied variables are used for the underlying checkpointing.
self._object_map = None
# A list of TPUEmbedding objects included in the checkpoint items.
self._tpu_embedding_objects = None
# A list of highest level `Trackable`s we will copy; does not contain
# TPUEmbedding objects
self._saveable_trackables = None
self._default_device = device_util.current() or "CPU:0"
self._default_device = device_util.canonicalize(self._default_device)
self._save_file_prefix = None
self._use_checkpoint_save = False
self._async_save_thread = None
# Concurrent queue that coordinates the events for writing/reading the
# cpu-copied variables. A 'True' in the queue triggers the async thread to
# perform saving; a 'False' breaks the while loop so that the async thread
# exits; no other values will be added to the queue.
# Maxsize is set to 1 only to ensure the exit procedure. We could have used
# queue.join() in _join_async_save_thread(), but queue.join() does not have
# a timeout argument. Hence we use queue.put(timeout=300), in case the last
# checkpoint takes forever. To achieve that, maxsize needs to be 1.
self._queue = queue.Queue(maxsize=1)
# Register to join the async save thread upon exit.
atexit.register(self._join_async_save_thread)
self._async_error = None
global _END_TIME_OF_LAST_ASYNC_WRITE
with _END_TIME_OF_LAST_ASYNC_WRITE_LOCK:
if _END_TIME_OF_LAST_ASYNC_WRITE is None:
_END_TIME_OF_LAST_ASYNC_WRITE = time.time()
@def_function.function
def _copy_to_cpu(self):
"""Copy the checkpointed variables from the accelerator to the host CPU.
TODO(chienchunh): Get the concrete function before firstly called to avoid
hangining the accelerators idle during function tracing.
"""
for t in self._saveable_trackables:
try:
t._copy_trackable_to_cpu(object_map=self._object_map) # pylint: disable=protected-access
except NotImplementedError as e:
logging.warning("Trackable %s skipped due to: %s", t, e)
for tpu_embedding in self._tpu_embedding_objects:
tpu_embedding._retrieve_variables() # pylint: disable=protected-access
def checkpointer(self):
"""Gets or creates the underlying Checkpoint instance."""
if self._checkpoint is None:
self._checkpoint = self._checkpointer_impl(**self._checkpoint_items)
return self._checkpoint
def _ensure_initialized(self):
"""Initialize the async checkpoint internal state."""
# This map will be used to store the CPU copy of all checkpointable objects
self._object_map = object_identity.ObjectIdentityDictionary()
self._tpu_embedding_objects = []
# Populate self._all_tracakbles, but exclude the checkpoint instance itself
# and its save_counter, as they will be returned by `descendants()`.
exclude_set = object_identity.ObjectIdentitySet()
exclude_set.add(self.checkpointer())
exclude_set.add(self.checkpointer().save_counter)
self._saveable_trackables, all_trackables = _get_all_trackables(
root=self.checkpointer(), exclude_set=exclude_set)
# Handle special cases: TPU Embedding, and slot variables.
# 1. TPUEmbedding: Different from other trackables, TPUEmbedding needs to
# call `_retrieve_variables` to checkpoint, while populating a dummy copy to
# the object map.
# 2. Slot variables: they need to be handled differently as they cannot be
# retrieved from `TrackableView.descendants()`.
# Note: dir() is used rather than hasattr() here to avoid triggering
# custom __getattr__ code, see b/152031870 for context.
for t in all_trackables:
# Special case 1: TPU Embedding, populate object_map here
# Special case 1: Handle TPU Embedding by addnig a dummy instance to the
# object map. Also add TPUEmbedding to separate list for special handling
# with values copy.
if hasattr(type(t), _TPU_EMBEDDING_ATTR):
self._handle_tpu_embedding(t)
# Special case 2: handle slot variables. The object_map is populated later
# when the variable values are being copied to host CPU for the first
# time.
if "get_slot_names" in dir(t):
slot_names = t.get_slot_names()
for slot_name in slot_names:
for original_variable in all_trackables:
if not isinstance(original_variable, variables.Variable):
continue
try:
# Usage of hasattr may result in KeyError
original_slot_variable = t.get_slot(original_variable, slot_name)
except (AttributeError, KeyError):
continue
if isinstance(original_slot_variable, base.Trackable):
self._saveable_trackables.append(original_slot_variable)
# Initiate the underlying Checkpoint instance's save_counter.
save_counter = self.checkpointer().save_counter.numpy()
logging.info("Initializing async checkpoint's save_counter: %d",
save_counter)
# Pass the object map of the copied variables to the underlying Checkpoint.
self.checkpointer()._saver._object_map = self._object_map # pylint: disable=protected-access
# We perform a `_copy_to_cpu()` to populate `self._object_map`,
# initializing copies. We do not call `self._copy_to_cpu()` directly
# because it is a tf function, which leads to access out of scope error.
# TODO(charlieruan) Figure out a better work around to solve the access
# out of scope error.
for t in self._saveable_trackables:
try:
t._copy_trackable_to_cpu(object_map=self._object_map) # pylint: disable=protected-access
except NotImplementedError as e:
logging.warning("Trackable %s skipped due to: %s", t, e)
for tpu_embedding in self._tpu_embedding_objects:
tpu_embedding._retrieve_variables() # pylint: disable=protected-access
# Initiate the async thread for checkpoint saving.
self._async_save_thread = threading.Thread(
target=self._async_save, daemon=True)
self._async_save_thread.start()
self._initialized = True
def _check_async_thread_error(self):
"""Expose the most recent error from the async saving thread to the caller.
"""
if self._async_error:
e = self._async_error
self._async_error = None
logging.error("Propagating the most recent error from the async thread "
"before joining: %s", str(e))
raise e
def _join_async_save_thread(self):
"""Join the async save thread.
The steps for terminating the async save thread:
1). Put will succeed when the last async save event is done. Putting a false
triggers the async save thread's while loop to end. We use put instead
of sync because sync does not have a timeout argument.
2). Join the async save thread. (The thread may finish before joining.)
"""
try:
self._queue.put(False, timeout=300) # Step-1.
logging.info("Joining the async save thread.")
if self._async_save_thread is not None:
self._async_save_thread.join() # Step-2.
except queue.Full:
logging.error("Timeout waiting for the async save thread; terminating the"
" thread instead. The last checkpoint may be incomeplete.")
finally:
self._check_async_thread_error()
def _async_save(self):
"""The thread function for the async checkpoint save."""
with context.executor_scope(
executor.new_executor(
enable_async=False, enable_streaming_enqueue=False)):
# The main thread inserts: a True to the queue when the user calls save,
# triggering async save; and a False when we exit the Checkpoint instance.
while self._queue.get():
logging.info("Starting async checkpoint save on the device: %s",
self._default_device)
async_save_start_time = time.time()
# Specify the ops placement on the worker if running with
# coordinator-worker mode. This is required as launching a new thread
# would clear the placement policy and make localhost the default
# placement, while the main thread's default placement would be the
# master worker's CPU:0.
try:
with ops.device(self._default_device):
with checkpoint_context.async_metrics_context():
if self._use_checkpoint_save:
self.checkpointer().save(
self._save_file_prefix, self._checkpoint_options
)
else:
self.checkpointer()._write( # pylint: disable=protected-access
self._save_file_prefix,
options=self._checkpoint_options,
)
except Exception as e: # # pylint: disable=broad-except
self._async_error = e
finally:
self._queue.task_done()
async_save_end_time = time.time()
metrics.AddAsyncCheckpointWriteDuration(
api_label=_ASYNC_CHECKPOINT,
microseconds=_get_duration_microseconds(async_save_start_time,
async_save_end_time))
# Measure the elapsed time since the last checkpoint.
# Due to the nature of async checkpoint, here it actually captures the
# duration between the start_time of the previous checkpoint and the
# start time of this checkpoint. As a result, the duration of the final
# async checkpoint is excluded, which is fine since it does not take
# much time.
global _END_TIME_OF_LAST_ASYNC_WRITE
with _END_TIME_OF_LAST_ASYNC_WRITE_LOCK:
metrics.AddTrainingTimeSaved(
api_label=_ASYNC_CHECKPOINT,
microseconds=_get_duration_microseconds(
_END_TIME_OF_LAST_ASYNC_WRITE, async_save_start_time))
_END_TIME_OF_LAST_ASYNC_WRITE = async_save_start_time
logging.info("Async save thread reached the end of the execution.")
def _handle_tpu_embedding(self, tpu_embedding):
"""Handle TPUEmbedding.
This is the only place where we populate object map in the class of
`AsyncCheckpointHelper`. For all other checkpointable trackables, we
populate object map using the trackable's own `_copy_trackable_to_cpu()`.
Args:
tpu_embedding: TPUEmbedding object to be handled.
Raises:
AttributeError: if the input trackable is not TPUEmbedding type.
"""
if not hasattr(type(tpu_embedding), _TPU_EMBEDDING_ATTR) or not callable(
tpu_embedding._create_copy_for_async_checkpoint # pylint: disable=protected-access
):
raise AttributeError(
"Expecting TPUEmbedding type; got %s" % type(tpu_embedding)
)
# Create a dummy TPUEmbedding object and add it to the object_map. This is
# to prevent the TPUEmbedding's save_callback from being triggered because
# the embedding values have already being retrieved by AsyncCheckpoint.
# pylint: disable=protected-access
new_embedding = tpu_embedding._create_copy_for_async_checkpoint(
feature_config=tpu_embedding._feature_config,
optimizer=tpu_embedding._table_config[0]
if tpu_embedding._table_config
else None,
pipeline_execution_with_tensor_core=tpu_embedding._pipeline_execution_with_tensor_core,
)
self._object_map[tpu_embedding] = new_embedding
# pylint: enable=protected-access
if tpu_embedding not in self._tpu_embedding_objects:
self._tpu_embedding_objects.append(tpu_embedding)
@property
def save_counter(self):
"""An integer variable numbering the checkpoint events.
This is maintained by the underlying tf.train.Checkpoint object employed by
AsyncCheckpoint class. The number starts at 0 and gets incremented for each
checkpoint event.
Returns:
The save counter variable.
"""
return self.checkpointer().save_counter
def write(self, save_path, options=None):
"""Save the checkpointed variables.
Args:
save_path: The file prefix of the checkpoint file.
options: Optional CheckpointOption instance.
Returns:
The full path of the checkpoint file.
"""
return self._write(save_path, options)
def _write(self, save_path, options=None):
"""Save the checkpointed variables.
This method has exactly the same logic as save(), except it does not
increment the underlying save_counter, which is done by the caller, e.g.,
CheckpointManager.
Args:
save_path: The file prefix of the checkpoint file.
options: Optional CheckpointOption instance.
Returns:
The full path of the checkpoint file.
"""
write_start_time = time.time()
if not self._initialized:
self._ensure_initialized()
else:
# First wait for async thread to finish the previous save, then copy the
# variable values to the host CPU.
self._queue.join()
self._copy_to_cpu()
# Surface the error from the async thread, if any.
# This step should come after the sem acquisition step in the above, so that
# it makes sure it waits until the previous async save finishes storing the
# error.
self._check_async_thread_error()
# Trigger the async thread to checkpoint the cpu-copied variables.
# Need to wait until the weight copying finishes before checkpoint save.
context.async_wait()
self._save_file_prefix = save_path
self._use_checkpoint_save = False
# Ensure that we do not request async checkpointing to the underlying
# checkpointer as this could lead to an infinite loop.
self._checkpoint_options = copy.copy(options) if options else None
if self._checkpoint_options:
self._checkpoint_options.experimental_enable_async_checkpoint = False
self._queue.put(True) # Trigger save in async thread
write_end_time = time.time()
metrics.AddCheckpointWriteDuration(
api_label=_ASYNC_CHECKPOINT,
microseconds=_get_duration_microseconds(write_start_time,
write_end_time))
return save_path
def save(self, save_path, options=None):
"""Save the checkpointed variables.
Args:
save_path: The file prefix of the checkpoint file.
options: Optional CheckpointOption instance.
Returns:
The full path of the checkpoint file.
"""
save_start_time = time.time()
# If this is the first time that AsyncCheckpoint.save() is called,
# initialize the internal states like `self._saveable_trackables`. We also
# populate `self._object_map` (i.e. initializing the cpu-copied variables
# and copy over the value for the first time) by essentially performing a
# `self._copy_to_cpu()`, hence the if-else logic here.
#
# This is not performed in the initializer because some variables, e.g.,
# slot variables of the optimizer, were not created until actually running
# the train function, so we could only get the complete list of the
# variables after some train steps were run.
if not self._initialized:
self._ensure_initialized()
else:
# First wait for async thread to finish the previous save, then copy the
# variable values to the host CPU.
self._queue.join()
self._copy_to_cpu()
# Surface the error from the async thread, if any.
# This step should come after the sem acquisition step in the above, so that
# it makes sure it waits until the previous async save finishes storing the
# error.
self._check_async_thread_error()
# Retrieve the save counter from the underlying checkpoint object to
# re-construct the full path of the checkpoint file.
# This step has to happen before triggering the underlying checkpoint;
# otherwise, the save_counter value may or may not have been updated.
save_counter = self.checkpointer().save_counter.numpy() + 1
full_path = "{}-{}".format(save_path, save_counter)
# Trigger the async thread to checkpoint the cpu-copied variables.
# Need to wait until the weight copying finishes before checkpoint save.
context.async_wait()
self._save_file_prefix = save_path
self._use_checkpoint_save = True
# Ensure that we do not request async checkpointing to the underlying
# checkpointer as this could lead to an infinite loop.
self._checkpoint_options = copy.copy(options) if options else None
if self._checkpoint_options:
self._checkpoint_options.experimental_enable_async_checkpoint = False
self._queue.put(True) # Trigger save in async thread
save_end_time = time.time()
metrics.AddCheckpointWriteDuration(
api_label=_ASYNC_CHECKPOINT,
microseconds=_get_duration_microseconds(save_start_time, save_end_time))
return full_path
def read(self, save_path, options=None):
"""Restore the checkpointed variables.
This method has exactly the same logic as restore(). This method is
implemented only to fulfill the duty of subclassing tf.train.Checkpoint.
Args:
save_path: The full name of the checkpoint file to be restored.
options: CheckpointOption instance.
Returns:
A load status object, which can be used to make assertions about the
status of a checkpoint restoration. See tf.train.Checkpoint.restore()
for more details.
"""
return self.restore(save_path, options)
def restore(self, save_path, options=None):
"""Restore the checkpointed variables.
Args:
save_path: The full name of the checkpoint file to be restored.
options: CheckpointOption instance.
Returns:
A load status object, which can be used to make assertions about the
status of a checkpoint restoration. See tf.train.Checkpoint.restore()
for more details.
"""
# Ensure that we do not request async checkpointing to the underlying
# checkpointer as this could lead to an infinite loop.
self._checkpoint_options = (
copy.copy(options) if options else self._checkpoint_options)
if self._checkpoint_options:
self._checkpoint_options.experimental_enable_async_checkpoint = False
# Wait for any ongoing checkpoint event to finish.
self._queue.join()
# Restore values of the cpu-copied variables directly back to accelerators
status = self.checkpointer().restore(save_path, self._checkpoint_options)
return status
def sync(self):
"""Sync on any ongoing save or restore events."""
self._queue.join()
logging.info("Sync on ongoing save/restore.")
tensorflow/python/checkpoint/benchmarks_test.py
+113
View File
@@ -0,0 +1,113 @@
# Copyright 2019 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.
# ==============================================================================
"""Benchmarks for checkpoint-related APIs."""
import os
import time
from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.framework import ops
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_io_ops
from tensorflow.python.platform import test
from tensorflow.python.trackable import base
from tensorflow.python.training import py_checkpoint_reader
class _TrivialRestore(base.Trackable):
def _serialize_to_tensors(self):
return {base.VARIABLE_VALUE_KEY: array_ops.ones([])}
def _restore_from_tensors(self, restored_tensors):
return control_flow_ops.no_op()
class _LazyTrivialObjects(module.Module):
def __init__(self):
self.existing = [_TrivialRestore() for _ in range(5)]
self.lazy = []
def __call__(self):
if not self.lazy:
self.lazy.extend(_TrivialRestore() for _ in range(5))
return
def _save_checkpoint():
original_checkpoint = util.Checkpoint(m=_LazyTrivialObjects())
original_checkpoint.m()
return original_checkpoint.write(os.path.join(test.get_temp_dir(), "ckpt"))
class SavingBenchmarks(test.Benchmark):
def _run(self, func, num_iters, execution_mode=None):
func()
start = time.time()
for _ in range(num_iters):
func()
end = time.time()
mean_us = (end - start) * 1e6 / num_iters
self.report_benchmark(
iters=num_iters,
wall_time=mean_us,
extras={"examples_per_sec": num_iters / (end - start)})
def benchmark_baseline_no_restore(self):
def _create_and_call():
checkpoint = util.Checkpoint(m=_LazyTrivialObjects())
checkpoint.m()
self._run(_create_and_call, 3)
def benchmark_batch_restore(self):
checkpoint_path = _save_checkpoint()
def _create_and_call():
checkpoint = util.Checkpoint(m=_LazyTrivialObjects())
checkpoint.m()
checkpoint.restore(checkpoint_path)
self._run(_create_and_call, 3)
def benchmark_restore_on_create(self):
checkpoint_path = _save_checkpoint()
def _create_and_call():
checkpoint = util.Checkpoint(m=_LazyTrivialObjects())
checkpoint.restore(checkpoint_path)
checkpoint.m()
self._run(_create_and_call, 3)
def benchmark_raw_restore(self):
checkpoint_path = _save_checkpoint()
all_names, all_dtypes = zip(*py_checkpoint_reader.NewCheckpointReader(
checkpoint_path).get_variable_to_dtype_map().items())
def _call_restore_v2():
gen_io_ops.restore_v2(checkpoint_path, all_names, [""] * len(all_names),
all_dtypes)
self._run(_call_restore_v2, 3)
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/checkpoint/checkpoint.py
+2753
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/checkpoint/checkpoint_adapter.py
+119
View File
@@ -0,0 +1,119 @@
# Copyright 2024 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.
# ==============================================================================
"""Experimental API for checkpoint adapter."""
import abc
from typing import List, Optional
from tensorflow.python.framework import tensor
from tensorflow.python.trackable import base
class ReshardCallback:
"""API to reshard a checkpoint value during restore.
When a ReshardCallback is attached to a CheckpointPosition, the restored value
of the checkpoint position is resharded based on this callback.
"""
def object_name(self) -> str:
"""Returns the local name of the object being restored.
Override this method when the local name of object is different than in the
checkpoint.
"""
return None
def reshard(
self,
checkpoint_values: List[tensor.Tensor],
shape_and_slice_spec: List[str],
) -> tensor.Tensor:
"""Reshards the checkpoint values as read from the checkpoint file.
Override this to reshard/modify the restored values
Args:
checkpoint_values: The values returned by the restore op, as read from
file.
shape_and_slice_spec: The shape and slice spec required by the caller.
Returns:
List of restored Tensor values after being resharded.
"""
del shape_and_slice_spec # unused
# Default reshard is a trivial one.
if len(checkpoint_values) != 1:
raise ValueError("Default reshard expects a single checkpoint value.")
return checkpoint_values[0]
def update_restore_inputs(
self, checkpoint_key, shape_and_slice_spec
) -> tuple[List[str], List[str]]:
"""Updates the specs to restore op.
Override this method if the arguments to restore op need to be updated as
per the resharding required.
Args:
checkpoint_key: The checkpoint key as requested by the caller
shape_and_slice_spec: The shape and slice spec as requested by caller
Returns:
Tuple of list of checkpoint_keys and specs that the restore op should fetch
as per the resharding requirement. The length of checkpoint keys returned by
this method will match the length of checkpoint_values that are input to
`reshard`.
"""
return ([checkpoint_key], [shape_and_slice_spec])
class AbstractCheckpointAdapter(abc.ABC):
"""Abstract API for checkpoint adapter.
This is an experimental API that specifies how checkpoint restore should be
adapted for specific trackable objects.
"""
@classmethod
@abc.abstractmethod
def create_from_checkpoint(cls, path: str):
"""Create factory to create an Adapter from checkpoint.
Args:
path: Path to checkpoint.
"""
@abc.abstractmethod
def is_applicable(self, trackable: base.Trackable) -> bool:
"""Returns whether the adapter is applicable to trackable for resharding.
Args:
trackable: A Trackable object that is being restored.
Returns:
A Boolean indicating if the checkpoint value for this Trackable should be
resharded.
"""
@abc.abstractmethod
def get_reshard_callback(self, name: str) -> Optional[ReshardCallback]:
"""Returns the reshard callback for the trackable with `name`."""
def maybe_reshard(self, name: str) -> tuple[str, Optional[ReshardCallback]]:
"""Returns the updated name and ReshardCallback applicable to it."""
callback = self.get_reshard_callback(name)
if callback is None:
return name, None
if callback.object_name():
return callback.object_name(), callback
return name, callback
tensorflow/python/checkpoint/checkpoint_context.py
+85
View File
@@ -0,0 +1,85 @@
# Copyright 2023 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.
# ==============================================================================
"""Context for saving checkpoint."""
import contextlib
import threading
class PreemptionSaveContext(threading.local):
"""A context for saving checkpoint upon preemption."""
def __init__(self):
super().__init__()
self._in_preemption_save_context = False
def enter_preemption_save_context(self):
self._in_preemption_save_context = True
def exit_preemption_save_context(self):
self._in_preemption_save_context = False
def in_preemption_save_context(self):
return self._in_preemption_save_context
_preemption_save_context = PreemptionSaveContext()
@contextlib.contextmanager
def preemption_save_context():
_preemption_save_context.enter_preemption_save_context()
try:
yield
finally:
_preemption_save_context.exit_preemption_save_context()
def in_preemption_save_context():
return _preemption_save_context.in_preemption_save_context()
class AsyncMetricsContext(threading.local):
"""A context for controlling metrics recording when async checkpoint is used.
"""
def __init__(self):
super().__init__()
self._in_async_metrics_context = False
def enter_async_metrics_context(self):
self._in_async_metrics_context = True
def exit_async_metrics_context(self):
self._in_async_metrics_context = False
def in_async_metrics_context(self):
return self._in_async_metrics_context
_async_metrics_context = AsyncMetricsContext()
@contextlib.contextmanager
def async_metrics_context():
_async_metrics_context.enter_async_metrics_context()
try:
yield
finally:
_async_metrics_context.exit_async_metrics_context()
def in_async_metrics_context():
return _async_metrics_context.in_async_metrics_context()
tensorflow/python/checkpoint/checkpoint_management.py
+902
View File
@@ -0,0 +1,902 @@
# Copyright 2015 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.
# ==============================================================================
# pylint: disable=invalid-name
"""Checkpoint Manager and other utilities for managing checkpoints."""
import collections
import copy
import os.path
import re
import time
from google.protobuf import text_format
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.eager import context
from tensorflow.python.framework import errors
from tensorflow.python.framework import ops
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import variable_scope
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import training_util
from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState
from tensorflow.python.util import compat
from tensorflow.python.util import deprecation
from tensorflow.python.util.tf_export import tf_export
def _evaluate(tensor):
"""Returns the numpy value of a tensor."""
if context.executing_eagerly():
return tensor.numpy()
return ops.get_default_session().run(tensor)
def _GetCheckpointFilename(save_dir, latest_filename):
"""Returns a filename for storing the CheckpointState.
Args:
save_dir: The directory for saving and restoring checkpoints.
latest_filename: Name of the file in 'save_dir' that is used
to store the CheckpointState.
Returns:
The path of the file that contains the CheckpointState proto.
"""
if latest_filename is None:
latest_filename = "checkpoint"
return os.path.join(save_dir, latest_filename)
@tf_export(v1=["train.generate_checkpoint_state_proto"])
def generate_checkpoint_state_proto(save_dir,
model_checkpoint_path,
all_model_checkpoint_paths=None,
all_model_checkpoint_timestamps=None,
last_preserved_timestamp=None):
"""Generates a checkpoint state proto.
Args:
save_dir: Directory where the model was saved.
model_checkpoint_path: The checkpoint file.
all_model_checkpoint_paths: List of strings. Paths to all not-yet-deleted
checkpoints, sorted from oldest to newest. If this is a non-empty list,
the last element must be equal to model_checkpoint_path. These paths
are also saved in the CheckpointState proto.
all_model_checkpoint_timestamps: A list of floats, indicating the number of
seconds since the Epoch when each checkpoint was generated.
last_preserved_timestamp: A float, indicating the number of seconds since
the Epoch when the last preserved checkpoint was written, e.g. due to a
`keep_checkpoint_every_n_hours` parameter (see
`tf.train.CheckpointManager` for an implementation).
Returns:
CheckpointState proto with model_checkpoint_path and
all_model_checkpoint_paths updated to either absolute paths or
relative paths to the current save_dir.
Raises:
ValueError: If `all_model_checkpoint_timestamps` was provided but its length
does not match `all_model_checkpoint_paths`.
"""
if all_model_checkpoint_paths is None:
all_model_checkpoint_paths = []
if (not all_model_checkpoint_paths or
all_model_checkpoint_paths[-1] != model_checkpoint_path):
logging.info("%s is not in all_model_checkpoint_paths. Manually adding it.",
model_checkpoint_path)
all_model_checkpoint_paths.append(model_checkpoint_path)
if (all_model_checkpoint_timestamps
and (len(all_model_checkpoint_timestamps)
!= len(all_model_checkpoint_paths))):
raise ValueError(
("Checkpoint timestamps, if provided, must match checkpoint paths (got "
"paths %s and timestamps %s)")
% (all_model_checkpoint_paths, all_model_checkpoint_timestamps))
# Relative paths need to be rewritten to be relative to the "save_dir"
# if model_checkpoint_path already contains "save_dir".
if not os.path.isabs(save_dir):
if not os.path.isabs(model_checkpoint_path):
model_checkpoint_path = os.path.relpath(model_checkpoint_path, save_dir)
for i, p in enumerate(all_model_checkpoint_paths):
if not os.path.isabs(p):
all_model_checkpoint_paths[i] = os.path.relpath(p, save_dir)
coord_checkpoint_proto = CheckpointState(
model_checkpoint_path=model_checkpoint_path,
all_model_checkpoint_paths=all_model_checkpoint_paths,
all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
last_preserved_timestamp=last_preserved_timestamp)
return coord_checkpoint_proto
@deprecation.deprecated(
date=None,
instructions=("Use `tf.train.CheckpointManager` to manage checkpoints "
"rather than manually editing the Checkpoint proto."))
@tf_export(v1=["train.update_checkpoint_state"])
def update_checkpoint_state(save_dir,
model_checkpoint_path,
all_model_checkpoint_paths=None,
latest_filename=None,
all_model_checkpoint_timestamps=None,
last_preserved_timestamp=None):
"""Updates the content of the 'checkpoint' file.
This updates the checkpoint file containing a CheckpointState
proto.
Args:
save_dir: Directory where the model was saved.
model_checkpoint_path: The checkpoint file.
all_model_checkpoint_paths: List of strings. Paths to all not-yet-deleted
checkpoints, sorted from oldest to newest. If this is a non-empty list,
the last element must be equal to model_checkpoint_path. These paths
are also saved in the CheckpointState proto.
latest_filename: Optional name of the checkpoint file. Default to
'checkpoint'.
all_model_checkpoint_timestamps: Optional list of timestamps (floats,
seconds since the Epoch) indicating when the checkpoints in
`all_model_checkpoint_paths` were created.
last_preserved_timestamp: A float, indicating the number of seconds since
the Epoch when the last preserved checkpoint was written, e.g. due to a
`keep_checkpoint_every_n_hours` parameter (see
`tf.train.CheckpointManager` for an implementation).
Raises:
RuntimeError: If any of the model checkpoint paths conflict with the file
containing CheckpointSate.
"""
update_checkpoint_state_internal(
save_dir=save_dir,
model_checkpoint_path=model_checkpoint_path,
all_model_checkpoint_paths=all_model_checkpoint_paths,
latest_filename=latest_filename,
save_relative_paths=False,
all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
last_preserved_timestamp=last_preserved_timestamp)
@tf_export("__internal__.train.update_checkpoint_state", v1=[])
def update_checkpoint_state_internal(save_dir,
model_checkpoint_path,
all_model_checkpoint_paths=None,
latest_filename=None,
save_relative_paths=False,
all_model_checkpoint_timestamps=None,
last_preserved_timestamp=None):
"""Updates the content of the 'checkpoint' file.
This updates the checkpoint file containing a CheckpointState
proto.
Args:
save_dir: Directory where the model was saved.
model_checkpoint_path: The checkpoint file.
all_model_checkpoint_paths: List of strings. Paths to all not-yet-deleted
checkpoints, sorted from oldest to newest. If this is a non-empty list,
the last element must be equal to model_checkpoint_path. These paths
are also saved in the CheckpointState proto.
latest_filename: Optional name of the checkpoint file. Default to
'checkpoint'.
save_relative_paths: If `True`, will write relative paths to the checkpoint
state file.
all_model_checkpoint_timestamps: Optional list of timestamps (floats,
seconds since the Epoch) indicating when the checkpoints in
`all_model_checkpoint_paths` were created.
last_preserved_timestamp: A float, indicating the number of seconds since
the Epoch when the last preserved checkpoint was written, e.g. due to a
`keep_checkpoint_every_n_hours` parameter (see
`tf.train.CheckpointManager` for an implementation).
Raises:
RuntimeError: If any of the model checkpoint paths conflict with the file
containing CheckpointSate.
"""
# Writes the "checkpoint" file for the coordinator for later restoration.
coord_checkpoint_filename = _GetCheckpointFilename(save_dir, latest_filename)
if save_relative_paths:
if os.path.isabs(model_checkpoint_path):
rel_model_checkpoint_path = os.path.relpath(
model_checkpoint_path, save_dir)
else:
rel_model_checkpoint_path = model_checkpoint_path
rel_all_model_checkpoint_paths = []
for p in all_model_checkpoint_paths:
if os.path.isabs(p):
rel_all_model_checkpoint_paths.append(os.path.relpath(p, save_dir))
else:
rel_all_model_checkpoint_paths.append(p)
ckpt = generate_checkpoint_state_proto(
save_dir,
rel_model_checkpoint_path,
all_model_checkpoint_paths=rel_all_model_checkpoint_paths,
all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
last_preserved_timestamp=last_preserved_timestamp)
else:
ckpt = generate_checkpoint_state_proto(
save_dir,
model_checkpoint_path,
all_model_checkpoint_paths=all_model_checkpoint_paths,
all_model_checkpoint_timestamps=all_model_checkpoint_timestamps,
last_preserved_timestamp=last_preserved_timestamp)
if coord_checkpoint_filename == ckpt.model_checkpoint_path:
raise RuntimeError("Save path '%s' conflicts with path used for "
"checkpoint state. Please use a different save path." %
model_checkpoint_path)
# Preventing potential read/write race condition by *atomically* writing to a
# file.
file_io.atomic_write_string_to_file(coord_checkpoint_filename,
text_format.MessageToString(ckpt))
@tf_export("train.get_checkpoint_state")
def get_checkpoint_state(checkpoint_dir, latest_filename=None):
"""Returns CheckpointState proto from the "checkpoint" file.
If the "checkpoint" file contains a valid CheckpointState
proto, returns it.
Args:
checkpoint_dir: The directory of checkpoints.
latest_filename: Optional name of the checkpoint file. Default to
'checkpoint'.
Returns:
A CheckpointState if the state was available, None
otherwise.
Raises:
ValueError: if the checkpoint read doesn't have model_checkpoint_path set.
"""
if isinstance(checkpoint_dir, os.PathLike):
checkpoint_dir = os.fspath(checkpoint_dir)
ckpt = None
coord_checkpoint_filename = _GetCheckpointFilename(checkpoint_dir,
latest_filename)
f = None
try:
# Check that the file exists before opening it to avoid
# many lines of errors from colossus in the logs.
if file_io.file_exists(coord_checkpoint_filename):
file_content = file_io.read_file_to_string(
coord_checkpoint_filename)
ckpt = CheckpointState()
text_format.Merge(file_content, ckpt)
if not ckpt.model_checkpoint_path:
raise ValueError("Invalid checkpoint state loaded from "
+ checkpoint_dir)
# For relative model_checkpoint_path and all_model_checkpoint_paths,
# prepend checkpoint_dir.
if not os.path.isabs(ckpt.model_checkpoint_path):
ckpt.model_checkpoint_path = os.path.join(checkpoint_dir,
ckpt.model_checkpoint_path)
for i, p in enumerate(ckpt.all_model_checkpoint_paths):
if not os.path.isabs(p):
ckpt.all_model_checkpoint_paths[i] = os.path.join(checkpoint_dir, p)
except errors.OpError as e:
# It's ok if the file cannot be read
logging.warning("%s: %s", type(e).__name__, e)
logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)
return None
except text_format.ParseError as e:
logging.warning("%s: %s", type(e).__name__, e)
logging.warning("%s: Checkpoint ignored", coord_checkpoint_filename)
return None
finally:
if f:
f.close()
return ckpt
def _prefix_to_checkpoint_path(prefix, format_version):
"""Returns the pathname of a checkpoint file, given the checkpoint prefix.
For V1 checkpoint, simply returns the prefix itself (the data file). For V2,
returns the pathname to the index file.
Args:
prefix: a string, the prefix of a checkpoint.
format_version: the checkpoint format version that corresponds to the
prefix.
Returns:
The pathname of a checkpoint file, taking into account the checkpoint
format version.
"""
if format_version == saver_pb2.SaverDef.V2:
return prefix + ".index" # The index file identifies a checkpoint.
return prefix # Just the data file.
@tf_export("train.latest_checkpoint")
def latest_checkpoint(checkpoint_dir, latest_filename=None):
"""Finds the filename of latest saved checkpoint file.
Gets the checkpoint state given the provided checkpoint_dir and looks for a
corresponding TensorFlow 2 (preferred) or TensorFlow 1.x checkpoint path.
The latest_filename argument is only applicable if you are saving checkpoint
using `v1.train.Saver.save`
See the [Training Checkpoints
Guide](https://www.tensorflow.org/guide/checkpoint) for more details and
examples.`
Args:
checkpoint_dir: Directory where the variables were saved.
latest_filename: Optional name for the protocol buffer file that
contains the list of most recent checkpoint filenames.
See the corresponding argument to `v1.train.Saver.save`.
Returns:
The full path to the latest checkpoint or `None` if no checkpoint was found.
"""
# Pick the latest checkpoint based on checkpoint state.
ckpt = get_checkpoint_state(checkpoint_dir, latest_filename)
if ckpt and ckpt.model_checkpoint_path:
# Look for either a V2 path or a V1 path, with priority for V2.
v2_path = _prefix_to_checkpoint_path(ckpt.model_checkpoint_path,
saver_pb2.SaverDef.V2)
v1_path = _prefix_to_checkpoint_path(ckpt.model_checkpoint_path,
saver_pb2.SaverDef.V1)
if file_io.get_matching_files(v2_path) or file_io.get_matching_files(
v1_path):
return ckpt.model_checkpoint_path
else:
logging.error("Couldn't match files for checkpoint %s",
ckpt.model_checkpoint_path)
return None
def checkpoint_exists_internal(checkpoint_prefix):
"""Checks whether a V1 or V2 checkpoint exists with the specified prefix.
This is an internal function to check if a checkpoint exists,
since it takes into account the naming difference between V1 and V2 formats.
Args:
checkpoint_prefix: the prefix of a V1 or V2 checkpoint, with V2 taking
priority. Typically the result of `Saver.save()` or that of
`tf.train.latest_checkpoint()`, regardless of sharded/non-sharded or
V1/V2.
Returns:
A bool, true if a checkpoint referred to by `checkpoint_prefix` exists.
"""
pathname = _prefix_to_checkpoint_path(checkpoint_prefix,
saver_pb2.SaverDef.V2)
if file_io.get_matching_files(pathname):
return True
elif file_io.get_matching_files(checkpoint_prefix):
return True
else:
return False
@deprecation.deprecated(
date=None,
instructions="Use standard file APIs to check for files with this prefix.")
@tf_export(v1=["train.checkpoint_exists"])
def checkpoint_exists(checkpoint_prefix):
"""Checks whether a V1 or V2 checkpoint exists with the specified prefix.
This is the recommended way to check if a checkpoint exists, since it takes
into account the naming difference between V1 and V2 formats.
Args:
checkpoint_prefix: the prefix of a V1 or V2 checkpoint, with V2 taking
priority. Typically the result of `Saver.save()` or that of
`tf.train.latest_checkpoint()`, regardless of sharded/non-sharded or
V1/V2.
Returns:
A bool, true if a checkpoint referred to by `checkpoint_prefix` exists.
"""
return checkpoint_exists_internal(checkpoint_prefix)
@deprecation.deprecated(
date=None,
instructions="Use standard file utilities to get mtimes.")
@tf_export(v1=["train.get_checkpoint_mtimes"])
def get_checkpoint_mtimes(checkpoint_prefixes):
"""Returns the mtimes (modification timestamps) of the checkpoints.
Globs for the checkpoints pointed to by `checkpoint_prefixes`. If the files
exist, collect their mtime. Both V2 and V1 checkpoints are considered, in
that priority.
This is the recommended way to get the mtimes, since it takes into account
the naming difference between V1 and V2 formats.
Note: If not all checkpoints exist, the length of the returned mtimes list
will be smaller than the length of `checkpoint_prefixes` list, so mapping
checkpoints to corresponding mtimes will not be possible.
Args:
checkpoint_prefixes: a list of checkpoint paths, typically the results of
`Saver.save()` or those of `tf.train.latest_checkpoint()`, regardless of
sharded/non-sharded or V1/V2.
Returns:
A list of mtimes (in microseconds) of the found checkpoints.
"""
mtimes = []
def match_maybe_append(pathname):
fnames = file_io.get_matching_files(pathname)
if fnames:
mtimes.append(file_io.stat(fnames[0]).mtime_nsec / 1e9)
return True
return False
for checkpoint_prefix in checkpoint_prefixes:
# Tries V2's metadata file first.
pathname = _prefix_to_checkpoint_path(checkpoint_prefix,
saver_pb2.SaverDef.V2)
if match_maybe_append(pathname):
continue
# Otherwise, tries V1, where the prefix is the complete pathname.
match_maybe_append(checkpoint_prefix)
return mtimes
@deprecation.deprecated(
date=None,
instructions="Use standard file APIs to delete files with this prefix.")
@tf_export(v1=["train.remove_checkpoint"])
def remove_checkpoint(checkpoint_prefix,
checkpoint_format_version=saver_pb2.SaverDef.V2,
meta_graph_suffix="meta"):
"""Removes a checkpoint given by `checkpoint_prefix`.
Args:
checkpoint_prefix: The prefix of a V1 or V2 checkpoint. Typically the result
of `Saver.save()` or that of `tf.train.latest_checkpoint()`, regardless of
sharded/non-sharded or V1/V2.
checkpoint_format_version: `SaverDef.CheckpointFormatVersion`, defaults to
`SaverDef.V2`.
meta_graph_suffix: Suffix for `MetaGraphDef` file. Defaults to 'meta'.
"""
_delete_file_if_exists(
meta_graph_filename(checkpoint_prefix, meta_graph_suffix))
if checkpoint_format_version == saver_pb2.SaverDef.V2:
# V2 has a metadata file and some data files.
_delete_file_if_exists(checkpoint_prefix + ".index")
_delete_file_if_exists(checkpoint_prefix + ".data-?????-of-?????")
else:
# V1, Legacy. Exact match on the data file.
_delete_file_if_exists(checkpoint_prefix)
def _delete_file_if_exists(filespec):
"""Deletes files matching `filespec`."""
for pathname in file_io.get_matching_files(filespec):
try:
file_io.delete_file(pathname)
except errors.NotFoundError:
logging.warning(
"Hit NotFoundError when deleting '%s', possibly because another "
"process/thread is also deleting/moving the same file", pathname)
def meta_graph_filename(checkpoint_filename, meta_graph_suffix="meta"):
"""Returns the meta graph filename.
Args:
checkpoint_filename: Name of the checkpoint file.
meta_graph_suffix: Suffix for `MetaGraphDef` file. Defaults to 'meta'.
Returns:
MetaGraph file name.
"""
# If the checkpoint_filename is sharded, the checkpoint_filename could
# be of format model.ckpt-step#-?????-of-shard#. For example,
# model.ckpt-123456-?????-of-00005, or model.ckpt-123456-00001-of-00002.
basename = re.sub(r"-[\d\?]+-of-\d+$", "", checkpoint_filename)
suffixed_filename = ".".join([basename, meta_graph_suffix])
return suffixed_filename
# TODO(allenl): Allow tf.keras.Model instances in the constructor directly?
@tf_export("train.CheckpointManager")
class CheckpointManager(object):
"""Manages multiple checkpoints by keeping some and deleting unneeded ones.
Example usage:
```python
import tensorflow as tf
checkpoint = tf.train.Checkpoint(optimizer=optimizer, model=model)
manager = tf.train.CheckpointManager(
checkpoint, directory="/tmp/model", max_to_keep=5)
status = checkpoint.restore(manager.latest_checkpoint)
while True:
# train
manager.save()
```
`CheckpointManager` preserves its own state across instantiations (see the
`__init__` documentation for details). Only one should be active in a
particular directory at a time.
"""
def __init__(
self,
checkpoint,
directory,
max_to_keep,
keep_checkpoint_every_n_hours=None,
checkpoint_name="ckpt",
step_counter=None,
checkpoint_interval=None,
init_fn=None,
last_checkpoint_step=None,
):
"""Configure a `CheckpointManager` for use in `directory`.
If a `CheckpointManager` was previously used in `directory`, its
state will be restored. This includes the list of managed checkpoints and
the timestamp bookkeeping necessary to support
`keep_checkpoint_every_n_hours`. The behavior of the new `CheckpointManager`
will be the same as the previous `CheckpointManager`, including cleaning up
existing checkpoints if appropriate.
Checkpoints are only considered for deletion just after a new checkpoint has
been added. At that point, `max_to_keep` checkpoints will remain in an
"active set". Once a checkpoint is preserved by
`keep_checkpoint_every_n_hours` it will not be deleted by this
`CheckpointManager` or any future `CheckpointManager` instantiated in
`directory` (regardless of the new setting of
`keep_checkpoint_every_n_hours`). The `max_to_keep` checkpoints in the
active set may be deleted by this `CheckpointManager` or a future
`CheckpointManager` instantiated in `directory` (subject to its
`max_to_keep` and `keep_checkpoint_every_n_hours` settings).
`CheckpointManager` can be also used for initializing the model if
there is no checkpoints for restoring in `directory`. An example usage is:
>>> import tempfile
>>> tmp_dir = tempfile.mkdtemp()
>>> checkpoint = tf.train.Checkpoint()
>>> init_path = checkpoint.save(os.path.join(tmp_dir, 'init'))
>>> def init_fn():
... # Partially restore the checkpoint from `init_path`.
... checkpoint.restore(init_path)
>>> manager = tf.train.CheckpointManager(
... checkpoint,
... directory=os.path.join(tmp_dir, 'ckpt'),
... max_to_keep=None,
... init_fn=init_fn)
>>> # `restore_or_initialize` will call `init_fn` if there is no existing
>>> # checkpoint in `directory`.
>>> manager.restore_or_initialize()
Args:
checkpoint: The `tf.train.Checkpoint` instance to save and manage
checkpoints for.
directory: The path to a directory in which to write checkpoints. A
special file named "checkpoint" is also written to this directory (in a
human-readable text format) which contains the state of the
`CheckpointManager`.
max_to_keep: An integer, the number of checkpoints to keep. Unless
preserved by `keep_checkpoint_every_n_hours`, checkpoints will be
deleted from the active set, oldest first, until only `max_to_keep`
checkpoints remain. If `None`, no checkpoints are deleted and everything
stays in the active set. Note that `max_to_keep=None` will keep all
checkpoint paths in memory and in the checkpoint state protocol buffer
on disk.
keep_checkpoint_every_n_hours: Upon removal from the active set, a
checkpoint will be preserved if it has been at least
`keep_checkpoint_every_n_hours` since the last preserved checkpoint. The
default setting of `None` does not preserve any checkpoints in this way.
checkpoint_name: Custom name for the checkpoint file.
step_counter: A `tf.Variable` instance for checking the current step
counter value, in case users want to save checkpoints every N steps. It
should be passed if `checkpoint_interval` is not None.
checkpoint_interval: An integer, indicates the minimum step interval
between two checkpoints.
init_fn: Callable. A function to do customized initialization if no
checkpoints are in the directory.
last_checkpoint_step: An integer, indicating the step number of the last
checkpoint saved. This will be used as the starting point for checking
checkpoint_interval against the current step. If None, the last
checkpoint step will be set to None.
Raises:
ValueError: If `max_to_keep` is not a positive integer.
"""
self._checkpoint = checkpoint
self._save_counter_assign = None
if max_to_keep is not None and max_to_keep <= 0:
raise ValueError(
("Expected a positive integer or `None` for `max_to_keep`, "
"got %d.")
% (max_to_keep,))
self._max_to_keep = max_to_keep
self._keep_checkpoint_every_n_hours = keep_checkpoint_every_n_hours
if isinstance(directory, os.PathLike):
directory = os.fspath(directory)
self._directory = directory
self._checkpoint_prefix = os.path.join(directory, checkpoint_name)
self._init_fn = init_fn
if checkpoint_interval is not None:
if step_counter is None:
raise ValueError("`step_counter` should be passed if "
"`checkpoint_interval` is not None.")
self._last_checkpoint_step = last_checkpoint_step
self._step_counter = step_counter
self._checkpoint_interval = checkpoint_interval
recovered_state = get_checkpoint_state(directory)
current_clock = time.time()
self._maybe_delete = collections.OrderedDict()
if recovered_state is None:
self._latest_checkpoint = None
# Set the clock back slightly to avoid race conditions when quickly
# re-creating a CheckpointManager.
self._last_preserved_timestamp = current_clock - 1.
else:
self._latest_checkpoint = recovered_state.model_checkpoint_path
self._last_preserved_timestamp = recovered_state.last_preserved_timestamp
if current_clock < self._last_preserved_timestamp:
# Time seems to have reversed itself. In addition to this warning, we'll
# min() saved checkpoint timestamps with the current time to ensure that
# old checkpoints don't get deleted accidentally.
logging.warning(
("time.time() returned a value %f seconds behind the last "
"preserved checkpoint timestamp.")
% (self._last_preserved_timestamp - current_clock,))
self._last_preserved_timestamp = current_clock
all_timestamps = recovered_state.all_model_checkpoint_timestamps
all_paths = recovered_state.all_model_checkpoint_paths
del recovered_state # Uses modified values from now on
if not all_timestamps:
all_timestamps = [self._last_preserved_timestamp] * len(all_paths)
for filename, timestamp in zip(all_paths, all_timestamps):
timestamp = min(timestamp, current_clock)
if timestamp > self._last_preserved_timestamp:
self._maybe_delete[filename] = timestamp
@property
def directory(self):
return self._directory
@property
def checkpoint_interval(self):
return self._checkpoint_interval
@property
def latest_checkpoint(self):
"""The prefix of the most recent checkpoint in `directory`.
Equivalent to `tf.train.latest_checkpoint(directory)` where `directory` is
the constructor argument to `CheckpointManager`.
Suitable for passing to `tf.train.Checkpoint.restore` to resume training.
Returns:
The checkpoint prefix. If there are no checkpoints, returns `None`.
"""
return self._latest_checkpoint
@property
def checkpoints(self):
"""A list of managed checkpoints.
Note that checkpoints saved due to `keep_checkpoint_every_n_hours` will not
show up in this list (to avoid ever-growing filename lists).
Returns:
A list of filenames, sorted from oldest to newest.
"""
return list(self._maybe_delete.keys())
def _sweep(self):
"""Deletes or preserves managed checkpoints."""
if not self._max_to_keep:
# Does not update self._last_preserved_timestamp, since everything is kept
# in the active set.
return
while len(self._maybe_delete) > self._max_to_keep:
filename, timestamp = self._maybe_delete.popitem(last=False)
# Even if we're keeping this checkpoint due to
# keep_checkpoint_every_n_hours, we won't reference it to avoid
# infinitely-growing CheckpointState protos.
if (self._keep_checkpoint_every_n_hours
and (timestamp - self._keep_checkpoint_every_n_hours * 3600.
>= self._last_preserved_timestamp)):
self._last_preserved_timestamp = timestamp
continue
_delete_file_if_exists(filename + ".index")
_delete_file_if_exists(filename + ".data-?????-of-?????")
def _record_state(self):
"""Saves the `CheckpointManager`'s state in `directory`."""
filenames, timestamps = zip(*self._maybe_delete.items())
update_checkpoint_state_internal(
self._directory,
model_checkpoint_path=self.latest_checkpoint,
all_model_checkpoint_paths=filenames,
all_model_checkpoint_timestamps=timestamps,
last_preserved_timestamp=self._last_preserved_timestamp,
save_relative_paths=True)
@property
def _prefix(self):
"""A common prefix for all checkpoints saved with this manager.
For example, if `directory` (a constructor argument) were `"/tmp/tf-model"`,
`prefix` would be `"/tmp/tf-model/ckpt"` and checkpoints would generally be
numbered `"/tmp/tf-model/ckpt-1"`, `"/tmp/tf-model/ckpt-2"`, and so on. Each
checkpoint has several associated files
(e.g. `"/tmp/tf-model/ckpt-2.index"`).
Returns:
A string prefix.
"""
return self._checkpoint_prefix
@property
def checkpoint(self):
"""Returns the `tf.train.Checkpoint` object."""
return self._checkpoint
def save(self, checkpoint_number=None, check_interval=True, options=None):
"""Creates a new checkpoint and manages it.
Args:
checkpoint_number: An optional integer, or an integer-dtype `Variable` or
`Tensor`, used to number the checkpoint. If `None` (default),
checkpoints are numbered using `checkpoint.save_counter`. Even if
`checkpoint_number` is provided, `save_counter` is still incremented. A
user-provided `checkpoint_number` is not incremented even if it is a
`Variable`.
check_interval: An optional boolean. The argument is only effective when
`checkpoint_interval` is passed into the manager. If `True`, the manager
will only save the checkpoint if the interval between checkpoints is
larger than `checkpoint_interval`. Otherwise it will always save the
checkpoint unless a checkpoint has already been saved for the current
step.
options: Optional `tf.train.CheckpointOptions` object. This argument only
works with TF2 checkpoint objects. For example, options =
tf.saved_model.SaveOptions(experimental_io_device='/job:localhost')
Returns:
The path to the new checkpoint. It is also recorded in the `checkpoints`
and `latest_checkpoint` properties. `None` if no checkpoint is saved.
"""
if self._checkpoint_interval is not None:
current_step = _evaluate(self._step_counter)
if self._last_checkpoint_step is not None:
if current_step == self._last_checkpoint_step:
return None
if check_interval and current_step < (
self._last_checkpoint_step + self._checkpoint_interval):
return None
self._last_checkpoint_step = current_step
# Save counter logic duplicated from tf.train.Checkpoint, soon to diverge
# slightly with a custom numbering option.
if context.executing_eagerly():
save_counter = self._checkpoint.save_counter
save_counter.assign_add(1)
session = None
else:
session = ops.get_default_session()
def _initializing_creator(next_creator, **kwargs):
"""Initialize the save counter if it has been newly created."""
v = next_creator(**kwargs)
session.run(v.initializer)
return v
with variable_scope.variable_creator_scope(_initializing_creator):
save_counter = self._checkpoint.save_counter
if self._save_counter_assign is None:
self._save_counter_assign = save_counter.assign_add(1, read_value=False)
session.run(self._save_counter_assign)
if checkpoint_number is None:
checkpoint_number = save_counter
if not isinstance(checkpoint_number, compat.integral_types):
checkpoint_number = training_util.global_step(
sess=session, global_step_tensor=checkpoint_number)
prefix = "%s-%d" % (self._prefix, checkpoint_number)
def _record_and_sweep_state(save_path):
timestamp = time.time()
# If this is an overwritten checkpoint we were previously tracking, delete
# and reinsert it to make sure it goes to the end of the queue.
if save_path in self._maybe_delete:
del self._maybe_delete[save_path]
self._maybe_delete[save_path] = timestamp
self._latest_checkpoint = save_path
# Before deleting anything we update the Checkpoint proto with the new
# checkpoint. We'll go back and correct it after cleaning up old files,
# but a preemption while deleting will be more likely to see the new
# checkpoint this way.
self._record_state()
self._sweep()
# Write out the Checkpoint proto a second time, now without the deleted
# checkpoints.
self._record_state()
# Register `_record_and_sweep_state` as a callback in `CheckpointOptions`
if options is None:
options = checkpoint_options.CheckpointOptions(
experimental_write_callbacks=[_record_and_sweep_state]
)
else:
# We create a copy so that user's `options` instance would not be mutated
# by internal mechanisms.
options = copy.copy(options)
if options.experimental_write_callbacks is None:
options.experimental_write_callbacks = [_record_and_sweep_state]
else:
options.experimental_write_callbacks.append(_record_and_sweep_state)
save_path = self._checkpoint._write(prefix, options=options) # pylint: disable=protected-access
return save_path
def restore_or_initialize(self):
"""Restore items in `checkpoint` from the latest checkpoint file.
This method will first try to restore from the most recent checkpoint in
`directory`. If no checkpoints exist in `directory`, and `init_fn` is
specified, this method will call `init_fn` to do customized
initialization. This can be used to support initialization from pretrained
models.
Note that unlike `tf.train.Checkpoint.restore()`, this method doesn't return
a load status object that users can run assertions on
(e.g. assert_consumed()). Thus to run assertions, users should directly use
`tf.train.Checkpoint.restore()` method.
Returns:
The restored checkpoint path if the latest checkpoint is found and
restored. Otherwise None.
"""
# TODO(chienchunh): When AsyncCheckpoint is used, we may need to force to
# sync until any ongoing async save is done. Otherwise, if this is the first
# checkpoint and _latest_checkpoint has not been updated due to async write,
# this would resort to init_fn instead of restoring from the checkpoin file.
# This should be fixed once AsyncCheckpoint is integrated with the public
# API so that we can rely on CheckpointOptions to tell whether we should
# sync for AsyncCheckpoint.
if self._latest_checkpoint is not None:
self._checkpoint.restore(self._latest_checkpoint)
if self._checkpoint_interval is not None:
self._last_checkpoint_step = _evaluate(self._step_counter)
return self._latest_checkpoint
if self._init_fn is not None:
self._init_fn()
logging.info(
"Customized initialization is done through the passed `init_fn`.")
return None
def sync(self):
"""Wait for any outstanding save or restore operations."""
if self._checkpoint:
self._checkpoint.sync()
tensorflow/python/checkpoint/checkpoint_management_test.py
+736
View File
@@ -0,0 +1,736 @@
# Copyright 2015 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 tensorflow.python.training.saver.py."""
import contextlib
import os
import pathlib
import shutil
import tempfile
from google.protobuf import text_format
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.checkpoint import checkpoint_management
from tensorflow.python.eager import context
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops as ops_lib
from tensorflow.python.framework import test_util
from tensorflow.python.lib.io import file_io
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import saver as saver_module
from tensorflow.python.training.checkpoint_state_pb2 import CheckpointState
class LatestCheckpointWithRelativePaths(test.TestCase):
@staticmethod
@contextlib.contextmanager
def tempWorkingDir(temppath):
cwd = os.getcwd()
os.chdir(temppath)
try:
yield
finally:
os.chdir(cwd)
@staticmethod
@contextlib.contextmanager
def tempDir():
tempdir = tempfile.mkdtemp()
try:
yield tempdir
finally:
shutil.rmtree(tempdir)
@test_util.run_deprecated_v1
def testNameCollision(self):
# Make sure we have a clean directory to work in.
with self.tempDir() as tempdir:
# Jump to that directory until this test is done.
with self.tempWorkingDir(tempdir):
# Save training snapshots to a relative path.
traindir = "train"
os.mkdir(traindir)
# Collides with the default name of the checkpoint state file.
filepath = os.path.join(traindir, "checkpoint")
with self.cached_session() as sess:
unused_a = variables.Variable(0.0) # So that Saver saves something.
self.evaluate(variables.global_variables_initializer())
# Should fail.
saver = saver_module.Saver(sharded=False)
with self.assertRaisesRegex(ValueError, "collides with"):
saver.save(sess, filepath)
# Succeeds: the file will be named "checkpoint-<step>".
saver.save(sess, filepath, global_step=1)
self.assertIsNotNone(
checkpoint_management.latest_checkpoint(traindir))
# Succeeds: the file will be named "checkpoint-<i>-of-<n>".
saver = saver_module.Saver(sharded=True)
saver.save(sess, filepath)
self.assertIsNotNone(
checkpoint_management.latest_checkpoint(traindir))
# Succeeds: the file will be named "checkpoint-<step>-<i>-of-<n>".
saver = saver_module.Saver(sharded=True)
saver.save(sess, filepath, global_step=1)
self.assertIsNotNone(
checkpoint_management.latest_checkpoint(traindir))
@test_util.run_deprecated_v1
def testRelativePath(self):
# Make sure we have a clean directory to work in.
with self.tempDir() as tempdir:
# Jump to that directory until this test is done.
with self.tempWorkingDir(tempdir):
# Save training snapshots to a relative path.
traindir = "train"
os.mkdir(traindir)
filename = "snapshot"
filepath = os.path.join(traindir, filename)
with self.cached_session() as sess:
# Build a simple graph.
v0 = variables.Variable(0.0)
inc = v0.assign_add(1.0)
save = saver_module.Saver({"v0": v0})
# Record a short training history.
self.evaluate(variables.global_variables_initializer())
save.save(sess, filepath, global_step=0)
self.evaluate(inc)
save.save(sess, filepath, global_step=1)
self.evaluate(inc)
save.save(sess, filepath, global_step=2)
with self.cached_session() as sess:
# Build a new graph with different initialization.
v0 = variables.Variable(-1.0)
# Create a new saver.
save = saver_module.Saver({"v0": v0})
self.evaluate(variables.global_variables_initializer())
# Get the most recent checkpoint name from the training history file.
name = checkpoint_management.latest_checkpoint(traindir)
self.assertIsNotNone(name)
# Restore "v0" from that checkpoint.
save.restore(sess, name)
self.assertEqual(self.evaluate(v0), 2.0)
class CheckpointStateTest(test.TestCase):
def _get_test_dir(self, dirname):
test_dir = os.path.join(self.get_temp_dir(), dirname)
gfile.MakeDirs(test_dir)
return test_dir
def testAbsPath(self):
save_dir = self._get_test_dir("abs_paths")
abs_path = os.path.join(save_dir, "model-0")
ckpt = checkpoint_management.generate_checkpoint_state_proto(
save_dir, abs_path)
self.assertEqual(ckpt.model_checkpoint_path, abs_path)
self.assertTrue(os.path.isabs(ckpt.model_checkpoint_path))
self.assertEqual(len(ckpt.all_model_checkpoint_paths), 1)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], abs_path)
def testRelPath(self):
train_dir = "train"
model = os.path.join(train_dir, "model-0")
# model_checkpoint_path should have no "train" directory part.
new_rel_path = "model-0"
ckpt = checkpoint_management.generate_checkpoint_state_proto(
train_dir, model)
self.assertEqual(ckpt.model_checkpoint_path, new_rel_path)
self.assertEqual(len(ckpt.all_model_checkpoint_paths), 1)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], new_rel_path)
def testAllModelCheckpointPaths(self):
save_dir = self._get_test_dir("all_models_test")
abs_path = os.path.join(save_dir, "model-0")
for paths in [None, [], ["model-2"]]:
ckpt = checkpoint_management.generate_checkpoint_state_proto(
save_dir, abs_path, all_model_checkpoint_paths=paths)
self.assertEqual(ckpt.model_checkpoint_path, abs_path)
self.assertTrue(os.path.isabs(ckpt.model_checkpoint_path))
self.assertEqual(
len(ckpt.all_model_checkpoint_paths), len(paths) if paths else 1)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], abs_path)
def testUpdateCheckpointState(self):
save_dir = self._get_test_dir("update_checkpoint_state")
os.chdir(save_dir)
# Make a temporary train directory.
train_dir = "train"
os.mkdir(train_dir)
abs_path = os.path.join(save_dir, "model-0")
rel_path = os.path.join("train", "model-2")
checkpoint_management.update_checkpoint_state(
train_dir, rel_path, all_model_checkpoint_paths=[abs_path, rel_path])
ckpt = checkpoint_management.get_checkpoint_state(train_dir)
self.assertEqual(ckpt.model_checkpoint_path, rel_path)
self.assertEqual(len(ckpt.all_model_checkpoint_paths), 2)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], rel_path)
self.assertEqual(ckpt.all_model_checkpoint_paths[0], abs_path)
def testFSPath(self):
save_dir = self._get_test_dir("fspath")
os.chdir(save_dir)
# Make a temporary train directory.
train_dir = "train"
os.mkdir(train_dir)
abs_path = os.path.join(save_dir, "model-0")
rel_path = os.path.join("train", "model-2")
checkpoint_management.update_checkpoint_state(
train_dir, rel_path, all_model_checkpoint_paths=[abs_path, rel_path])
ckpt = checkpoint_management.get_checkpoint_state(pathlib.Path(train_dir))
self.assertEqual(ckpt.model_checkpoint_path, rel_path)
def testUpdateCheckpointStateSaveRelativePaths(self):
save_dir = self._get_test_dir("update_checkpoint_state")
os.chdir(save_dir)
abs_path2 = os.path.join(save_dir, "model-2")
rel_path2 = "model-2"
abs_path0 = os.path.join(save_dir, "model-0")
rel_path0 = "model-0"
checkpoint_management.update_checkpoint_state_internal(
save_dir=save_dir,
model_checkpoint_path=abs_path2,
all_model_checkpoint_paths=[rel_path0, abs_path2],
save_relative_paths=True)
# File should contain relative paths.
file_content = file_io.read_file_to_string(
os.path.join(save_dir, "checkpoint"))
ckpt = CheckpointState()
text_format.Merge(file_content, ckpt)
self.assertEqual(ckpt.model_checkpoint_path, rel_path2)
self.assertEqual(len(ckpt.all_model_checkpoint_paths), 2)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], rel_path2)
self.assertEqual(ckpt.all_model_checkpoint_paths[0], rel_path0)
# get_checkpoint_state should return absolute paths.
ckpt = checkpoint_management.get_checkpoint_state(save_dir)
self.assertEqual(ckpt.model_checkpoint_path, abs_path2)
self.assertEqual(len(ckpt.all_model_checkpoint_paths), 2)
self.assertEqual(ckpt.all_model_checkpoint_paths[-1], abs_path2)
self.assertEqual(ckpt.all_model_checkpoint_paths[0], abs_path0)
def testCheckPointStateFailsWhenIncomplete(self):
save_dir = self._get_test_dir("checkpoint_state_fails_when_incomplete")
os.chdir(save_dir)
ckpt_path = os.path.join(save_dir, "checkpoint")
ckpt_file = open(ckpt_path, "w")
ckpt_file.write("")
ckpt_file.close()
with self.assertRaises(ValueError):
checkpoint_management.get_checkpoint_state(save_dir)
def testCheckPointCompletesRelativePaths(self):
save_dir = self._get_test_dir("checkpoint_completes_relative_paths")
os.chdir(save_dir)
ckpt_path = os.path.join(save_dir, "checkpoint")
ckpt_file = open(ckpt_path, "w")
ckpt_file.write("""
model_checkpoint_path: "./model.ckpt-687529"
all_model_checkpoint_paths: "./model.ckpt-687500"
all_model_checkpoint_paths: "./model.ckpt-687529"
""")
ckpt_file.close()
ckpt = checkpoint_management.get_checkpoint_state(save_dir)
self.assertEqual(ckpt.model_checkpoint_path,
os.path.join(save_dir, "./model.ckpt-687529"))
self.assertEqual(ckpt.all_model_checkpoint_paths[0],
os.path.join(save_dir, "./model.ckpt-687500"))
self.assertEqual(ckpt.all_model_checkpoint_paths[1],
os.path.join(save_dir, "./model.ckpt-687529"))
class SaverUtilsTest(test.TestCase):
def setUp(self):
self._base_dir = os.path.join(self.get_temp_dir(), "saver_utils_test")
gfile.MakeDirs(self._base_dir)
def tearDown(self):
gfile.DeleteRecursively(self._base_dir)
@test_util.run_deprecated_v1
def testCheckpointExists(self):
for sharded in (False, True):
for version in (saver_pb2.SaverDef.V2, saver_pb2.SaverDef.V1):
with self.session(graph=ops_lib.Graph()) as sess:
unused_v = variables.Variable(1.0, name="v")
self.evaluate(variables.global_variables_initializer())
saver = saver_module.Saver(sharded=sharded, write_version=version)
path = os.path.join(self._base_dir, "%s-%s" % (sharded, version))
self.assertFalse(
checkpoint_management.checkpoint_exists(path)) # Not saved yet.
ckpt_prefix = saver.save(sess, path)
self.assertTrue(checkpoint_management.checkpoint_exists(ckpt_prefix))
ckpt_prefix = checkpoint_management.latest_checkpoint(self._base_dir)
self.assertTrue(checkpoint_management.checkpoint_exists(ckpt_prefix))
@test_util.run_deprecated_v1
def testGetCheckpointMtimes(self):
prefixes = []
for version in (saver_pb2.SaverDef.V2, saver_pb2.SaverDef.V1):
with self.session(graph=ops_lib.Graph()) as sess:
unused_v = variables.Variable(1.0, name="v")
self.evaluate(variables.global_variables_initializer())
saver = saver_module.Saver(write_version=version)
prefixes.append(
saver.save(sess, os.path.join(self._base_dir, str(version))))
mtimes = checkpoint_management.get_checkpoint_mtimes(prefixes)
self.assertEqual(2, len(mtimes))
self.assertTrue(mtimes[1] >= mtimes[0])
@test_util.run_deprecated_v1
def testRemoveCheckpoint(self):
for sharded in (False, True):
for version in (saver_pb2.SaverDef.V2, saver_pb2.SaverDef.V1):
with self.session(graph=ops_lib.Graph()) as sess:
unused_v = variables.Variable(1.0, name="v")
self.evaluate(variables.global_variables_initializer())
saver = saver_module.Saver(sharded=sharded, write_version=version)
path = os.path.join(self._base_dir, "%s-%s" % (sharded, version))
ckpt_prefix = saver.save(sess, path)
self.assertTrue(checkpoint_management.checkpoint_exists(ckpt_prefix))
checkpoint_management.remove_checkpoint(ckpt_prefix, version)
self.assertFalse(checkpoint_management.checkpoint_exists(ckpt_prefix))
class CheckpointManagerTest(test.TestCase):
@test_util.run_in_graph_and_eager_modes
def testDeletion(self):
checkpoint = util.Checkpoint()
manager = checkpoint_management.CheckpointManager(
checkpoint, self.get_temp_dir(), max_to_keep=3)
first_path = manager.save()
second_path = manager.save()
third_path = manager.save()
fourth_path = manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(fourth_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertFalse(checkpoint_management.checkpoint_exists(first_path))
@test_util.run_in_graph_and_eager_modes
def testKeepAll(self):
checkpoint = util.Checkpoint()
directory = os.path.join(
self.get_temp_dir(),
# Avoid sharing directories between eager and graph
# TODO(allenl): stop run_in_graph_and_eager_modes reusing directories
str(context.executing_eagerly()))
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=None)
first_path = manager.save()
second_path = manager.save()
third_path = manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertTrue(checkpoint_management.checkpoint_exists(first_path))
self.assertEqual(third_path, manager.latest_checkpoint)
self.assertEqual([first_path, second_path, third_path],
manager.checkpoints)
del manager
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=None)
fourth_path = manager.save()
self.assertEqual([first_path, second_path, third_path, fourth_path],
manager.checkpoints)
del manager
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=3)
self.assertEqual([first_path, second_path, third_path, fourth_path],
manager.checkpoints)
self.assertTrue(checkpoint_management.checkpoint_exists(fourth_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertTrue(checkpoint_management.checkpoint_exists(first_path))
fifth_path = manager.save()
self.assertEqual([third_path, fourth_path, fifth_path],
manager.checkpoints)
self.assertTrue(checkpoint_management.checkpoint_exists(fifth_path))
self.assertTrue(checkpoint_management.checkpoint_exists(fourth_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertFalse(checkpoint_management.checkpoint_exists(second_path))
self.assertFalse(checkpoint_management.checkpoint_exists(first_path))
@test_util.run_in_graph_and_eager_modes
@test.mock.patch.object(checkpoint_management, "time")
def testSaveRestoreState(self, mock_time):
directory = self.get_temp_dir()
mock_time.time.return_value = 3.
checkpoint = util.Checkpoint()
first_manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2)
first_time = 10000.
first_name = os.path.join(directory, "ckpt-1")
mock_time.time.return_value = first_time
first_manager.save()
state = checkpoint_management.get_checkpoint_state(directory)
second_time = first_time + 3610.
second_name = os.path.join(directory, "ckpt-2")
mock_time.time.return_value = second_time
first_manager.save()
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual([first_time, second_time],
state.all_model_checkpoint_timestamps)
self.assertEqual([first_name, second_name], first_manager.checkpoints)
self.assertEqual(second_name, first_manager.latest_checkpoint)
del first_manager
second_manager = checkpoint_management.CheckpointManager(
checkpoint, directory,
max_to_keep=2, keep_checkpoint_every_n_hours=1.5)
self.assertEqual([first_name, second_name], second_manager.checkpoints)
self.assertEqual(second_name, second_manager.latest_checkpoint)
third_name = os.path.join(directory, "ckpt-3")
third_time = second_time + 3600. * 0.2
mock_time.time.return_value = third_time
second_manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(first_name))
self.assertTrue(checkpoint_management.checkpoint_exists(second_name))
self.assertEqual([second_name, third_name],
second_manager.checkpoints)
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual(first_time, state.last_preserved_timestamp)
fourth_time = third_time + 3600. * 0.5
mock_time.time.return_value = fourth_time
fourth_name = os.path.join(directory, "ckpt-4")
second_manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(first_name))
self.assertFalse(checkpoint_management.checkpoint_exists(second_name))
self.assertEqual([third_name, fourth_name],
second_manager.checkpoints)
fifth_time = fourth_time + 3600. * 0.5
mock_time.time.return_value = fifth_time
fifth_name = os.path.join(directory, "ckpt-5")
second_manager.save()
self.assertEqual([fourth_name, fifth_name],
second_manager.checkpoints)
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual(first_time, state.last_preserved_timestamp)
del second_manager
third_manager = checkpoint_management.CheckpointManager(
checkpoint, directory,
max_to_keep=2, keep_checkpoint_every_n_hours=1.5)
self.assertEqual(fifth_name, third_manager.latest_checkpoint)
mock_time.time.return_value += 10.
third_manager.save()
sixth_name = os.path.join(directory, "ckpt-6")
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual(fourth_time, state.last_preserved_timestamp)
self.assertTrue(checkpoint_management.checkpoint_exists(first_name))
self.assertTrue(checkpoint_management.checkpoint_exists(fourth_name))
self.assertTrue(checkpoint_management.checkpoint_exists(fifth_name))
self.assertTrue(checkpoint_management.checkpoint_exists(sixth_name))
self.assertFalse(checkpoint_management.checkpoint_exists(second_name))
self.assertFalse(checkpoint_management.checkpoint_exists(third_name))
self.assertEqual([fifth_name, sixth_name],
third_manager.checkpoints)
@test_util.run_in_graph_and_eager_modes
def testContinueFromUnmanaged(self):
directory = self.get_temp_dir()
prefix = os.path.join(directory, "unusual_prefix")
checkpoint = util.Checkpoint()
first_path = checkpoint.save(prefix)
second_path = checkpoint.save(prefix)
del checkpoint
checkpoint = util.Checkpoint()
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2)
checkpoint.restore(manager.latest_checkpoint).run_restore_ops()
self.assertEqual(2, self.evaluate(checkpoint.save_counter))
third_path = manager.save()
self.assertEqual([third_path], manager.checkpoints)
fourth_path = manager.save()
self.assertEqual([third_path, fourth_path],
manager.checkpoints)
fifth_path = manager.save()
self.assertEqual([fourth_path, fifth_path],
manager.checkpoints)
self.assertTrue(checkpoint_management.checkpoint_exists(first_path))
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertFalse(checkpoint_management.checkpoint_exists(third_path))
self.assertTrue(checkpoint_management.checkpoint_exists(fourth_path))
self.assertTrue(checkpoint_management.checkpoint_exists(fifth_path))
@test_util.run_in_graph_and_eager_modes
@test.mock.patch.object(checkpoint_management, "time")
def testClockReset(self, mock_time):
directory = self.get_temp_dir()
mock_time.time.return_value = 10000.
checkpoint = util.Checkpoint()
first_manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=1, keep_checkpoint_every_n_hours=1.)
first_path = first_manager.save()
mock_time.time.return_value += 3600.
second_path = first_manager.save()
mock_time.time.return_value += 3600.
third_path = first_manager.save()
self.assertFalse(checkpoint_management.checkpoint_exists(first_path))
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertEqual([third_path], first_manager.checkpoints)
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual(13600., state.last_preserved_timestamp)
# Set the clock back in time
mock_time.time.return_value = 5000.
del first_manager
with test.mock.patch.object(logging, "warning") as mock_log:
second_manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=1)
self.assertRegex(
str(mock_log.call_args),
"behind the last preserved checkpoint timestamp")
# We should err on the side of keeping checkpoints around when we're not
# sure whether they were preserved or not due to clock funkiness.
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
# We know about the existing checkpoints, but they'll never be deleted and
# so won't go in the CheckpointState proto on save.
self.assertEqual(third_path, second_manager.latest_checkpoint)
self.assertEqual([], second_manager.checkpoints)
mock_time.time.return_value += 10.
fourth_path = second_manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertEqual(fourth_path, second_manager.latest_checkpoint)
self.assertEqual([fourth_path], second_manager.checkpoints)
mock_time.time.return_value += 10.
fifth_path = second_manager.save()
self.assertTrue(checkpoint_management.checkpoint_exists(second_path))
self.assertTrue(checkpoint_management.checkpoint_exists(third_path))
self.assertEqual([fifth_path], second_manager.checkpoints)
state = checkpoint_management.get_checkpoint_state(directory)
self.assertEqual(5000., state.last_preserved_timestamp)
self.assertEqual([5020.],
state.all_model_checkpoint_timestamps)
@test_util.run_in_graph_and_eager_modes
def testCustomNumbering(self):
directory = self.get_temp_dir()
step = variables.Variable(0, dtype=dtypes.int64)
checkpoint = util.Checkpoint(step=step)
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2)
self.evaluate(step.initializer)
for i in range(5):
path = manager.save(checkpoint_number=step)
expected_suffix = "-%d" % (2 * i,)
if not path.endswith(expected_suffix):
self.fail("%s should have suffix %s" % (path, expected_suffix))
self.evaluate(step.assign_add(2))
self.assertEqual(5, self.evaluate(checkpoint.save_counter))
# Test regular integers
last_path = manager.save(checkpoint_number=32)
self.assertIn("-32", last_path)
self.assertEqual(last_path, manager.latest_checkpoint)
self.assertEqual(
last_path, checkpoint_management.latest_checkpoint(directory))
state = checkpoint_management.get_checkpoint_state(directory)
# Only the most recent two checkpoints are saved
self.assertEqual([path, last_path], state.all_model_checkpoint_paths)
@test_util.run_in_graph_and_eager_modes
def testCustomCheckpointPrefix(self):
directory = self.get_temp_dir()
checkpoint = util.Checkpoint()
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2, checkpoint_name="ckpt_name")
path = manager.save(checkpoint_number=5)
self.assertEqual(os.path.basename(path), "ckpt_name-5")
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2)
path = manager.save(checkpoint_number=5)
self.assertEqual(os.path.basename(path), "ckpt-5")
@test_util.run_in_graph_and_eager_modes
def testRestoreOrInitialize(self):
directory = self.get_temp_dir()
# Create a checkpoint for initializing.
init_prefix = os.path.join(directory, "init")
init_v = variables.Variable(2.0)
init_ckpt = util.Checkpoint(v=init_v)
self.evaluate(init_v.initializer)
init_path = init_ckpt.save(init_prefix)
# Create the checkpoint manager.
ckpt_dir = os.path.join(directory, "ckpt")
v = variables.Variable(1.0)
checkpoint = util.Checkpoint(v=v)
manager = checkpoint_management.CheckpointManager(
checkpoint,
ckpt_dir,
max_to_keep=None,
init_fn=lambda: checkpoint.restore(init_path).run_restore_ops())
self.evaluate(v.initializer)
# First call should call `init_fn`.
self.assertIsNone(manager.restore_or_initialize())
self.assertEqual(2.0, self.evaluate(v))
# Save a checkpoint and second call should restore from the checkpoints.
manager.save()
self.assertIsNotNone(manager.restore_or_initialize())
@test_util.run_in_graph_and_eager_modes
def testCheckpointManagerFSpathDirectory(self):
directory = pathlib.Path(self.get_temp_dir())
v = variables.Variable(0.0)
checkpoint = util.Checkpoint(v=v)
self.evaluate(v.initializer)
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2, checkpoint_name="ckpt_name")
save_path = manager.save()
expected = str(directory / "ckpt_name-1")
self.assertEqual(expected, save_path)
restore_path = manager.restore_or_initialize()
self.assertEqual(str(directory / "ckpt_name-1"), restore_path)
@test_util.run_in_graph_and_eager_modes
def testLatestCheckpointFSpathDirectory(self):
directory = pathlib.Path(self.get_temp_dir())
checkpoint = util.Checkpoint()
manager = checkpoint_management.CheckpointManager(
checkpoint, directory, max_to_keep=2, checkpoint_name="ckpt_name")
manager.save()
cp_dir = checkpoint_management.latest_checkpoint(directory)
self.assertEqual(str(directory / "ckpt_name-1"), cp_dir)
@test_util.run_in_graph_and_eager_modes
def testCheckpointInterval(self):
v = variables.Variable(1.0)
step_counter = variables.Variable(0)
self.evaluate([v.initializer, step_counter.initializer])
checkpoint = util.Checkpoint(v=v)
manager = checkpoint_management.CheckpointManager(
checkpoint,
self.get_temp_dir(),
max_to_keep=None,
step_counter=step_counter,
checkpoint_interval=2)
# step_counter: 0, save an initial checkpoint.
path = manager.save(check_interval=True)
self.assertTrue(checkpoint_management.checkpoint_exists(path))
# step_counter: 1, no checkpoint saved.
self.evaluate(step_counter.assign_add(1))
path = manager.save(check_interval=True)
self.assertIsNone(path)
# step_counter: 2, checkpoint saved.
self.evaluate(step_counter.assign_add(1))
path = manager.save(check_interval=True)
self.assertTrue(checkpoint_management.checkpoint_exists(path))
# no checkpoint saved when calling `save` with the same step counter.
path = manager.save(check_interval=True)
self.assertIsNone(path)
# step_counter: 3, no checkpoint saved.
self.evaluate(step_counter.assign_add(1))
path = manager.save(check_interval=True)
self.assertIsNone(path)
# Always save the checkpoint.
path = manager.save(check_interval=False)
self.assertTrue(checkpoint_management.checkpoint_exists(path))
@test_util.run_in_graph_and_eager_modes
def testCheckpointIntervalWithLastCheckpointStep(self):
v = variables.Variable(1.0)
step_counter = variables.Variable(1)
self.evaluate([v.initializer, step_counter.initializer])
checkpoint = util.Checkpoint(v=v)
manager = checkpoint_management.CheckpointManager(
checkpoint,
self.get_temp_dir(),
max_to_keep=None,
step_counter=step_counter,
checkpoint_interval=2,
last_checkpoint_step=1)
# step_counter: 1, no checkpoint saved.
path = manager.save(check_interval=True)
self.assertIsNone(path)
# step_counter: 2, no checkpoint saved since the interval starts at 1.
self.evaluate(step_counter.assign_add(1))
path = manager.save(check_interval=True)
self.assertIsNone(path)
# step_counter: 3, checkpoint saved.
self.evaluate(step_counter.assign_add(1))
path = manager.save(check_interval=True)
self.assertTrue(checkpoint_management.checkpoint_exists(path))
@test_util.run_in_graph_and_eager_modes
def testCheckpointIntervalWithRestore(self):
directory = self.get_temp_dir()
v = variables.Variable(1.0)
step_counter = variables.Variable(0)
self.evaluate([v.initializer, step_counter.initializer])
# Prepare a checkpoint.
checkpoint = util.Checkpoint(v=v)
checkpoint.save(os.path.join(directory, "ckpt"))
manager = checkpoint_management.CheckpointManager(
checkpoint,
directory,
max_to_keep=None,
step_counter=step_counter,
checkpoint_interval=2)
# Restore from the checkpoint.
self.assertIsNotNone(manager.restore_or_initialize())
# step_counter: 0, no checkpoint saved because it is restored from the
# checkpoint with the same step.
path = manager.save()
self.assertIsNone(path)
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/checkpoint_metrics_test.py
+108
View File
@@ -0,0 +1,108 @@
# Copyright 2021 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 checking the checkpoint reading and writing metrics."""
import os
import time
from tensorflow.core.framework import summary_pb2
from tensorflow.python.checkpoint import checkpoint as util
from tensorflow.python.eager import context
from tensorflow.python.ops import variables as variables_lib
from tensorflow.python.platform import test
from tensorflow.python.saved_model.pywrap_saved_model import metrics
class CheckpointMetricTests(test.TestCase):
def _get_write_histogram_proto(self, api_label):
proto_bytes = metrics.GetCheckpointWriteDurations(api_label=api_label)
histogram_proto = summary_pb2.HistogramProto()
histogram_proto.ParseFromString(proto_bytes)
return histogram_proto
def _get_read_histogram_proto(self, api_label):
proto_bytes = metrics.GetCheckpointReadDurations(api_label=api_label)
histogram_proto = summary_pb2.HistogramProto()
histogram_proto.ParseFromString(proto_bytes)
return histogram_proto
def _get_time_saved(self, api_label):
return metrics.GetTrainingTimeSaved(api_label=api_label)
def _get_checkpoint_size(self, api_label, filesize):
return metrics.GetCheckpointSize(api_label=api_label, filesize=filesize)
def test_metrics_v2(self):
api_label = util._CHECKPOINT_V2
prefix = os.path.join(self.get_temp_dir(), 'ckpt')
with context.eager_mode():
ckpt = util.Checkpoint(v=variables_lib.Variable(1.))
self.assertEqual(self._get_time_saved(api_label), 0.0)
self.assertEqual(self._get_write_histogram_proto(api_label).num, 0.0)
for i in range(3):
time_saved = self._get_time_saved(api_label)
time.sleep(1)
ckpt_path = ckpt.write(file_prefix=prefix)
filesize = util._get_checkpoint_size(ckpt_path)
self.assertEqual(self._get_checkpoint_size(api_label, filesize), i + 1)
self.assertGreater(self._get_time_saved(api_label), time_saved)
self.assertEqual(self._get_write_histogram_proto(api_label).num, 3.0)
self.assertEqual(self._get_read_histogram_proto(api_label).num, 0.0)
time_saved = self._get_time_saved(api_label)
with context.eager_mode():
ckpt.restore(ckpt_path)
self.assertEqual(self._get_read_histogram_proto(api_label).num, 1.0)
# Restoring a checkpoint in the same "job" does not increase training time
# saved.
self.assertEqual(self._get_time_saved(api_label), time_saved)
def test_metrics_v1(self):
api_label = util._CHECKPOINT_V1
prefix = os.path.join(self.get_temp_dir(), 'ckpt')
with self.cached_session():
ckpt = util.CheckpointV1()
v = variables_lib.Variable(1.)
self.evaluate(v.initializer)
ckpt.v = v
self.assertEqual(self._get_time_saved(api_label), 0.0)
self.assertEqual(self._get_write_histogram_proto(api_label).num, 0.0)
for i in range(3):
time_saved = self._get_time_saved(api_label)
time.sleep(1)
ckpt_path = ckpt.write(file_prefix=prefix)
filesize = util._get_checkpoint_size(ckpt_path)
self.assertEqual(self._get_checkpoint_size(api_label, filesize), i + 1)
self.assertGreater(self._get_time_saved(api_label), time_saved)
self.assertEqual(self._get_write_histogram_proto(api_label).num, 3.0)
self.assertEqual(self._get_read_histogram_proto(api_label).num, 0.0)
time_saved = self._get_time_saved(api_label)
ckpt.restore(ckpt_path)
self.assertEqual(self._get_read_histogram_proto(api_label).num, 1.0)
# Restoring a checkpoint in the same "job" does not increase training time
# saved.
self.assertEqual(self._get_time_saved(api_label), time_saved)
if __name__ == '__main__':
test.main()
tensorflow/python/checkpoint/checkpoint_options.py
+132
View File
@@ -0,0 +1,132 @@
# Copyright 2020 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.
# ==============================================================================
"""Options for saving Checkpoints."""
import copy
import inspect
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.util.deprecation import deprecated_args
from tensorflow.python.util.tf_export import tf_export
@tf_export("train.CheckpointOptions")
class CheckpointOptions(object):
"""Options for constructing a Checkpoint.
Used as the `options` argument to either `tf.train.Checkpoint.save()` or
`tf.train.Checkpoint.restore()` methods to adjust how variables are
saved/restored.
Example: Run IO ops on "localhost" while saving a checkpoint:
```
step = tf.Variable(0, name="step")
checkpoint = tf.train.Checkpoint(step=step)
options = tf.train.CheckpointOptions(experimental_io_device="/job:localhost")
checkpoint.save("/tmp/ckpt", options=options)
```
"""
# Define object attributes in __slots__ for improved memory and performance.
__slots__ = (
"experimental_io_device",
"experimental_enable_async_checkpoint",
"experimental_write_callbacks",
"enable_async",
"experimental_sharding_callback",
"experimental_skip_slot_variables",
)
@deprecated_args(
None, "Use enable_async instead", "experimental_enable_async_checkpoint"
)
def __init__(
self,
experimental_io_device=None,
experimental_enable_async_checkpoint=False,
experimental_write_callbacks=None,
enable_async=False,
experimental_skip_slot_variables=False,
experimental_sharding_callback=None
):
"""Creates an object that stores options for a Checkpoint.
Args:
experimental_io_device: string. Applies in a distributed setting.
Tensorflow device to use to access the filesystem. If `None` (default)
then for each variable the filesystem is accessed from the CPU:0 device
of the host where that variable is assigned. If specified, the
filesystem is instead accessed from that device for all variables. This
is for example useful if you want to save to a local directory, such as
"/tmp" when running in a distributed setting. In that case pass a device
for the host where the "/tmp" directory is accessible.
experimental_enable_async_checkpoint: bool Type. Deprecated, please use
the enable_async option.
experimental_write_callbacks: List[Callable]. A list of callback functions
that will be executed after each saving event finishes (i.e. after
`save()` or `write()`). For async checkpoint, the callbacks will be
executed only after the async thread finishes saving. The return values
of the callback(s) will be ignored. The callback(s) can optionally take
the `save_path` (the result of `save()` or `write()`) as an argument.
The callbacks will be executed in the same order of this list after the
checkpoint has been written.
enable_async: bool Type. Indicates whether async checkpointing is enabled.
Default is False, i.e., no async checkpoint. Async checkpoint moves the
checkpoint file writing off the main thread, so that the model can
continue to train while the checkpoint file writing runs in the
background. Async checkpoint reduces TPU device idle cycles and speeds
up model training process, while memory consumption may increase.
experimental_skip_slot_variables: bool Type. If true, ignores slot
variables during restore. Context: TPU Embedding layers for Serving do
not properly restore slot variables. This option is a way to omit
restoring slot variables which are not required for Serving usecase
anyways.(b/315912101)
experimental_sharding_callback: `tf.train.experimental.ShardingCallback`.
A pre-made or custom callback that determines how checkpoints are
sharded on disk. Pre-made callback options are
`tf.train.experimental.ShardByDevicePolicy` and
`tf.train.experimental.MaxShardSizePolicy`. You may also write a custom
callback, see `tf.train.experimental.ShardingCallback`.
"""
self.experimental_io_device = experimental_io_device
self.enable_async = experimental_enable_async_checkpoint or enable_async
self.experimental_enable_async_checkpoint = self.enable_async
# Ensure that each callback only has either 0 or 1 parameter
if experimental_write_callbacks is not None:
for callback in experimental_write_callbacks:
assert len(inspect.signature(callback).parameters) <= 1
self.experimental_write_callbacks = experimental_write_callbacks
if experimental_sharding_callback is not None:
if not isinstance(
experimental_sharding_callback, sharding_util.ShardingCallback):
raise ValueError("The experimental_sharding_callback checkpoint option"
"must be of type ShardingCallback. The option provided"
f"was of type {type(experimental_sharding_callback)}.")
self.experimental_sharding_callback = experimental_sharding_callback
self.experimental_skip_slot_variables = experimental_skip_slot_variables
def __copy__(self):
# Only `experimental_write_callbacks` needs special treatment to Ensure that
# the list is deep-copied, but the callbacks are not deep-copied.
cls = self.__class__
result = cls.__new__(cls)
for name in self.__slots__:
if hasattr(self, name):
setattr(result, name, getattr(self, name))
result.experimental_write_callbacks = copy.copy(
self.experimental_write_callbacks
)
return result
tensorflow/python/checkpoint/checkpoint_test.py
+1504
View File
File diff suppressed because it is too large Load Diff
tensorflow/python/checkpoint/checkpoint_view.py
+301
View File
@@ -0,0 +1,301 @@
"""Manages a Checkpoint View."""
# 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.
# ==============================================================================
import collections
from tensorflow.core.protobuf import trackable_object_graph_pb2
from tensorflow.python.checkpoint import trackable_view
from tensorflow.python.framework import errors_impl
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.trackable import base
from tensorflow.python.training import py_checkpoint_reader
from tensorflow.python.util import object_identity
from tensorflow.python.util.tf_export import tf_export
@tf_export("train.CheckpointView", v1=[])
class CheckpointView(object):
"""Gathers and serializes a checkpoint view.
This is for loading specific portions of a module from a
checkpoint, and be able to compare two modules by matching components.
Example usage:
>>> class SimpleModule(tf.Module):
... def __init__(self, name=None):
... super().__init__(name=name)
... self.a_var = tf.Variable(5.0)
... self.b_var = tf.Variable(4.0)
... self.vars = [tf.Variable(1.0), tf.Variable(2.0)]
>>> root = SimpleModule(name="root")
>>> root.leaf = SimpleModule(name="leaf")
>>> ckpt = tf.train.Checkpoint(root)
>>> save_path = ckpt.save('/tmp/tf_ckpts')
>>> checkpoint_view = tf.train.CheckpointView(save_path)
Pass `node_id=0` to `tf.train.CheckpointView.children()` to get the dictionary
of all children directly linked to the checkpoint root.
>>> for name, node_id in checkpoint_view.children(0).items():
... print(f"- name: '{name}', node_id: {node_id}")
- name: 'a_var', node_id: 1
- name: 'b_var', node_id: 2
- name: 'vars', node_id: 3
- name: 'leaf', node_id: 4
- name: 'root', node_id: 0
- name: 'save_counter', node_id: 5
"""
def __init__(self, save_path):
"""Configure the checkpoint view.
Args:
save_path: The path to the checkpoint.
Raises:
ValueError: If the save_path does not lead to a TF2 checkpoint.
"""
reader = py_checkpoint_reader.NewCheckpointReader(save_path)
try:
object_graph_string = reader.get_tensor(base.OBJECT_GRAPH_PROTO_KEY)
except errors_impl.NotFoundError as not_found_error:
raise ValueError(
f"The specified checkpoint \"{save_path}\" does not appear to be "
"object-based (saved with TF2) since it is missing the key "
f"\"{base.OBJECT_GRAPH_PROTO_KEY}\". Likely it was created with the "
"TF1 name-based saver and does not contain an object dependency graph."
) from not_found_error
object_graph_proto = (trackable_object_graph_pb2.TrackableObjectGraph())
object_graph_proto.ParseFromString(object_graph_string)
self._object_graph_proto = object_graph_proto
def children(self, node_id):
"""Returns all child trackables attached to obj.
Args:
node_id: Id of the node to return its children.
Returns:
Dictionary of all children attached to the object with name to node_id.
"""
return {
child.local_name: child.node_id
for child in self._object_graph_proto.nodes[node_id].children
}
def descendants(self):
"""Returns a list of trackables by node_id attached to obj."""
return list(self._descendants_with_paths().keys())
def _descendants_with_paths(self):
"""Returns a dict of descendants by node_id and paths to node.
The names returned by this private method are subject to change.
"""
all_nodes_with_paths = {}
to_visit = collections.deque([0])
# node_id:0 will always be "root".
all_nodes_with_paths[0] = "root"
path = all_nodes_with_paths.get(0)
while to_visit:
node_id = to_visit.popleft()
obj = self._object_graph_proto.nodes[node_id]
for child in obj.children:
if child.node_id == 0 or child.node_id in all_nodes_with_paths.keys():
continue
path = all_nodes_with_paths.get(node_id)
if child.node_id not in all_nodes_with_paths.keys():
to_visit.append(child.node_id)
all_nodes_with_paths[child.node_id] = path + "." + child.local_name
return all_nodes_with_paths
def match(self, obj):
"""Returns all matching trackables between CheckpointView and Trackable.
Matching trackables represents trackables with the same name and position in
graph.
Args:
obj: `Trackable` root.
Returns:
Dictionary containing all overlapping trackables that maps `node_id` to
`Trackable`.
Example usage:
>>> class SimpleModule(tf.Module):
... def __init__(self, name=None):
... super().__init__(name=name)
... self.a_var = tf.Variable(5.0)
... self.b_var = tf.Variable(4.0)
... self.vars = [tf.Variable(1.0), tf.Variable(2.0)]
>>> root = SimpleModule(name="root")
>>> leaf = root.leaf = SimpleModule(name="leaf")
>>> leaf.leaf3 = tf.Variable(6.0, name="leaf3")
>>> leaf.leaf4 = tf.Variable(7.0, name="leaf4")
>>> ckpt = tf.train.Checkpoint(root)
>>> save_path = ckpt.save('/tmp/tf_ckpts')
>>> checkpoint_view = tf.train.CheckpointView(save_path)
>>> root2 = SimpleModule(name="root")
>>> leaf2 = root2.leaf2 = SimpleModule(name="leaf2")
>>> leaf2.leaf3 = tf.Variable(6.0)
>>> leaf2.leaf4 = tf.Variable(7.0)
Pass `node_id=0` to `tf.train.CheckpointView.children()` to get the
dictionary of all children directly linked to the checkpoint root.
>>> checkpoint_view_match = checkpoint_view.match(root2).items()
>>> for item in checkpoint_view_match:
... print(item)
(0, ...)
(1, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=5.0>)
(2, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=4.0>)
(3, ListWrapper([<tf.Variable 'Variable:0' shape=() dtype=float32,
numpy=1.0>, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>]))
(6, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>)
(7, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>)
"""
if not isinstance(obj, base.Trackable):
raise ValueError(f"Expected a Trackable, got {obj} of type {type(obj)}.")
overlapping_nodes = {}
# Root node is always matched.
overlapping_nodes[0] = obj
# Queue of tuples of node_id and trackable.
to_visit = collections.deque([(0, obj)])
visited = set()
view = trackable_view.TrackableView(obj)
while to_visit:
current_node_id, current_trackable = to_visit.popleft()
trackable_children = view.children(current_trackable)
for child_name, child_node_id in self.children(current_node_id).items():
if child_node_id in visited or child_node_id == 0:
continue
if child_name in trackable_children:
current_assignment = overlapping_nodes.get(child_node_id)
if current_assignment is None:
overlapping_nodes[child_node_id] = trackable_children[child_name]
to_visit.append((child_node_id, trackable_children[child_name]))
else:
# The object was already mapped for this checkpoint load, which
# means we don't need to do anything besides check that the mapping
# is consistent (if the dependency DAG is not a tree then there are
# multiple paths to the same object).
if current_assignment is not trackable_children[child_name]:
logging.warning(
"Inconsistent references when matching the checkpoint into "
"this object graph. The referenced objects are: "
f"({current_assignment} and "
f"{trackable_children[child_name]}).")
visited.add(current_node_id)
return overlapping_nodes
def diff(self, obj):
"""Returns diff between CheckpointView and Trackable.
This method is intended to be used to compare the object stored in a
checkpoint vs a live model in Python. For example, if checkpoint
restoration fails the `assert_consumed()` or
`assert_existing_objects_matched()` checks, you can use this to list out
the objects/checkpoint nodes which were not restored.
Example Usage:
>>> class SimpleModule(tf.Module):
... def __init__(self, name=None):
... super().__init__(name=name)
... self.a_var = tf.Variable(5.0)
... self.b_var = tf.Variable(4.0)
... self.vars = [tf.Variable(1.0), tf.Variable(2.0)]
>>> root = SimpleModule(name="root")
>>> leaf = root.leaf = SimpleModule(name="leaf")
>>> leaf.leaf3 = tf.Variable(6.0, name="leaf3")
>>> leaf.leaf4 = tf.Variable(7.0, name="leaf4")
>>> ckpt = tf.train.Checkpoint(root)
>>> save_path = ckpt.save('/tmp/tf_ckpts')
>>> checkpoint_view = tf.train.CheckpointView(save_path)
>>> root2 = SimpleModule(name="root")
>>> leaf2 = root2.leaf2 = SimpleModule(name="leaf2")
>>> leaf2.leaf3 = tf.Variable(6.0)
>>> leaf2.leaf4 = tf.Variable(7.0)
Pass `node_id=0` to `tf.train.CheckpointView.children()` to get the
dictionary of all children directly linked to the checkpoint root.
>>> checkpoint_view_diff = checkpoint_view.diff(root2)
>>> checkpoint_view_match = checkpoint_view_diff[0].items()
>>> for item in checkpoint_view_match:
... print(item)
(0, ...)
(1, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=5.0>)
(2, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=4.0>)
(3, ListWrapper([<tf.Variable 'Variable:0' shape=() dtype=float32,
numpy=1.0>, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>]))
(6, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>)
(7, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>)
>>> only_in_checkpoint_view = checkpoint_view_diff[1]
>>> print(only_in_checkpoint_view)
[4, 5, 8, 9, 10, 11, 12, 13, 14]
>>> only_in_trackable = checkpoint_view_diff[2]
>>> print(only_in_trackable)
[..., <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=5.0>,
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=4.0>,
ListWrapper([<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>]),
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=6.0>,
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=7.0>,
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=1.0>,
<tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>]
Args:
obj: `Trackable` root.
Returns:
Tuple of (
- Overlaps: Dictionary containing all overlapping trackables that maps
`node_id` to `Trackable`, same as CheckpointView.match().
- Only in CheckpointView: List of `node_id` that only exist in
CheckpointView.
- Only in Trackable: List of `Trackable` that only exist in Trackable.
)
"""
overlapping_nodes = self.match(obj)
only_in_checkpoint_view = []
only_in_trackable = []
for node_id in self.descendants():
if node_id not in overlapping_nodes.keys():
only_in_checkpoint_view.append(node_id)
for trackable in trackable_view.TrackableView(obj).descendants():
if trackable not in object_identity.ObjectIdentitySet(
overlapping_nodes.values()):
only_in_trackable.append(trackable)
return overlapping_nodes, only_in_checkpoint_view, only_in_trackable
tensorflow/python/checkpoint/checkpoint_view_test.py
+148
View File
@@ -0,0 +1,148 @@
# 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 the checkpoint view."""
import os
from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.checkpoint import checkpoint_view
from tensorflow.python.eager import test
from tensorflow.python.trackable import autotrackable
class CheckpointViewTest(test.TestCase):
def test_children(self):
root = autotrackable.AutoTrackable()
root.leaf = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
current_name, node_id = next(
iter(
checkpoint_view.CheckpointView(root_save_path).children(0).items()))
self.assertEqual("leaf", current_name)
self.assertEqual(1, node_id)
def test_all_nodes(self):
root = autotrackable.AutoTrackable()
root.leaf = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
all_nodes = checkpoint_view.CheckpointView(root_save_path).descendants()
self.assertEqual(3, len(all_nodes))
self.assertEqual(0, all_nodes[0])
self.assertEqual(1, all_nodes[1])
def test_all_nodes_with_paths(self):
root = autotrackable.AutoTrackable()
leaf1 = root.leaf1 = autotrackable.AutoTrackable()
leaf2 = root.leaf2 = autotrackable.AutoTrackable()
leaf1.leaf3 = autotrackable.AutoTrackable()
leaf1.leaf4 = autotrackable.AutoTrackable()
leaf2.leaf5 = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
all_nodes_with_paths = checkpoint_view.CheckpointView(
root_save_path)._descendants_with_paths()
self.assertEqual(
{
"root", "root.leaf1", "root.leaf2", "root.save_counter",
"root.leaf1.leaf3", "root.leaf1.leaf4", "root.leaf2.leaf5"
}, set(all_nodes_with_paths.values()))
def test_match(self):
root1 = autotrackable.AutoTrackable()
leaf1 = root1.leaf1 = autotrackable.AutoTrackable()
leaf2 = root1.leaf2 = autotrackable.AutoTrackable()
leaf1.leaf3 = autotrackable.AutoTrackable()
leaf1.leaf4 = autotrackable.AutoTrackable()
leaf2.leaf5 = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root1)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = autotrackable.AutoTrackable()
leaf11 = root2.leaf1 = autotrackable.AutoTrackable()
leaf12 = root2.leaf2 = autotrackable.AutoTrackable()
leaf13 = leaf11.leaf3 = autotrackable.AutoTrackable()
leaf15 = leaf12.leaf5 = autotrackable.AutoTrackable()
matching_nodes = checkpoint_view.CheckpointView(root_save_path).match(root2)
self.assertDictEqual(matching_nodes, {
0: root2,
1: leaf11,
2: leaf12,
4: leaf13,
6: leaf15
})
def test_match_overlapping_nodes(self):
root1 = autotrackable.AutoTrackable()
root1.a = root1.b = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root1)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = autotrackable.AutoTrackable()
a1 = root2.a = autotrackable.AutoTrackable()
root2.b = autotrackable.AutoTrackable()
with self.assertLogs(level="WARNING") as logs:
matching_nodes = checkpoint_view.CheckpointView(root_save_path).match(
root2)
self.assertDictEqual(
matching_nodes,
{
0: root2,
1: a1,
# Only the first element at the same position will be matched.
})
expected_message = (
"Inconsistent references when matching the checkpoint into this object"
" graph.")
self.assertIn(expected_message, logs.output[0])
def test_diff(self):
root1 = autotrackable.AutoTrackable()
leaf1 = root1.leaf1 = autotrackable.AutoTrackable()
leaf2 = root1.leaf2 = autotrackable.AutoTrackable()
leaf1.leaf3 = autotrackable.AutoTrackable()
leaf1.leaf4 = autotrackable.AutoTrackable()
leaf2.leaf5 = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root1)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = autotrackable.AutoTrackable()
leaf11 = root2.leaf1 = autotrackable.AutoTrackable()
leaf12 = root2.leaf2 = autotrackable.AutoTrackable()
leaf13 = leaf11.leaf3 = autotrackable.AutoTrackable()
leaf15 = leaf12.leaf5 = autotrackable.AutoTrackable()
leaf16 = leaf12.leaf6 = autotrackable.AutoTrackable()
diff = checkpoint_view.CheckpointView(root_save_path).diff(root2)
self.assertEqual(len(diff), 3)
self.assertDictEqual(diff[0], {
0: root2,
1: leaf11,
2: leaf12,
4: leaf13,
6: leaf15
})
self.assertListEqual(diff[1], [3, 5])
self.assertListEqual(diff[2], [leaf16])
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/checkpoint_with_v1_optimizers_test.py
+311
View File
@@ -0,0 +1,311 @@
# 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.
# ==============================================================================
"""Tests for object-based saving which use tf.train.* optimizers."""
import os
from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.checkpoint import checkpoint_options as options
from tensorflow.python.client import session as session_lib
from tensorflow.python.eager import context
from tensorflow.python.eager import test
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import template
from tensorflow.python.ops import variable_scope
from tensorflow.python.trackable import autotrackable
from tensorflow.python.training import adam
from tensorflow.python.training import checkpoint_utils
class CheckpointingTests(test.TestCase):
@test_util.run_in_graph_and_eager_modes
def testDeferredSlotRestoration(self):
checkpoint_directory = self.get_temp_dir()
root = trackable_utils.Checkpoint()
root.var = trackable_utils.add_variable(
root, name="var", initializer=0.)
optimizer = adam.AdamOptimizer(0.1)
if context.executing_eagerly():
optimizer.minimize(root.var.read_value)
else:
train_op = optimizer.minimize(root.var)
# Note that `optimizer` has not been added as a dependency of
# `root`. Create a one-off grouping so that slot variables for `root.var`
# get initialized too.
self.evaluate(trackable_utils.gather_initializers(
trackable_utils.Checkpoint(root=root, optimizer=optimizer)))
self.evaluate(train_op)
self.evaluate(state_ops.assign(root.var, 12.))
no_slots_path = root.save(os.path.join(checkpoint_directory, "no_slots"))
root.optimizer = optimizer
self.evaluate(state_ops.assign(root.var, 13.))
self.evaluate(state_ops.assign(optimizer.get_slot(name="m", var=root.var),
14.))
slots_path = root.save(os.path.join(checkpoint_directory, "with_slots"))
new_root = trackable_utils.Checkpoint()
# Load the slot-containing checkpoint (deferred), then immediately overwrite
# the non-slot variable (also deferred).
slot_status = new_root.restore(slots_path)
no_slot_status = new_root.restore(no_slots_path)
with self.assertRaises(AssertionError):
no_slot_status.assert_consumed()
new_root.var = trackable_utils.add_variable(
new_root, name="var", shape=[])
no_slot_status.assert_consumed()
no_slot_status.run_restore_ops()
self.assertEqual(12., self.evaluate(new_root.var))
new_root.optimizer = adam.AdamOptimizer(0.1)
slot_status.assert_existing_objects_matched()
with self.assertRaisesRegex(AssertionError, "beta1_power"):
slot_status.assert_consumed()
self.assertEqual(12., self.evaluate(new_root.var))
if context.executing_eagerly():
# Slot variables are only created with restoring initializers when
# executing eagerly.
self.assertEqual(14., self.evaluate(
new_root.optimizer.get_slot(name="m", var=new_root.var)))
else:
self.assertIs(new_root.optimizer.get_slot(name="m", var=new_root.var),
None)
if context.executing_eagerly():
new_root.optimizer.minimize(new_root.var.read_value)
else:
train_op = new_root.optimizer.minimize(new_root.var)
# The slot variable now exists; restore() didn't create it, but we should
# now have a restore op for it.
slot_status.run_restore_ops()
self.assertEqual(14., self.evaluate(
new_root.optimizer.get_slot(name="m", var=new_root.var)))
self.evaluate(train_op)
slot_status.assert_consumed()
def testSkipSlotVarRestore(self):
def create_trackable():
root = trackable_utils.Checkpoint()
root.var = trackable_utils.add_variable(root, name="var", initializer=0.0)
root.optimizer = adam.AdamOptimizer(0.1)
root.optimizer._create_slots([root.var])
return root
root = create_trackable()
checkpoint_directory = self.get_temp_dir()
self.evaluate(state_ops.assign(root.var, 12.0))
self.evaluate(
state_ops.assign(root.optimizer.get_slot(name="m", var=root.var), 14.0)
)
chkpt_path = root.save(os.path.join(checkpoint_directory, "chkpt"))
new_root = create_trackable()
chkpt_keys = [x[0] for x in checkpoint_utils.list_variables(chkpt_path)]
self.assertEqual(
chkpt_keys,
[
"_CHECKPOINTABLE_OBJECT_GRAPH",
"optimizer/beta1_power/.ATTRIBUTES/VARIABLE_VALUE",
"optimizer/beta2_power/.ATTRIBUTES/VARIABLE_VALUE",
"save_counter/.ATTRIBUTES/VARIABLE_VALUE",
"var/.ATTRIBUTES/VARIABLE_VALUE",
"var/.OPTIMIZER_SLOT/optimizer/m/.ATTRIBUTES/VARIABLE_VALUE",
"var/.OPTIMIZER_SLOT/optimizer/v/.ATTRIBUTES/VARIABLE_VALUE",
],
)
# Restore without slot variables
slot_restore = new_root.restore(
chkpt_path,
options.CheckpointOptions(experimental_skip_slot_variables=True),
)
slot_restore.assert_consumed().run_restore_ops()
# regular variable is restored
self.assertEqual(self.evaluate(new_root.var), 12.0)
# Slot variable is not restored
self.assertEqual(
self.evaluate(new_root.optimizer.get_slot(name="m", var=new_root.var)),
0.0,
)
# Restore everything
slot_restore = new_root.restore(
chkpt_path,
options.CheckpointOptions(experimental_skip_slot_variables=False),
)
slot_restore.assert_consumed().run_restore_ops()
self.assertEqual(self.evaluate(new_root.var), 12.0)
# Slot variable also restored
self.assertEqual(
self.evaluate(new_root.optimizer.get_slot(name="m", var=new_root.var)),
14.0,
)
def testManySavesGraph(self):
"""Saves after the first should not modify the graph."""
with context.graph_mode():
graph = ops.Graph()
with graph.as_default(), self.session(graph):
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
obj = trackable_utils.Checkpoint()
obj.var = variable_scope.get_variable(name="v", initializer=0.)
obj.opt = adam.AdamOptimizer(0.1)
obj.opt.minimize(obj.var.read_value())
self.evaluate(trackable_utils.gather_initializers(obj))
obj.save(checkpoint_prefix)
before_ops = graph.get_operations()
obj.save(checkpoint_prefix)
self.assertEqual(before_ops, graph.get_operations())
def testManyRestoresGraph(self):
"""Restores after the first should not modify the graph."""
with context.graph_mode():
graph = ops.Graph()
with graph.as_default(), self.session(graph):
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
obj = trackable_utils.Checkpoint()
obj.var = variable_scope.get_variable(name="v", initializer=0.)
obj.opt = adam.AdamOptimizer(0.1)
obj.opt.minimize(obj.var.read_value())
self.evaluate(trackable_utils.gather_initializers(obj))
save_path = obj.save(checkpoint_prefix)
obj.restore(save_path)
before_ops = graph.get_operations()
obj.restore(save_path)
self.assertEqual(before_ops, graph.get_operations())
def testMultipleGraphsNonSlotVariables(self):
with context.graph_mode():
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
optimizer = adam.AdamOptimizer(0.001)
# Construct a model in one graph
first_graph = ops.Graph()
first_session = session_lib.Session(graph=first_graph)
with first_graph.as_default(), first_session.as_default():
first_variable = resource_variable_ops.ResourceVariable([1.])
first_root_trackable = trackable_utils.Checkpoint(
optimizer=optimizer, variable=first_variable)
train_op = optimizer.minimize(first_variable.read_value)
self.evaluate(trackable_utils.gather_initializers(
first_root_trackable))
self.evaluate(train_op)
self.evaluate(first_variable.assign([1.]))
self.evaluate(optimizer.get_slot(
var=first_variable, name="m").assign([2.]))
beta1_power, _ = optimizer._get_beta_accumulators()
self.evaluate(beta1_power.assign(3.))
# Save and load in a second graph
second_graph = ops.Graph()
with second_graph.as_default(), session_lib.Session(graph=second_graph):
second_variable = resource_variable_ops.ResourceVariable([1.])
second_root_trackable = trackable_utils.Checkpoint(
optimizer=optimizer, variable=second_variable)
train_op = optimizer.minimize(second_variable.read_value)
second_root_trackable.restore(None).initialize_or_restore()
self.evaluate(train_op)
self.evaluate(second_variable.assign([4.]))
self.evaluate(optimizer.get_slot(
var=second_variable, name="m").assign([5.]))
beta1_power, _ = optimizer._get_beta_accumulators()
self.evaluate(beta1_power.assign(6.))
save_path = second_root_trackable.save(checkpoint_prefix)
self.evaluate(second_variable.assign([7.]))
self.evaluate(optimizer.get_slot(
var=second_variable, name="m").assign([8.]))
beta1_power, _ = optimizer._get_beta_accumulators()
self.assertAllEqual(6., self.evaluate(beta1_power))
status = second_root_trackable.restore(save_path)
status.assert_consumed().run_restore_ops()
self.assertAllEqual([4.], self.evaluate(second_variable))
self.assertAllEqual([5.], self.evaluate(optimizer.get_slot(
var=second_variable, name="m")))
beta1_power, _ = optimizer._get_beta_accumulators()
self.assertAllEqual(6., self.evaluate(beta1_power))
# Check that the first graph is unmolested
with first_graph.as_default(), first_session.as_default():
self.assertAllEqual([1.], self.evaluate(first_variable))
self.assertAllEqual([2.], self.evaluate(optimizer.get_slot(
var=first_variable, name="m")))
beta1_power, _ = optimizer._get_beta_accumulators()
self.assertAllEqual(3., self.evaluate(beta1_power))
class _ManualScope(autotrackable.AutoTrackable):
def __call__(self):
with variable_scope.variable_scope("ManualScope") as vs:
self.variable_scope = vs
with trackable_utils.capture_dependencies(template=self):
return self._build()
def _build(self):
return variable_scope.get_variable(name="in_manual_scope", shape=[])
class TemplateTests(test.TestCase):
@test_util.run_in_graph_and_eager_modes
def test_trackable_save_restore(self):
def _templated():
v = variable_scope.get_variable(
"v", shape=[1], initializer=init_ops.zeros_initializer(),
use_resource=True)
v2 = variable_scope.get_variable(
"v2", shape=[1], initializer=init_ops.zeros_initializer(),
use_resource=True)
manual = _ManualScope()
return v, v + 1., v2, manual, manual()
save_template = template.make_template("s1", _templated)
v1_save, _, v2_save, manual_scope, manual_scope_v = save_template()
self.assertCountEqual([
id(obj) for obj in
[v1_save, v2_save, manual_scope, manual_scope_v, save_template]
], [id(obj) for obj in trackable_utils.list_objects(save_template)])
self.assertDictEqual({"in_manual_scope": manual_scope_v},
manual_scope._trackable_children())
optimizer = adam.AdamOptimizer(0.0)
save_root = trackable_utils.Checkpoint(
my_template=save_template, optimizer=optimizer)
optimizer.minimize(v1_save.read_value)
self.evaluate([v.initializer for v in save_template.variables])
self.evaluate([v.initializer for v in optimizer.variables()])
self.evaluate(v1_save.assign([12.]))
self.evaluate(v2_save.assign([14.]))
checkpoint_directory = self.get_temp_dir()
checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
save_path = save_root.save(checkpoint_prefix)
load_template = template.make_template("s2", _templated)
load_optimizer = adam.AdamOptimizer(0.0)
load_root = trackable_utils.Checkpoint(
my_template=load_template, optimizer=load_optimizer)
status = load_root.restore(save_path)
var, var_plus_one, var2, _, _ = load_template()
load_optimizer.minimize(var.read_value)
self.assertEqual(3, len(load_template._trackable_children()))
self.assertEqual(set(["v", "v2", "ManualScope"]),
load_template._trackable_children().keys())
status.assert_consumed().run_restore_ops()
self.assertAllEqual([12.], self.evaluate(var))
self.assertAllEqual([13.], self.evaluate(var_plus_one))
self.assertAllEqual([14.], self.evaluate(var2))
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/functional_saver.py
+661
View File
@@ -0,0 +1,661 @@
# Copyright 2015 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.
# ==============================================================================
"""Saves and restore variables inside traced @tf.functions."""
import dataclasses
import math
import time
from typing import Callable, Mapping, MutableMapping, MutableSequence, Sequence
from absl import logging
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.checkpoint.sharding import sharding_policies
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import context
from tensorflow.python.eager import def_function
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_spec
from tensorflow.python.framework import tensor_util
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_io_ops
from tensorflow.python.ops import io_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.saved_model import registration
from tensorflow.python.saved_model.pywrap_saved_model import metrics
from tensorflow.python.trackable import base
from tensorflow.python.trackable import trackable_utils
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.types import core
from tensorflow.python.util import nest
from tensorflow.python.util import object_identity
RegisteredSaversDict = Mapping[
registration.RegisteredSaver, Mapping[str, base.Trackable]]
MappedCapturesCallable = Callable[
[core.ConcreteFunction, Sequence[tensor_lib.Tensor]], tensor_lib.Tensor]
def _single_shard_save(
file_prefix: tensor_lib.Tensor,
shard: sharding_util.Shard,
task: device_lib.DeviceSpec,
options: "checkpoint_options.CheckpointOptions | None" = None,
) -> ops.Operation:
"""Save the saveable objects to a checkpoint with `file_prefix`.
Args:
file_prefix: A string or scalar string Tensor containing the prefix to
save under.
shard: Dict containing tensors. {checkpoint key: {slice_spec: tensor} }
task: The device spec task of the tensors in the shard.
options: Optional `CheckpointOptions` object.
Returns:
An `Operation`, or None when executing eagerly.
"""
options = options or checkpoint_options.CheckpointOptions()
tensor_names = []
tensors = []
slice_specs = []
for checkpoint_key, tensor_slices in shard.items():
for slice_spec, tensor in tensor_slices.items():
# A tensor value of `None` indicates that this SaveableObject gets
# recorded in the object graph, but that no value is saved in the
# checkpoint.
if tensor is not None:
# See `MultiDeviceSaver._get_shards_by_task` for an explanation on the
# wrapped properties.
name = (tensor._wrapped_name # pylint: disable=protected-access
if hasattr(tensor, "_wrapped_name")
else checkpoint_key)
spec = (tensor._wrapped_slice_spec # pylint: disable=protected-access
if hasattr(tensor, "_wrapped_slice_spec")
else slice_spec)
tensor_names.append(name)
tensors.append(tensor)
slice_specs.append(spec)
save_device = options.experimental_io_device or (tensors and task)
with ops.device(save_device or "CPU:0"):
return io_ops.save_v2(file_prefix, tensor_names, slice_specs, tensors)
def _single_shard_restore(
file_prefix: tensor_lib.Tensor,
shardable_tensors: Sequence[sharding_util.ShardableTensor],
options: "checkpoint_options.CheckpointOptions | None" = None
) -> sharding_util.Shard:
"""Restore the saveable objects from a checkpoint with `file_prefix`.
Args:
file_prefix: A string or scalar string Tensor containing the prefix for
files to read from.
shardable_tensors: A list of ShardableTensors to restore.
options: Optional `CheckpointOptions` object.
Returns:
A restored tensor dict (maps checkpoint_key -> slice_spec -> tensor).
"""
options = options or checkpoint_options.CheckpointOptions()
tensor_names = []
tensor_dtypes = []
slice_specs = []
for shardable_tensor in shardable_tensors:
if shardable_tensor._tensor_save_spec: # pylint: disable=protected-access
name = shardable_tensor._tensor_save_spec.name # pylint: disable=protected-access
spec = shardable_tensor._tensor_save_spec.slice_spec # pylint: disable=protected-access
else:
name, spec = shardable_tensor.checkpoint_key, shardable_tensor.slice_spec
tensor_names.append(name)
slice_specs.append(spec)
tensor_dtypes.append(shardable_tensor.dtype)
restore_device = options.experimental_io_device or "cpu:0"
with ops.device(restore_device):
restored_tensors = io_ops.restore_v2(
file_prefix, tensor_names, slice_specs, tensor_dtypes)
restored_tensor_dict = {}
for shardable_tensor in shardable_tensors:
restored_tensor = restored_tensors.pop(0)
(restored_tensor_dict
.setdefault(shardable_tensor.checkpoint_key, {}
)[shardable_tensor.slice_spec]) = restored_tensor
return restored_tensor_dict
def sharded_filename(
filename_tensor: tensor_lib.Tensor,
shard: int,
num_shards: tensor_lib.Tensor
) -> tensor_lib.Tensor:
"""Append sharding information to a filename.
Args:
filename_tensor: A string tensor.
shard: Integer. The shard for the filename.
num_shards: An int Tensor for the number of shards.
Returns:
A string tensor.
"""
return gen_io_ops.sharded_filename(filename_tensor, shard, num_shards)
def registered_saver_filename(
filename_tensor: tensor_lib.Tensor,
saver_name: registration.RegisteredSaver
) -> tensor_lib.Tensor:
return string_ops.string_join(
[filename_tensor, constant_op.constant(f"-{saver_name}")])
def _get_mapped_registered_save_fn(
fn: Callable[..., tensor_lib.Tensor],
trackables: Sequence[base.Trackable],
call_with_mapped_captures: MappedCapturesCallable
) -> Callable[[tensor_lib.Tensor], MappedCapturesCallable]:
"""Converts the function to a python or tf.function with a single file arg."""
def save_fn(file_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
return fn(trackables=trackables, file_prefix=file_prefix)
if call_with_mapped_captures is None:
return save_fn
else:
tf_fn = def_function.function(save_fn, autograph=False)
concrete = tf_fn.get_concrete_function(
file_prefix=tensor_spec.TensorSpec(shape=(), dtype=dtypes.string))
def save_fn_with_replaced_captures(
file_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
return call_with_mapped_captures(concrete, [file_prefix])
return save_fn_with_replaced_captures
def _get_mapped_registered_restore_fn(
fn: Callable[..., tensor_lib.Tensor],
trackables: Sequence[base.Trackable],
call_with_mapped_captures: MappedCapturesCallable
) -> Callable[..., tensor_lib.Tensor]:
"""Converts the function to a python or tf.function with a single file arg."""
def restore_fn(merged_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
return fn(trackables=trackables, merged_prefix=merged_prefix)
if call_with_mapped_captures is None:
return restore_fn
else:
tf_fn = def_function.function(restore_fn, autograph=False)
concrete = tf_fn.get_concrete_function(
merged_prefix=tensor_spec.TensorSpec(shape=(), dtype=dtypes.string))
def restore_fn_with_replaced_captures(
merged_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
return call_with_mapped_captures(concrete, [merged_prefix])
return restore_fn_with_replaced_captures
_restore_noop = lambda *args, **kwargs: None
TensorKeyAndSliceSpec = tuple[str, str]
RestoreFn = Callable[[Mapping[str, tensor_lib.Tensor]], ops.Operation]
class MultiDeviceSaver:
"""Saves checkpoints directly from multiple devices.
Note that this is a low-level utility which stores Tensors in the keys
specified by `SaveableObject`s. Higher-level utilities for object-based
checkpointing are built on top of it.
"""
def __init__(
self,
serialized_tensors: Mapping[base.Trackable, sharding_util.Shard],
registered_savers: "RegisteredSaversDict | None" = None,
call_with_mapped_captures: "MappedCapturesCallable | None" = None):
"""Specify a list of `SaveableObject`s to save and restore.
Args:
serialized_tensors: A dictionary mapping `Trackable` to a tensor dict,
which maps checkpoint_key -> (slice_spec ->) -> Tensor/SaveSpec. The
`Trackable` key is used to get the `restore_from_tensors` function,
and may be `None` if the tensor is not meant to be restored.
registered_savers: A dictionary mapping `registration.RegisteredSaver`
namedtuples to a dictionary of named Trackables. The keys of the
Trackable dictionary are string names that uniquely identify the
Trackable in the checkpoint.
call_with_mapped_captures: TODO
"""
self._shardable_tensors_by_task: MutableMapping[
device_lib.DeviceSpec,
MutableSequence[sharding_util.ShardableTensor]] = {}
# Keep these two data structures so that we can map restored tensors to
# the Trackable restore functions.
self._keys_to_restore_fn: MutableMapping[
TensorKeyAndSliceSpec, RestoreFn] = {}
self._restore_fn_to_keys: MutableMapping[
RestoreFn, MutableSequence[TensorKeyAndSliceSpec]] = {}
unique_tasks = set()
for obj, tensor_dict in serialized_tensors.items():
restore_fn = _restore_noop if obj is None else obj._restore_from_tensors
# Divide tensor_dict by task.
for checkpoint_key, tensor_slice_dict in tensor_dict.items():
if not isinstance(tensor_slice_dict, dict):
# Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
tensor_slice_dict = {"": tensor_slice_dict}
for slice_spec, tensor_save_spec in tensor_slice_dict.items():
tensor_value = None
if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
tensor_value = tensor_save_spec
tensor_save_spec = saveable_object.SaveSpec(
tensor=tensor_value,
slice_spec=slice_spec,
name=checkpoint_key,
dtype=tensor_save_spec.dtype,
device=tensor_save_spec.device)
if (checkpoint_key, slice_spec) in self._keys_to_restore_fn:
raise ValueError(
"Received multiple tensors with the same checkpoint key and "
"slice spec. This is invalid because one will overwrite the "
"other in the checkpoint. This indicates a bug in the "
"Checkpoint key-generation.")
self._keys_to_restore_fn[(checkpoint_key, slice_spec)] = restore_fn
self._restore_fn_to_keys.setdefault(restore_fn, []).append(
(checkpoint_key, slice_spec))
if isinstance(tensor_save_spec.device, str):
device = device_lib.DeviceSpec.from_string(tensor_save_spec.device)
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(tensor_save_spec.device))
else:
device = tensor_save_spec.device
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(device.to_string()))
self._shardable_tensors_by_task.setdefault(task, []).append(
sharding_util.ShardableTensor(
_tensor_save_spec=tensor_save_spec,
tensor=tensor_value,
dtype=tensor_save_spec.dtype,
device=device,
name=tensor_save_spec.name,
shape=None,
slice_spec=slice_spec.strip(),
checkpoint_key=checkpoint_key,
trackable=obj))
unique_tasks.add(
saveable_object_util.set_cpu0(device.to_string()))
self._num_unique_tasks = len(unique_tasks)
self._registered_savers = {}
if registered_savers:
for registered_name, trackables in registered_savers.items():
save_fn = _get_mapped_registered_save_fn(
registration.get_save_function(registered_name),
trackables, call_with_mapped_captures)
restore_fn = _get_mapped_registered_restore_fn(
registration.get_restore_function(registered_name),
trackables, call_with_mapped_captures)
self._registered_savers[registered_name] = (save_fn, restore_fn)
@classmethod
def from_saveables(
cls,
saveables: Sequence[base.Trackable],
registered_savers: "RegisteredSaversDict | None" = None,
call_with_mapped_captures: "MappedCapturesCallable | None" = None
) -> "MultiDeviceSaver":
"""Constructs a MultiDeviceSaver from a list of `SaveableObject`s."""
serialized_tensors = object_identity.ObjectIdentityDictionary()
for saveable in saveables:
trackable = saveable_object_util.SaveableCompatibilityConverter(
saveable, saveables=[saveable])
serialized_tensors[trackable] = trackable._serialize_to_tensors() # pylint: disable=protected-access
return cls(serialized_tensors, registered_savers, call_with_mapped_captures)
def to_proto(self) -> saver_pb2.SaverDef:
"""Serializes to a SaverDef referencing the current graph."""
filename_tensor = array_ops.placeholder(
shape=[], dtype=dtypes.string, name="saver_filename")
save_tensor = self._traced_save(filename_tensor)
restore_op = self._traced_restore(filename_tensor).op
return saver_pb2.SaverDef(
filename_tensor_name=filename_tensor.name,
save_tensor_name=save_tensor.name,
restore_op_name=restore_op.name,
version=saver_pb2.SaverDef.V2)
@def_function.function(
input_signature=(tensor_spec.TensorSpec(shape=(), dtype=dtypes.string),),
autograph=False)
def _traced_save(self, file_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
save_op = self.save(file_prefix)
with ops.device("cpu:0"):
with ops.control_dependencies([save_op]):
return array_ops.identity(file_prefix)
@def_function.function(
input_signature=(tensor_spec.TensorSpec(shape=(), dtype=dtypes.string),),
autograph=False)
def _traced_restore(
self, file_prefix: tensor_lib.Tensor) -> tensor_lib.Tensor:
restore_ops = self.restore(file_prefix)
with ops.device("cpu:0"):
with ops.control_dependencies(restore_ops.values()):
return array_ops.identity(file_prefix)
def _get_shards_by_task(
self,
sharding_callback: sharding_util.ShardingCallback
) -> Sequence[tuple[str, Sequence[sharding_util.Shard]]]:
"""Calls the sharding callback with shardable_tensors.
Args:
sharding_callback: ShardingCallback. The callback function wrapper that
splits shardable_tensors into shards.
Returns:
A list of (task, shards) tuples.
"""
def wrap_tensor(shardable_tensor):
tensor_val = shardable_tensor.tensor
tensor_shape = shardable_tensor.shape
save_spec = shardable_tensor._tensor_save_spec # pylint: disable=protected-access
with ops.device(shardable_tensor.device):
save_spec_tensor = save_spec.tensor
if tensor_val is None and save_spec_tensor is None:
# A tensor value of `None` indicates that this SaveableObject gets
# recorded in the object graph, but that no value is saved in the
# checkpoint.
return None
elif save_spec_tensor is not None:
# Pull the tensor value from _tensor_save_spec.
tensor_val = save_spec_tensor
tensor_shape = save_spec_tensor.shape
# Propagate the save spec name and/or slice spec when they are tensors.
# This makes sure properties like `layout` for dtensor names/slice specs
# are preserved during sharding.
if isinstance(save_spec.name, tensor_lib.Tensor):
tensor_val._wrapped_name = save_spec.name # pylint: disable=protected-access
if isinstance(shardable_tensor.slice_spec, tensor_lib.Tensor):
tensor_val._wrapped_slice_spec = save_spec.slice_spec # pylint: disable=protected-access
return dataclasses.replace(
shardable_tensor,
tensor=tensor_val,
shape=tensor_shape)
shardable_tensors_by_task = {
task: [shardable_tensor
for shardable_tensor in map(wrap_tensor, shardable_tensors)
if shardable_tensor is not None]
for task, shardable_tensors in self._shardable_tensors_by_task.items()}
sharding_callback = (
sharding_callback or sharding_policies.ShardByTaskPolicy())
metrics.SetShardingCallbackDescription(
description=sharding_callback.description)
callback_start_time = time.time() * 1e6
shards_by_task = []
for task, shardable_tensors in shardable_tensors_by_task.items():
shards_by_task.append((task, sharding_callback(shardable_tensors)))
callback_end_time = time.time() * 1e6
callback_duration = math.ceil(callback_end_time - callback_start_time)
metrics.AddShardingCallbackDuration(
callback_duration=max(1, callback_duration)) # in microseconds
logging.info("Sharding callback duration: %s microseconds",
callback_duration)
return shards_by_task
def save(
self,
file_prefix: tensor_lib.Tensor,
options: "checkpoint_options.CheckpointOptions | None" = None
) -> ops.Operation:
"""Save the saveable objects to a checkpoint with `file_prefix`.
Args:
file_prefix: A string or scalar string Tensor containing the prefix to
save under.
options: Optional `CheckpointOptions` object.
Returns:
An `Operation`, or None when executing eagerly.
"""
options = options or checkpoint_options.CheckpointOptions()
# IMPLEMENTATION DETAILS: most clients should skip.
#
# Suffix for any well-formed "checkpoint_prefix", when sharded.
# Transformations:
# * Users pass in "save_path" in save() and restore(). Say "myckpt".
# * checkpoint_prefix gets fed <save_path><sharded_suffix>.
#
# Example:
# During runtime, a temporary directory is first created, which contains
# files
#
# <train dir>/myckpt_temp/
# part-?????-of-?????{.index, .data-00000-of-00001}
#
# Before .save() finishes, they will be (hopefully, atomically) renamed to
#
# <train dir>/
# myckpt{.index, .data-?????-of-?????}
#
# Filesystems with eventual consistency (such as S3), don't need a
# temporary location. Using a temporary directory in those cases might
# cause situations where files are not available during copy.
#
# Users only need to interact with the user-specified prefix, which is
# "<train dir>/myckpt" in this case. Save() and Restore() work with the
# prefix directly, instead of any physical pathname. (On failure and
# subsequent restore, an outdated and orphaned temporary directory can be
# safely removed.)
with ops.device("CPU"):
sharded_suffix = array_ops.where(
string_ops.regex_full_match(file_prefix, "^s3://.*"),
constant_op.constant(".part"),
constant_op.constant("_temp/part"))
tmp_checkpoint_prefix = string_ops.string_join(
[file_prefix, sharded_suffix])
registered_paths = {
saver_name: registered_saver_filename(file_prefix, saver_name)
for saver_name in self._registered_savers
}
def save_fn() -> ops.Operation:
saved_prefixes = []
# Save with the registered savers. These run before default savers due to
# the API contract.
for saver_name, (save_fn, _) in self._registered_savers.items():
maybe_saved_prefixes = save_fn(registered_paths[saver_name])
if maybe_saved_prefixes is not None:
flattened_saved_prefixes = nest.flatten(maybe_saved_prefixes)
if not all(
tensor_util.is_tf_type(x) and x.dtype == dtypes.string
for x in flattened_saved_prefixes):
raise ValueError(
"Registered saver must return a (maybe empty) list of "
f"string type tensors. Got {maybe_saved_prefixes}.")
saved_prefixes.extend(flattened_saved_prefixes)
shards_by_task = self._get_shards_by_task(
options.experimental_sharding_callback)
num_shards = sum([len(shards) for _, shards in shards_by_task])
metrics.AddNumCheckpointShardsWritten(num_shards=num_shards)
num_shards_tensor = constant_op.constant(num_shards, name="num_shards")
sharded_saves = []
shard_idx = 0
for task, shards in shards_by_task:
for shard in shards:
with ops.device(task):
shard_prefix = sharded_filename(tmp_checkpoint_prefix, shard_idx,
num_shards_tensor)
shard_idx += 1
saved_prefixes.append(shard_prefix)
sharded_saves.append(
_single_shard_save(shard_prefix, shard, task, options))
with ops.control_dependencies(sharded_saves):
# Merge on the io_device if specified, otherwise co-locates the merge op
# with the last device used.
tensor_device_spec = list(self._shardable_tensors_by_task.keys())[-1]
merge_device_spec = (
options.experimental_io_device or
saveable_object_util.set_cpu0(tensor_device_spec.to_string()))
with ops.device(merge_device_spec):
# V2 format write path consists of a metadata merge step. Once
# merged, attempts to delete the temporary directory,
# "<user-fed prefix>_temp".
return gen_io_ops.merge_v2_checkpoints(
saved_prefixes, file_prefix, delete_old_dirs=True)
# Since this will causes a function re-trace on each save, limit this to the
# cases where it is needed: eager and when there are multiple tasks. Note
# that the retrace is needed to ensure we pickup the latest values of
# options like experimental_io_device.
if context.executing_eagerly() and self._num_unique_tasks > 1:
# Explicitly place the identity op on the first device.
@def_function.function(jit_compile=False)
def tf_function_save() -> None:
save_fn()
tf_function_save()
else:
return save_fn()
def restore(
self,
file_prefix: tensor_lib.Tensor,
options: "checkpoint_options.CheckpointOptions | None" = None
) -> Mapping[str, ops.Operation]:
"""Restore the saveable objects from a checkpoint with `file_prefix`.
Args:
file_prefix: A string or scalar string Tensor containing the prefix for
files to read from.
options: Optional `CheckpointOptions` object.
Returns:
When not run eagerly or when saving on a single device, returns a
dictionary mapping from SaveableObject names to restore operations;
otherwise, returns an empty dict.
"""
options = options or checkpoint_options.CheckpointOptions()
def restore_fn() -> Mapping[str, ops.Operation]:
restore_fn_inputs = {}
restore_fn_input_count = {
fn: len(keys) for fn, keys in self._restore_fn_to_keys.items()}
restore_ops = {}
for task, shard in self._shardable_tensors_by_task.items():
with ops.device(task):
# Load values from checkpoint
restored_tensor_dict = _single_shard_restore(
file_prefix, shard, options)
# Map restored tensors to the corresponding restore_fn, and see if
# all inputs have all been loaded. Call `restore_fn` if that is the
# case.
for ckpt_key, slice_and_tensor in restored_tensor_dict.items():
for slice_spec, tensor in slice_and_tensor.items():
restore_fn = self._keys_to_restore_fn[(ckpt_key,
slice_spec)]
# Processing the returned restored_tensor_dict to prepare for
# the Trackable `restore` function. The `restore` function
# expects a map of `string name (checkpoint_key) -> Tensor`.
# Unless there is a slice_spec, in which case the map will be of
# `string name (checkpoint_key)-> slice_spec -> Tensor`.
if slice_spec:
(restore_fn_inputs.setdefault(restore_fn, {}).setdefault(
ckpt_key, {})[slice_spec]) = tensor
else:
restore_fn_inputs.setdefault(restore_fn,
{})[ckpt_key] = tensor
restore_fn_input_count[restore_fn] -= 1
if restore_fn_input_count[restore_fn] == 0:
restored_tensors = {}
# Extracts the substring after the "/.ATTRIBUTES/" in the
# ckpt_key from restore_fn_inputs[restore_fn] to
# restored_tensors. For example, if
# restore_fn_input[restore_fn] is dict
# { "/.ATTIBUTES/a": Tensor}, restored_tensors will be
# changed to dict {"a": Tensor}
for ckpt_key, tensor in restore_fn_inputs[restore_fn].items():
restored_tensors[trackable_utils.extract_local_name(
ckpt_key)] = tensor
ret = restore_fn(restored_tensors)
if isinstance(ret, dict):
restore_ops.update(ret)
# Run registered restore methods after the default restore ops.
for _, (_, restore_fn) in self._registered_savers.items():
restore_fn(file_prefix)
return restore_ops
has_custom_device_saver = False
for sts in self._shardable_tensors_by_task.values():
if any([context.is_custom_device(st.device.to_string()) for st in sts]):
has_custom_device_saver = True
break
# Since this will cause a function re-trace on each restore, limit this to
# cases where it is needed: eager and when there are multiple tasks or any
# device_spec is a custom device. Note that the retrace is needed to ensure
# we pickup the latest values of options like experimental_io_device.
#
# We run in a function when there is a custom device saver because custom
# devices, such as DTensor, usually do a sharded save and restore.
# Doing a sharded save and restore requires knowledge about what shards
# of variables we are restoring to. In practice, this means that custom
# devices need the AssignVariableOps along with the Restore op within the
# same graph to infer shapes and shard specs for Restore op.
if context.executing_eagerly() and (self._num_unique_tasks > 1 or
has_custom_device_saver):
@def_function.function(jit_compile=False, autograph=False)
def tf_function_restore() -> Mapping[str, ops.Operation]:
restore_fn()
return {}
restore_ops = tf_function_restore()
else:
restore_ops = restore_fn()
return restore_ops
tensorflow/python/checkpoint/functional_saver_test.py
+257
View File
@@ -0,0 +1,257 @@
# 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 the functional saver."""
import os
import time
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.checkpoint import functional_saver
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.eager import context
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.eager import wrap_function
from tensorflow.python.framework import config
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import ops
from tensorflow.python.framework import test_util
from tensorflow.python.module import module
from tensorflow.python.ops import gen_io_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.platform import gfile
from tensorflow.python.training import server_lib
from tensorflow.python.training.saving import saveable_object_util
LOCALHOST = "/job:localhost/replica:0/task:0/device:CPU:0"
class SaverTest(test.TestCase):
def setUp(self):
super(SaverTest, self).setUp()
cpus = config.list_physical_devices("CPU")
# Set 3 virtual CPUs
config.set_logical_device_configuration(cpus[0], [
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration(),
context.LogicalDeviceConfiguration()
])
self.local_options = checkpoint_options.CheckpointOptions(
experimental_io_device=LOCALHOST)
def _get_tensors_by_task(self, root):
serialized_tensors, _, _, _ = (
checkpoint.TrackableSaver(graph_view.ObjectGraphView(root))
._gather_serialized_tensors(None))
tensors_by_task = {}
for tensor_dict in serialized_tensors.values():
for checkpoint_key, maybe_tensor in tensor_dict.items():
if not isinstance(maybe_tensor, dict):
maybe_tensor = {"": maybe_tensor}
for slice_spec, tensor in maybe_tensor.items():
tensor_task = saveable_object_util.set_cpu0(tensor.device)
(tensors_by_task
.setdefault(tensor_task, {})
.setdefault(checkpoint_key, {})[slice_spec]) = tensor
return tensors_by_task
@test_util.run_in_graph_and_eager_modes
def test_resource_variable(self):
v1 = resource_variable_ops.ResourceVariable(2.)
self.evaluate(v1.initializer)
saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v1, "x"))
prefix = os.path.join(self.get_temp_dir(), "ckpt")
self.evaluate(saver.save(constant_op.constant(prefix)))
self.assertEqual(2, len(gfile.Glob(prefix + "*")))
self.evaluate(v1.assign(1.))
self.evaluate(saver.restore(prefix))
self.assertEqual(2., self.evaluate(v1))
v2 = resource_variable_ops.ResourceVariable(3.)
self.evaluate(v2.initializer)
second_saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v2, "x"))
self.evaluate(second_saver.restore(prefix))
self.assertEqual(2., self.evaluate(v2))
@test_util.run_in_graph_and_eager_modes
def test_resource_variable_use_localhost(self):
v1 = resource_variable_ops.ResourceVariable(2.)
self.evaluate(v1.initializer)
saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v1, "x"))
prefix = os.path.join(self.get_temp_dir(), "ckpt")
self.evaluate(saver.save(constant_op.constant(prefix), self.local_options))
self.assertEqual(2, len(gfile.Glob(prefix + "*")))
self.evaluate(v1.assign(1.))
self.evaluate(saver.restore(prefix, self.local_options))
self.assertEqual(2., self.evaluate(v1))
v2 = resource_variable_ops.ResourceVariable(3.)
self.evaluate(v2.initializer)
second_saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v2, "x"))
self.evaluate(second_saver.restore(prefix, self.local_options))
self.assertEqual(2., self.evaluate(v2))
# In graph mode, verify that the save and restore ops were set to run on
# localhost.
if not context.executing_eagerly():
for op in ops.get_default_graph().get_operations():
if op.type in ("SaveV2", "RestoreV2"):
self.assertEqual(LOCALHOST, op.device)
def test_to_proto(self):
v1 = resource_variable_ops.ResourceVariable(2.)
saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v1, "x"))
prefix = os.path.join(self.get_temp_dir(), "ckpt")
proto_accumulator = []
wrapped = wrap_function.wrap_function(
lambda: proto_accumulator.append(saver.to_proto()), signature=())
self.assertEqual(1, len(proto_accumulator))
proto = proto_accumulator[0]
save = wrapped.prune(
feeds=wrapped.graph.get_tensor_by_name(proto.filename_tensor_name),
fetches=wrapped.graph.get_tensor_by_name(proto.save_tensor_name))
restore = wrapped.prune(
feeds=wrapped.graph.get_tensor_by_name(proto.filename_tensor_name),
fetches=wrapped.graph.get_operation_by_name(proto.restore_op_name))
save_path = save(constant_op.constant(prefix))
v1.assign(1.)
restore(constant_op.constant(save_path))
self.assertEqual(2., self.evaluate(v1))
v2 = resource_variable_ops.ResourceVariable(3.)
second_saver = functional_saver.MultiDeviceSaver.from_saveables(
saveable_object_util.saveable_objects_for_op(v2, "x"))
second_saver.restore(save_path)
self.assertEqual(2., self.evaluate(v2))
@test_util.disable_tfrt("b/171765113: server is not supported in TFRT yet.")
def test_checkpoint_is_sharded_by_task(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)
with ops.device("/job:worker/task:0/cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.)
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(1.)
with ops.device("/job:worker/task:2/cpu:0"):
v2 = resource_variable_ops.ResourceVariable(2.)
self.evaluate([v0.initializer, v1.initializer, v2.initializer])
saver = functional_saver.MultiDeviceSaver.from_saveables(
list(saveable_object_util.saveable_objects_for_op(v0, "v0")) +
list(saveable_object_util.saveable_objects_for_op(v1, "v1")) +
list(saveable_object_util.saveable_objects_for_op(v2, "v2")))
prefix = os.path.join(self.get_temp_dir(), "ckpt")
self.evaluate(saver.save(constant_op.constant(prefix)))
self.assertEqual(4, len(gfile.Glob(prefix + "*")))
self.evaluate(v0.assign(-1.))
self.evaluate(v1.assign(-1.))
self.evaluate(v2.assign(-1.))
self.evaluate(saver.restore(constant_op.constant(prefix)))
self.assertEqual(0., self.evaluate(v0))
self.assertEqual(1., self.evaluate(v1))
self.assertEqual(2., self.evaluate(v2))
@test_util.run_in_graph_and_eager_modes
def test_checkpoint_multi_device_using_localhost(self):
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.)
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.)
with ops.device("cpu:2"):
v2 = resource_variable_ops.ResourceVariable(2.)
self.evaluate([v0.initializer, v1.initializer, v2.initializer])
saver = functional_saver.MultiDeviceSaver.from_saveables(
list(saveable_object_util.saveable_objects_for_op(v0, "v0")) +
list(saveable_object_util.saveable_objects_for_op(v1, "v1")) +
list(saveable_object_util.saveable_objects_for_op(v2, "v2")))
prefix = os.path.join(self.get_temp_dir(), "ckpt")
self.evaluate(saver.save(constant_op.constant(prefix), self.local_options))
self.assertEqual(2, len(gfile.Glob(prefix + "*")))
self.evaluate(v0.assign(-1.))
self.evaluate(v1.assign(-1.))
self.evaluate(v2.assign(-1.))
self.evaluate(
saver.restore(constant_op.constant(prefix), self.local_options))
self.assertEqual(0., self.evaluate(v0))
self.assertEqual(1., self.evaluate(v1))
self.assertEqual(2., self.evaluate(v2))
# In graph mode, verify that the save and restore ops were set to run on
# localhost.
if not context.executing_eagerly():
for op in ops.get_default_graph().get_operations():
if op.type in ("SaveV2", "RestoreV2", "MergeV2Checkpoints"):
self.assertEqual(LOCALHOST, op.device)
def test_single_task_save_singlehost_multidevice(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.)
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.)
with ops.device("cpu:2"):
v2 = resource_variable_ops.ResourceVariable(2.)
root.v0 = v0
root.v1 = v1
root.v2 = v2
tensors_by_task = self._get_tensors_by_task(root)
var_names = [
"v0/.ATTRIBUTES/VARIABLE_VALUE",
"v1/.ATTRIBUTES/VARIABLE_VALUE",
"v2/.ATTRIBUTES/VARIABLE_VALUE"
]
vars_numpy = [v0.numpy(), v1.numpy(), v2.numpy()]
tmp_dir = self.get_temp_dir()
for device in ["cpu:0", "cpu:1", "cpu:2"]:
for shard, (_, tensor_slice_dict) in enumerate(
sorted(tensors_by_task.items())[1:]):
with ops.device(device):
shard_prefix = gen_io_ops.sharded_filename(
os.path.join(tmp_dir, str(shard)), shard, 3)
functional_saver._single_task_save(
shard_prefix, tensor_slice_dict)
start_time = time.time()
max_save_time = start_time + 5 # seconds
while not (gfile.ListDirectory(tmp_dir) or time.time() > max_save_time):
pass # eager execution is lovely
self.assertNotEmpty(gfile.ListDirectory(tmp_dir))
with ops.device(device):
restored_dict = functional_saver._single_task_restore(
shard_prefix, tensor_slice_dict)
self.evaluate(restored_dict)
self.assertEqual(
restored_dict[var_names[shard]][""].numpy(),
vars_numpy[shard])
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/checkpoint/graph_view.py
+199
View File
@@ -0,0 +1,199 @@
"""Manages a graph of Trackable objects."""
# 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.
# ==============================================================================
import copy
import weakref
from tensorflow.python.checkpoint import save_util_v1
from tensorflow.python.checkpoint import trackable_view
from tensorflow.python.trackable import base
from tensorflow.python.util.tf_export import tf_export
def _copy_attached_dependencies(attached_dependencies, memo):
"""Deep-copies attached dependencies, reconnecting weak refs to copies."""
if attached_dependencies is None:
return None
copied_dependencies = []
for dependency in attached_dependencies:
if isinstance(dependency, base.WeakTrackableReference):
strong_ref = dependency.ref
if strong_ref is None:
copied_dependencies.append(copy.deepcopy(dependency, memo))
else:
copied_dependencies.append(
base.WeakTrackableReference(
dependency.name, copy.deepcopy(strong_ref, memo)
)
)
else:
copied_dependencies.append(copy.deepcopy(dependency, memo))
if isinstance(attached_dependencies, tuple):
return tuple(copied_dependencies)
return copied_dependencies
@tf_export("__internal__.tracking.ObjectGraphView", v1=[])
class ObjectGraphView(trackable_view.TrackableView):
"""Gathers and serializes an object graph."""
def __init__(self, root, attached_dependencies=None):
"""Configure the graph view.
Args:
root: A `Trackable` object whose variables (including the variables of
dependencies, recursively) should be saved. May be a weak reference.
attached_dependencies: List of dependencies to attach to the root object.
Used when saving a Checkpoint with a defined root object. To avoid
reference cycles, this should use the WeakTrackableReference class.
"""
trackable_view.TrackableView.__init__(self, root)
# ObjectGraphView should never contain a strong reference to root, since it
# may result in a cycle:
# root -> deferred dependencies -> CheckpointPosition
# -> CheckpointRestoreCoordinator -> ObjectGraphView -> root
self._root_ref = (root if isinstance(root, weakref.ref)
else weakref.ref(root))
self._attached_dependencies = attached_dependencies
def __deepcopy__(self, memo):
# By default, weak references are not copied, which leads to surprising
# deepcopy behavior. To fix, we copy the strong root object first, then
# explicitly reconnect _root_ref to that copy instead of relying on
# copy.deepcopy(weakref.ref(...)), which may return the original weakref.
strong_root = self._root_ref()
strong_copy = None
if strong_root is not None:
strong_copy = copy.deepcopy(strong_root, memo)
# super() does not have a __deepcopy__, so we need to re-implement it
copied = super().__new__(type(self))
memo[id(self)] = copied
for key, value in vars(self).items():
if key == "_root_ref":
setattr(copied, key,
weakref.ref(strong_copy) if strong_copy is not None else value)
elif key == "_attached_dependencies":
setattr(copied, key, _copy_attached_dependencies(value, memo))
else:
setattr(copied, key, copy.deepcopy(value, memo))
return copied
def list_children(self, obj, save_type=base.SaveType.CHECKPOINT, **kwargs):
"""Returns list of all child trackables attached to obj.
Args:
obj: A `Trackable` object.
save_type: A string, can be 'savedmodel' or 'checkpoint'.
**kwargs: kwargs to use when retrieving the object's children.
Returns:
List of all children attached to the object.
"""
children = []
for name, ref in super(ObjectGraphView,
self).children(obj, save_type, **kwargs).items():
children.append(base.TrackableReference(name, ref))
# GraphView objects may define children of the root object that are not
# actually attached, e.g. a Checkpoint object's save_counter.
if obj is self.root and self._attached_dependencies:
children.extend(self._attached_dependencies)
return children
def children(self, obj, save_type=base.SaveType.CHECKPOINT, **kwargs):
"""Returns all child trackables attached to obj.
Args:
obj: A `Trackable` object.
save_type: A string, can be 'savedmodel' or 'checkpoint'.
**kwargs: kwargs to use when retrieving the object's children.
Returns:
Dictionary of all children attached to the object with name to trackable.
"""
children = {}
for name, ref in self.list_children(obj, **kwargs):
children[name] = ref
return children
@property
def attached_dependencies(self):
"""Returns list of dependencies that should be saved in the checkpoint.
These dependencies are not tracked by root, but are in the checkpoint.
This is defined when the user creates a Checkpoint with both root and kwargs
set.
Returns:
A list of TrackableReferences.
"""
return self._attached_dependencies
@property
def root(self):
if isinstance(self._root_ref, weakref.ref):
derefed = self._root_ref()
assert derefed is not None
return derefed
else:
return self._root_ref
def breadth_first_traversal(self):
return self._breadth_first_traversal()
def _breadth_first_traversal(self):
"""Find shortest paths to all dependencies of self.root."""
return super(ObjectGraphView, self)._descendants_with_paths()
def serialize_object_graph(self, saveables_cache=None):
"""Determine checkpoint keys for variables and build a serialized graph.
Non-slot variables are keyed based on a shortest path from the root saveable
to the object which owns the variable (i.e. the one which called
`Trackable._add_variable` to create it).
Slot variables are keyed based on a shortest path to the variable being
slotted for, a shortest path to their optimizer, and the slot name.
Args:
saveables_cache: An optional cache storing previously created
SaveableObjects created for each Trackable. Maps Trackables to a
dictionary of attribute names to Trackable.
Returns:
A tuple of (named_variables, object_graph_proto, feed_additions):
named_variables: A dictionary mapping names to variable objects.
object_graph_proto: A TrackableObjectGraph protocol buffer
containing the serialized object graph and variable references.
feed_additions: A dictionary mapping from Tensors to values which should
be fed when saving.
Raises:
ValueError: If there are invalid characters in an optimizer's slot names.
"""
named_saveable_objects, object_graph_proto, feed_additions, _ = (
save_util_v1.serialize_object_graph_with_registered_savers(
self, saveables_cache))
return named_saveable_objects, object_graph_proto, feed_additions
def frozen_saveable_objects(self,
object_map=None,
to_graph=None,
call_with_mapped_captures=None):
"""Creates SaveableObjects with the current object graph frozen."""
return save_util_v1.frozen_saveables_and_savers(
self, object_map, to_graph, call_with_mapped_captures)[0]
tensorflow/python/checkpoint/restore.py
+838
View File
@@ -0,0 +1,838 @@
# Copyright 2022 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.
# ==============================================================================
"""Logic for restoring checkpointed values for Trackables."""
import collections
from typing import Optional, Mapping, Any
from tensorflow.python.checkpoint import checkpoint_adapter
from tensorflow.python.checkpoint import checkpoint_view
from tensorflow.python.checkpoint import functional_saver
from tensorflow.python.checkpoint import save_util_v1
from tensorflow.python.checkpoint import saveable_compat
from tensorflow.python.eager import context
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_io_ops as io_ops
from tensorflow.python.ops import io_ops
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.saved_model import registration
from tensorflow.python.trackable import base
from tensorflow.python.trackable import constants
from tensorflow.python.trackable import python_state
from tensorflow.python.trackable import trackable_utils
from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util import object_identity
class CheckpointPosition(object):
"""Indicates a position within a `_CheckpointRestoreCoordinator`."""
__slots__ = ["_checkpoint", "_proto_id", "skip_restore", "callback"]
def __init__(self, checkpoint, proto_id):
"""Specify an object within a checkpoint.
Args:
checkpoint: A _CheckpointRestoreCoordinator object.
proto_id: The index of this object in TrackableObjectGraph.nodes.
"""
self._checkpoint = checkpoint
self._proto_id = proto_id
# This may be set to True if the registered saver cannot be used with this
# object.
self.skip_restore = False
self.callback = checkpoint_adapter.ReshardCallback()
def restore(self, trackable, reader=None):
"""Restore this value into `trackable`."""
with ops.init_scope():
if self.bind_object(trackable):
# This object's correspondence with a checkpointed object is new, so
# process deferred restorations for it and its dependencies.
restore_ops = self._restore_descendants(reader)
if restore_ops:
self._checkpoint.new_restore_ops(restore_ops)
def bind_object(self, trackable):
"""Set a checkpoint<->object correspondence.
Args:
trackable: The object to record a correspondence for.
Returns:
True if this is a new assignment, False if this object has already been
mapped to a checkpointed `Object` proto.
Raises:
AssertionError: If another object is already bound to the `Object` proto.
"""
checkpoint = self.checkpoint
checkpoint.all_python_objects.add(trackable)
current_assignment = checkpoint.object_by_proto_id.get(self._proto_id, None)
checkpoint.matched_proto_ids.add(self._proto_id)
if current_assignment is None:
checkpoint.object_by_proto_id[self._proto_id] = trackable
return True # New assignment
else:
# The object was already mapped for this checkpoint load, which means
# we don't need to do anything besides check that the mapping is
# consistent (if the dependency DAG is not a tree then there are
# multiple paths to the same object).
if current_assignment is not trackable:
logging.warning(
"Inconsistent references when loading the checkpoint into this "
"object graph. For example, in the saved checkpoint object, "
"`model.layer.weight` and `model.layer_copy.weight` reference the "
"same variable, while in the current object these are two different"
" variables. The referenced variables are:"
f"({current_assignment} and {trackable}).")
return False # Not a new assignment
def update_resharding_callback(
self, callback: checkpoint_adapter.ReshardCallback
):
"""Add a resharding callback to the checkpoint.
This will be applied to the checkpoint value before being supplied to the
restore ops.
Args:
callback: Reshard callback for resharding this checkpoint position. Maybe
None.
"""
if not issubclass(checkpoint_adapter.ReshardCallback, type(self.callback)):
raise TypeError(
"Cannot override resharding callback, already set to non trivial."
)
self.callback = callback
def has_non_trivial_reshard_callback(self) -> bool:
"""Determine whether this value has a non-trivial resharding callback."""
return not issubclass(
checkpoint_adapter.ReshardCallback, type(self.callback)
)
def is_simple_variable(self) -> bool:
"""Determine whether this value is restorable with a Tensor initializer."""
attributes = self.object_proto.attributes
return (
len(attributes) == 1
and attributes[0].name == constants.VARIABLE_VALUE_KEY
and not self.object_proto.children
)
def value_tensors(
self, shape_and_slices: Optional[str] = None
) -> Mapping[str, tensor.Tensor]:
"""Create value `Tensor`s for this object's attributes.
Does not require that the Python object has been created. Used for
restore-on-create when executing eagerly.
Args:
shape_and_slices: A dict mapping from object attribute names to a shape
and slice string that will be passed to a RestoreV2 op. If the dict is
None or if an object attribute is not in the dict, the full tensor will
be restored.
Returns:
A dictionary mapping from object attribute names to `Tensor`s.
"""
value_tensors = {}
for serialized_tensor in self.object_proto.attributes:
checkpoint_key = serialized_tensor.checkpoint_key
io_device = self._checkpoint.options.experimental_io_device or "cpu:0"
with ops.init_scope():
with ops.device(io_device):
# Run the restore itself on the io_device(CPU or specified).
if (
shape_and_slices is not None
and serialized_tensor.name in shape_and_slices
):
shape_and_slice = shape_and_slices[serialized_tensor.name]
else:
shape_and_slice = ""
checkpoint_keys, full_shape_and_slices = (
self.callback.update_restore_inputs(
checkpoint_key, shape_and_slice
)
)
dtypes = []
for key in checkpoint_keys:
dtype = self._checkpoint.dtype_map[key]
dtypes.append(dtype.base_dtype)
restored_values = io_ops.restore_v2(
prefix=self._checkpoint.save_path_tensor,
tensor_names=checkpoint_keys,
shape_and_slices=full_shape_and_slices,
dtypes=dtypes,
name="%s_checkpoint_read" % (serialized_tensor.name,),
)
value = self.callback.reshard(
restored_values, shape_and_slice
)
# Copy the value to the current device if necessary.
value_tensors[serialized_tensor.name] = array_ops.identity(value)
return value_tensors
def gather_ops_or_named_saveables(self):
"""Looks up or creates SaveableObjects which don't have cached ops.
Returns:
A tuple of (
existing_restore_ops: list,
named_saveables: dict,
python_positions: list,
registered_savers: dict)
"""
recorded_registered_saver = self.get_registered_saver_name()
if not (self.object_proto.attributes or recorded_registered_saver):
return [], {}, [], {}
existing_restore_ops = []
named_saveables = {}
python_positions = []
registered_savers = collections.defaultdict(dict)
saveable_factories = saveable_object_util.saveable_objects_from_trackable(
self.trackable)
saver_name = registration.get_registered_saver_name(self.trackable)
if recorded_registered_saver:
if not self.skip_restore:
name = self.object_proto.registered_saver.object_name
registered_savers[recorded_registered_saver][name] = self.trackable
# Else: Skip restoration of this Trackable. This skip only happens if the
# registered saver has enabled `option_restore`. Otherwise, an error would
# have been raised at `self.get_registered_saver_name()`.
elif saver_name:
# In this case, the checkpoint has a recorded serialized tensor but no
# registered saver, while the Trackable loading the checkpoint has
# migrated to the registered checkpoint functionality (TPUEmbedding is an
# example of this).
# Set the Trackable's object name to the first checkpoint key that is
# stored in checkpoint. If there is a use case that requires the other
# keys, then we can take another look at this.
registered_savers[saver_name] = {
self.object_proto.attributes[0].checkpoint_key: self.trackable
}
elif isinstance(self.trackable, python_state.PythonState):
python_positions.append(self)
elif saveable_factories.keys() == {
trackable_utils.SERIALIZE_TO_TENSORS_NAME
}:
existing_restore_ops, named_saveables = (
self._create_serialize_to_tensor_saveable(saveable_factories))
elif saveable_factories:
existing_restore_ops, named_saveables = (
self._create_saveables_by_attribute_name(saveable_factories))
else:
# If no registered savers were found, then it means that one or more
# serialized tensors were never used.
for serialized_tensor in self.object_proto.attributes:
self._checkpoint.unused_attributes.setdefault(
self._proto_id, []).append(serialized_tensor.name)
return (existing_restore_ops, named_saveables, python_positions,
registered_savers)
def _create_serialize_to_tensor_saveable(self, saveable_factories):
"""Creates a saveable using the _serialize_to_tensor method."""
# Extract the saveable name from the checkpoint key. This will be used as
# the cache key or the name to pass to the saveable factory.
suffix = saveable_compat.get_saveable_name(self.trackable) or ""
saveable_name = _extract_saveable_name(
self.object_proto.attributes[0].checkpoint_key) + suffix
# Try to find the cached saveable (only in graph mode).
if not context.executing_eagerly():
existing_op = self._checkpoint.restore_ops_by_name.get(
saveable_name, None)
if existing_op is not None:
return [existing_op], {}
saveables_cache = self._checkpoint.saveables_cache.setdefault(
self.trackable, {})
if saveable_name in saveables_cache:
return [], {saveable_name: saveables_cache[saveable_name]}
saveable = saveable_factories[trackable_utils.SERIALIZE_TO_TENSORS_NAME](
name=saveable_name)
if not context.executing_eagerly():
saveables_cache[saveable_name] = saveable
return [], {saveable_name: saveable}
def _create_saveables_by_attribute_name(self, saveable_factories):
"""Creates or caches SaveableObjects by matching the attribute names.
The attribute name keys in the `saveable_factories` is used to find the
corresponding attribute in the object proto. Attributes contain checkpoint
keys which are passed to the factory function to generate the
SaveableObject.
Args:
saveable_factories: a dict mapping attribute name to a callable factory
function that produces a SaveableObject.
Returns:
A tuple of (
existing_restore_ops: list,
named_saveables: dict)
"""
# Name saveables based on the name this object had when it was checkpointed.
named_saveables = {}
existing_restore_ops = []
# Forward compatibility code: when loading a future checkpoint, there may
# be multiple SerializedTensors mapped to a single saveable.
created_compat_names = set()
for serialized_tensor in self.object_proto.attributes:
if context.executing_eagerly():
existing_op = None
else:
existing_op = self._checkpoint.restore_ops_by_name.get(
serialized_tensor.checkpoint_key, None)
if existing_op is not None:
existing_restore_ops.append(existing_op)
continue
if any(serialized_tensor.name.startswith(name)
for name in created_compat_names):
continue # Saveable has already been created for this tensor.
# Only if we don't have cached ops for this SaveableObject, we'll see if
# the SaveableObject itself has been cached. If not, we'll make it, and
# either way we'll extract new ops from it (or if it has Python state to
# restore, we'll run that).
saveables_cache = self._checkpoint.saveables_cache
if saveables_cache is None:
# No SaveableObject caching when executing eagerly.
saveable = None
else:
# If we've already created and cached a SaveableObject for this
# attribute, we can re-use it to avoid re-creating some ops when graph
# building.
saveable_list = saveables_cache.get(self.trackable,
{}).get(serialized_tensor.name,
(None,))
if len(saveable_list) == 1:
# Almost every attribute will have exactly one SaveableObject.
saveable, = saveable_list
else:
# Don't use cached SaveableObjects for partitioned variables, which is
# the only case where we'd have a list of SaveableObjects. Op caching
# will catch them.
saveable = None
if saveable is not None:
# The name of this attribute has changed, so we need to re-generate
# the SaveableObject.
if serialized_tensor.checkpoint_key not in saveable.name:
saveable = None
del saveables_cache[self.trackable]
if saveable is None:
# If there was no cached SaveableObject, create one.
# Use the name to check if the Python object has the same attribute.
saveable = _get_saveable_from_factory(saveable_factories,
serialized_tensor,
created_compat_names)
if saveable is None:
# Purposefully does not throw an exception if attributes have been
# added or deleted. Stores unused attributes so an exception can be
# raised if the user decides to check that everything in the
# checkpoint was loaded.
self._checkpoint.unused_attributes.setdefault(
self._proto_id, []).append(serialized_tensor.name)
continue
if saveables_cache is not None:
saveables_cache.setdefault(self.trackable,
{})[serialized_tensor.name] = [saveable]
named_saveables[serialized_tensor.checkpoint_key] = saveable
return existing_restore_ops, named_saveables
def restore_ops(self, reader=None):
"""Create or fetch restore ops for this object's attributes.
Requires that the `Trackable` Python object has been bound to an object
ID in the checkpoint.
Args:
reader: A `CheckpointReader`. If None, a new instance will be created.
Returns:
A list of operations when graph building, or an empty list when executing
eagerly.
"""
if self._has_registered_saver():
raise ValueError("Unable to run individual checkpoint restore for objects"
" with registered savers.")
(restore_ops, tensor_saveables, python_positions,
_) = self.gather_ops_or_named_saveables()
restore_ops.extend(
self._checkpoint.restore_saveables(
tensor_saveables, python_positions, reader=reader))
return restore_ops
@property
def checkpoint(self):
return self._checkpoint
@property
def trackable(self):
return self._checkpoint.object_by_proto_id[self._proto_id]
@property
def object_proto(self):
return self._checkpoint.object_graph_proto.nodes[self._proto_id]
@property
def proto_id(self):
return self._proto_id
@property
def restore_uid(self):
return self._checkpoint.restore_uid
def __repr__(self):
return repr(self.object_proto)
def value_shape(self):
"""The shape of the VARIABLE_VALUE tensor.
Returns:
If found a TensorShape object, otherwise None.
"""
for serialized_tensor in self.object_proto.attributes:
if serialized_tensor.name == constants.VARIABLE_VALUE_KEY:
return self._checkpoint.shape_map[serialized_tensor.checkpoint_key]
return None
def _has_registered_saver(self):
return bool(self.object_proto.registered_saver.name)
def get_registered_saver_name(self):
"""Returns the registered saver name defined in the Checkpoint."""
if self._has_registered_saver():
saver_name = self.object_proto.registered_saver.name
try:
registration.validate_restore_function(self.trackable, saver_name)
except ValueError as e:
if registration.get_strict_predicate_restore(saver_name):
raise e
self.skip_restore = True
return saver_name
return None
def create_slot_variable_position(
self,
optimizer_object: Any,
variable: base.Trackable,
slot_variable_id: str,
slot_name: str,
reshard_callback: Optional[checkpoint_adapter.ReshardCallback] = None,
):
"""Generates CheckpointPosition for a slot variable.
Args:
optimizer_object: Optimizer that owns the slot variable.
variable: Variable associated with the slot variable.
slot_variable_id: ID of the slot variable.
slot_name: Name of the slot variable.
reshard_callback: A callback object for resharding value from checkpoint
at restore.
Returns:
If there is a slot variable in the `optimizer_object` that has not been
bound to the checkpoint, this function returns a tuple of (
new `CheckpointPosition` for the slot variable,
the slot variable itself).
"""
slot_variable_position = CheckpointPosition(
checkpoint=self.checkpoint, proto_id=slot_variable_id
)
# pylint: disable=protected-access
if reshard_callback is not None:
# slot_variable_shape kwarg is available only for optimizer_v2 objects.
slot_variable_position.update_resharding_callback(reshard_callback)
slot_variable = optimizer_object._create_or_restore_slot_variable(
slot_variable_position=slot_variable_position,
variable=variable,
slot_name=slot_name,
slot_variable_shape=variable.shape,
)
else:
slot_variable = optimizer_object._create_or_restore_slot_variable(
slot_variable_position=slot_variable_position,
variable=variable,
slot_name=slot_name,
)
# pylint: enable=protected-access
if slot_variable is not None and slot_variable_position.bind_object(
slot_variable
):
return slot_variable_position, slot_variable
else:
return None, None
def create_child_position(self, node_id):
return CheckpointPosition(checkpoint=self.checkpoint, proto_id=node_id)
def _restore_descendants(self, reader=None):
"""Restore the bound Trackable and dependencies (may be deferred)."""
# Attempt a breadth-first traversal, since presumably the user has more
# control over shorter paths. If we don't have all of the dependencies at
# this point, the end result is not breadth-first (since other deferred
# traversals will happen later).
# You may be wondering why elements in the `visit_queue` are tuples that
# contains both CheckpointPositions and their Trackable. The reason is that
# Optimizers will not keep a strong reference to slot vars for
# ShardedVariables. The slot variable must be kept in memory until the
# restore saveables have been created.
visit_queue = collections.deque([(self, self.trackable)])
restore_ops = []
tensor_saveables = {}
python_positions = []
registered_savers = collections.defaultdict(dict)
while visit_queue:
current_position, _ = visit_queue.popleft()
# Restore using the ops defined in a Saveable or registered function.
(
new_restore_ops,
new_tensor_saveables,
new_python_positions,
new_registered_savers,
) = current_position._single_restore() # pylint: disable=protected-access
restore_ops.extend(new_restore_ops)
tensor_saveables.update(new_tensor_saveables)
python_positions.extend(new_python_positions)
for saver_name, trackable_map in new_registered_savers.items():
registered_savers[saver_name].update(trackable_map)
# Pass the restoration to the dependencies.
_queue_children_for_restoration(current_position, visit_queue)
_queue_slot_variables(current_position, visit_queue)
restore_ops.extend(
current_position.checkpoint.restore_saveables(
tensor_saveables, python_positions, registered_savers, reader=reader
)
)
return restore_ops
def _single_restore(self):
"""Restores the trackable."""
trackable = self.trackable
trackable._maybe_initialize_trackable() # pylint: disable=protected-access
checkpoint = self.checkpoint
# If the UID of this restore is lower than our current update UID, we don't
# need to actually restore the object.
if checkpoint.restore_uid > trackable._update_uid: # pylint: disable=protected-access
restore_ops, tensor_saveables, python_positions, registered_savers = (
self.gather_ops_or_named_saveables()
)
trackable._update_uid = checkpoint.restore_uid # pylint: disable=protected-access
else:
restore_ops = ()
tensor_saveables = {}
python_positions = ()
registered_savers = {}
return restore_ops, tensor_saveables, python_positions, registered_savers
def restore_nodes(save_path, nodes_to_restore):
"""Restores nodes from a dict.
Requires that the `Trackable` Python object has been bound to an object
ID in the checkpoint.
Args:
save_path: a string represents path to the checkpoint.
nodes_to_restore: a dict maps `node_id` to `trackable` to be restored.
"""
if save_path is None:
raise ValueError("save_path cannot be empty.")
if not isinstance(nodes_to_restore, dict):
raise ValueError(
"Expecting a dictionary of node_id to Trackable for nodes_to_restore.")
ckpt_view = checkpoint_view.CheckpointView(save_path)
ckpt_view_descendants = ckpt_view.descendants()
for node_id, trackable in nodes_to_restore.items():
# node_id does not have a corresponding Checkpoint value.
if (node_id not in ckpt_view_descendants or
ckpt_view._object_graph_proto.nodes[ # pylint: disable=protected-access
node_id] is None):
raise ValueError(
f"The expected node_id: {node_id} to Trackable {trackable} to "
"restore does not exist in the checkpoint.")
# Trackable mapped to node_id to restore is empty.
if trackable is None or not isinstance(trackable, base.Trackable):
raise ValueError(
f"Expecting a valid Trackable to node_id: {node_id} but got "
f"trackable: {trackable}."
)
serialized_tensors = object_identity.ObjectIdentityDictionary()
for node_id, current_trackable in nodes_to_restore.items():
ckpt_contains_serialized_tensors = ckpt_view._object_graph_proto.nodes[ # pylint: disable=protected-access
node_id].attributes
node = ckpt_view._object_graph_proto.nodes[node_id] # pylint: disable=protected-access
trackable_has_serialize_to_tensor = (
saveable_object_util.trackable_has_serialize_to_tensor(
current_trackable
)
)
if not trackable_has_serialize_to_tensor:
if not node.attributes:
if saveable_object_util.saveable_objects_from_trackable(
current_trackable):
raise ValueError(
f"Trackable {current_trackable} expects checkpointed values but "
"checkpoint does not contain serialized tensors for node_id: "
f"{node_id}.")
else:
continue
object_names = object_identity.ObjectIdentityDictionary()
object_names[current_trackable] = trackable_utils.extract_object_name(
node.attributes[0].checkpoint_key)
checkpoint_factory_map, _ = (
save_util_v1.get_checkpoint_factories_and_keys(object_names, None)
)
saveable_objects = save_util_v1.generate_saveable_objects(
checkpoint_factory_map)[0]
if len(node.attributes) != len(saveable_objects):
raise ValueError("Size for saveable_objects for Trackable: "
f"{len(saveable_objects)} did not match the size for "
"serialized_tensors for checkpoint: "
f"{len(node.attributes)}.")
current_trackable = saveable_object_util.SaveableCompatibilityConverter(
current_trackable, saveable_objects)
serialized_tensors[
current_trackable] = current_trackable._serialize_to_tensors() # pylint: disable=protected-access
trackable_expects_ckpted_value = bool(serialized_tensors[current_trackable])
if trackable_expects_ckpted_value and not ckpt_contains_serialized_tensors:
raise ValueError(
f"Trackable {current_trackable} expects checkpointed values but "
"checkpoint does not contain serialized tensors for node_id: "
f"{node_id}.")
if not trackable_expects_ckpted_value and ckpt_contains_serialized_tensors:
raise ValueError(
f"Trackable {current_trackable} does not expect checkpointed "
"values but checkpoint contains serialized tensors: "
f"{ckpt_contains_serialized_tensors} for node_id: {node_id}.")
if len(node.attributes) != len(serialized_tensors[current_trackable]):
raise ValueError("Size for serialized_tensors for Trackable: "
f"{len(serialized_tensors[current_trackable])} did not "
"match size for serialized_tensors for checkpoint: "
f"{len(node.attributes)}.")
if not trackable_has_serialize_to_tensor:
functional_saver.MultiDeviceSaver(serialized_tensors).restore(save_path)
else:
# Converts attribute.name to attribute.checkpoint_key since that's what
# restore method is expecting. i.e., converts "a" to "/.ATTRIBUTES/a".
serialized_tensors_renamed = object_identity.ObjectIdentityDictionary()
serialized_tensors_renamed[current_trackable] = {}
for attribute in node.attributes:
name = attribute.name
checkpoint_key = attribute.checkpoint_key
serialized_tensors_renamed[current_trackable][
checkpoint_key] = serialized_tensors[current_trackable][name]
functional_saver.MultiDeviceSaver(serialized_tensors_renamed).restore(
save_path)
def _maybe_get_adapter(checkpoint_position, trackable):
adapter = trackable._checkpoint_adapter( # pylint: disable=protected-access
checkpoint_position.checkpoint.save_path_string
)
if adapter and adapter.is_applicable(trackable):
return adapter
return None
def _queue_children_for_restoration(checkpoint_position, visit_queue):
"""Queues the restoration of trackable's children or defers them."""
# pylint: disable=protected-access
trackable = checkpoint_position.trackable
trackable_children = trackable._trackable_children()
adapter = _maybe_get_adapter(checkpoint_position, trackable)
for child in checkpoint_position.object_proto.children:
# trackable._lookup_dependency can be expensive so first check if this node
# already has an object correspondence. If so we skip this node.
correspondence = checkpoint_position.checkpoint.object_by_proto_id.get(
child.node_id, None
)
if correspondence is not None:
continue
child_position = checkpoint_position.create_child_position(child.node_id)
local_object = trackable._lookup_dependency(child.local_name,
trackable_children)
child_proto = child_position.object_proto
if local_object is None:
# We don't yet have a dependency registered with this name. Save it
# in case we do.
if child_proto.HasField("has_checkpoint_values"):
has_value = child_proto.has_checkpoint_values.value
else:
# If the field is not set, do a simple check to see if the dependency
# has children and/or checkpointed values.
has_value = bool(
child_proto.children
or child_proto.attributes
or child_proto.slot_variables
or child_proto.HasField("registered_saver")
)
if has_value:
local_trackable_name = child.local_name
if adapter:
local_trackable_name, reshard_callback = adapter.maybe_reshard(
child.local_name
)
if reshard_callback:
child_position.update_resharding_callback(reshard_callback)
trackable._deferred_dependencies.setdefault(
local_trackable_name, []
).append(child_position)
else:
if child_position.bind_object(trackable=local_object):
# This object's correspondence is new, so dependencies need to be
# visited. Delay doing it so that we get a breadth-first dependency
# resolution order (shallowest paths first). The caller is responsible
# for emptying visit_queue.
visit_queue.append((child_position, local_object))
_DeferredSlotVariableRestoration = collections.namedtuple(
"_DeferredSlotVariableRestoration", [
"original_variable",
"slot_variable_id",
"slot_name",
])
def _queue_slot_variables(checkpoint_position, visit_queue):
"""Queues slot variables for restoration."""
trackable = checkpoint_position.trackable
checkpoint = checkpoint_position.checkpoint
for deferred_slot_restoration in (checkpoint.deferred_slot_restorations.pop(
checkpoint_position.proto_id, ())):
slot_variable_position, slot_variable = (
checkpoint_position.create_slot_variable_position(
trackable,
deferred_slot_restoration.original_variable,
deferred_slot_restoration.slot_variable_id,
deferred_slot_restoration.slot_name,
# If the corresponding variable has a non trivial resharding
# attached, the slot variable should be resharded in the same
# way.
checkpoint_position.callback
if checkpoint_position.has_non_trivial_reshard_callback()
else None,
)
)
if slot_variable_position is not None:
visit_queue.append((slot_variable_position, slot_variable))
for slot_restoration in checkpoint.slot_restorations.pop(
checkpoint_position.proto_id, ()):
optimizer_object = checkpoint.object_by_proto_id.get(
slot_restoration.optimizer_id, None)
if optimizer_object is None:
# The optimizer has not yet been created or tracked. Record in the
# checkpoint that the slot variables need to be restored when it is.
checkpoint.deferred_slot_restorations.setdefault(
slot_restoration.optimizer_id, []).append(
_DeferredSlotVariableRestoration(
original_variable=trackable,
slot_variable_id=slot_restoration.slot_variable_id,
slot_name=slot_restoration.slot_name))
# `optimizer_object` can be a `Checkpoint` when user only needs the
# attributes the optimizer holds, such as `iterations`. In those cases,
# it would not have the optimizer's `_create_or_restore_slot_variable`
# method.
elif hasattr(optimizer_object, "_create_or_restore_slot_variable"):
slot_variable_position, slot_variable = (
checkpoint_position.create_slot_variable_position(
optimizer_object,
trackable,
slot_restoration.slot_variable_id,
slot_restoration.slot_name,
# If the corresponding variable has a non trivial resharding
# attached, the slot variable should be resharded in the same
# way.
checkpoint_position.callback
if checkpoint_position.has_non_trivial_reshard_callback()
else None,
)
)
if slot_variable_position is not None:
visit_queue.append((slot_variable_position, slot_variable))
def _extract_saveable_name(checkpoint_key):
# Substring the checkpoint key to the end of the "{...}.ATTRIBUTES/"
search_key = trackable_utils.OBJECT_ATTRIBUTES_NAME + "/"
return checkpoint_key[:checkpoint_key.index(search_key) + len(search_key)]
def _get_saveable_from_factory(saveable_factories, serialized_tensor,
created_compat_names):
"""Returns the saveable generated from the factory method."""
matched_factory = None
# The `expected_factory_name` is used to find the right saveable factory,
# while the `factory_input_name` is the value that is passed to the factory
# method to instantiate the SaveableObject.
expected_factory_name = serialized_tensor.name
factory_input_name = serialized_tensor.checkpoint_key
# Case 1: the name already exactly matches a key in saveable_factories.
if expected_factory_name in saveable_factories:
matched_factory = saveable_factories[expected_factory_name]
# Case 2: (Forward compat) The serialized name is composed of
# "factory_name" + "SUFFIX". Get the matching factory name.
if matched_factory is None:
for factory_name, factory in saveable_factories.items():
if expected_factory_name.startswith(factory_name):
if matched_factory is not None:
# This condition is met in the extreme edge case where the object
# returns two saveable factories with similar names. This is very
# unlikely because there zero objects inside TensorFlow that use
# more than one saveable factory.
raise ValueError("Forward compatibility load error: Unable to load "
"checkpoint saved in future version of TensorFlow. "
"Please update your version of TensorFlow to the "
"version in which the checkpoint was saved.")
matched_factory = factory
factory_input_name = _extract_saveable_name(
serialized_tensor.checkpoint_key) + factory_name
created_compat_names.add(factory_name)
if callable(matched_factory):
return matched_factory(name=factory_input_name)
return matched_factory
tensorflow/python/checkpoint/restore_test.py
+309
View File
@@ -0,0 +1,309 @@
# Copyright 2022 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 restore.py."""
import os
from tensorflow.python.checkpoint import checkpoint as trackable_utils
from tensorflow.python.checkpoint import restore
from tensorflow.python.eager import test
from tensorflow.python.module import module
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import autotrackable
from tensorflow.python.trackable import base
from tensorflow.python.training.saving import saveable_object
class ExtractSaveablenameTest(test.TestCase):
def test_standard_saveable_name(self):
self.assertEqual(
"object_path/.ATTRIBUTES/",
restore._extract_saveable_name("object_path/.ATTRIBUTES/123"))
self.assertEqual(
"object/path/ATTRIBUTES/.ATTRIBUTES/",
restore._extract_saveable_name("object/path/ATTRIBUTES/.ATTRIBUTES/"))
def test_restore_nodes_error_cases_high_level(self):
root = autotrackable.AutoTrackable()
root.leaf = autotrackable.AutoTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = autotrackable.AutoTrackable()
root2.leaf = autotrackable.AutoTrackable()
with self.assertRaisesRegex(
ValueError,
"Expecting a dictionary of node_id to Trackable for nodes_to_restore."):
restore.restore_nodes(root_save_path, [0, 1])
with self.assertRaisesRegex(
ValueError,
"The expected node_id: 3 to Trackable <.*?> to restore does not exist "
"in the checkpoint."):
restore.restore_nodes(root_save_path, {3: root2})
with self.assertRaisesRegex(
ValueError,
"Expecting a valid Trackable to node_id: 0 but got trackable: None."):
restore.restore_nodes(root_save_path, {0: None})
def test_restore_nodes_error_cases_trackable_ckpt_view_mismatch(self):
class MyTrackable(base.Trackable):
def __init__(self):
self.a = module.Module()
class MyTrackable2(base.Trackable):
def __init__(self):
self.a = variables.Variable(5.0)
def _serialize_to_tensors(self):
return {"a": variables.Variable(5.0)}
root = MyTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackable2()
with self.assertRaisesRegex(
ValueError,
"Trackable <.*?> expects checkpointed values but checkpoint does not "
"contain serialized tensors for node_id: 0."):
restore.restore_nodes(root_save_path, {0: root2})
def test_restore_nodes_has_serialize_to_tensor(self):
class MyTrackable(base.Trackable):
def __init__(self):
self.a = variables.Variable(5.0)
def _restore_from_tensors(self, restored_tensors):
return self.a.assign(restored_tensors["a"])
def _serialize_to_tensors(self):
return {"a": self.a}
root = MyTrackable()
leaf = MyTrackable()
root._track_trackable(leaf, "leaf")
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackable()
leaf2 = MyTrackable()
root2._track_trackable(leaf2, "leaf")
root2.a.assign(3.0)
# Restore root
restore.restore_nodes(root_save_path, {0: root2})
self.assertEqual(root2.a.numpy(), 5.0) # Restored from 3.0 to 5.0
self.assertEqual(leaf2.a.numpy(), 5.0) # Unchanged
root3 = MyTrackable()
leaf3 = MyTrackable()
root3._track_trackable(leaf3, "leaf")
leaf3.a.assign(3.0)
# Restore leaf
restore.restore_nodes(root_save_path, {1: leaf3})
self.assertEqual(root3.a.numpy(), 5.0) # Unchanged
self.assertEqual(leaf3.a.numpy(), 5.0) # Restored from 3.0 to 5.0.
def test_restore_nodes_with_different_number_of_serialized_to_tensors(self):
class MyTrackableA(base.Trackable):
def __init__(self):
self.a = variables.Variable(5.0)
def _restore_from_tensors(self, restored_tensors):
return self.a.assign(restored_tensors["a"])
def _serialize_to_tensors(self):
return {"a": self.a}
class MyTrackableAandB(base.Trackable):
def __init__(self):
self.a = variables.Variable(5.0)
self.b = variables.Variable(6.0)
def _restore_from_tensors(self, restored_tensors):
return control_flow_ops.group(
self.a.assign(restored_tensors["a"]),
self.b.assign(restored_tensors["b"])
)
def _serialize_to_tensors(self):
return {"a": self.a, "b": self.b}
root = MyTrackableA()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackableAandB()
with self.assertRaisesRegex(
ValueError,
"Size for serialized_tensors for Trackable: 2 did not match size for "
"serialized_tensors for checkpoint: 1."):
restore.restore_nodes(root_save_path, {0: root2})
root = MyTrackableAandB()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackableA()
with self.assertRaisesRegex(
ValueError,
"Size for serialized_tensors for Trackable: 1 did not match size for "
"serialized_tensors for checkpoint: 2."):
restore.restore_nodes(root_save_path, {0: root2})
def test_restore_nodes_not_serialize_to_tensor(self):
class _VarSaveable(saveable_object.SaveableObject):
def __init__(self, obj, name):
self.obj = obj
specs = [saveable_object.SaveSpec(obj.a, "", name + "-a")]
super(_VarSaveable, self).__init__(None, specs, name)
def restore(self, restored_tensors, restored_shapes):
del restored_shapes # Unused.
self.obj.a.assign(restored_tensors[0])
class MyTrackable(base.Trackable):
def __init__(self):
self.a = variables.Variable(5.0)
def _gather_saveables_for_checkpoint(self):
return {"a": lambda name: _VarSaveable(self, name)}
root = MyTrackable()
leaf = MyTrackable()
root._track_trackable(leaf, "leaf")
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackable()
leaf2 = MyTrackable()
root2._track_trackable(leaf2, "leaf")
root2.a.assign(3.0)
# Restore root
restore.restore_nodes(root_save_path, {0: root2})
self.assertEqual(root2.a.numpy(), 5.0) # Restored from 3.0 to 5.0
self.assertEqual(leaf2.a.numpy(), 5.0) # Unchanged
root3 = MyTrackable()
leaf3 = MyTrackable()
root3._track_trackable(leaf3, "leaf")
leaf3.a.assign(3.0)
# Restore leaf
restore.restore_nodes(root_save_path, {1: leaf3})
self.assertEqual(root3.a.numpy(), 5.0) # Unchanged
self.assertEqual(leaf3.a.numpy(), 5.0) # Restored from 3.0 to 5.0.
def test_restore_nodes_not_serialize_to_tensor_error_cases(self):
class _VarSaveable(saveable_object.SaveableObject):
def __init__(self, obj, name):
self.obj = obj
specs = [saveable_object.SaveSpec(obj.a, "", name + "-a")]
super(_VarSaveable, self).__init__(None, specs, name)
def restore(self, restored_tensors, restored_shapes):
del restored_shapes # Unused.
self.obj.a.assign(restored_tensors[0])
class MyTrackable(base.Trackable):
def __init__(self):
self.a = module.Module()
class MyTrackableWithSingleSaveable(base.Trackable):
def __init__(self):
self.a = variables.Variable(1.0)
def _gather_saveables_for_checkpoint(self):
return {"foo": lambda name: _VarSaveable(self, name)}
class MyTrackableWithMultiSaveables(base.Trackable):
def __init__(self):
self.a = variables.Variable(1.0)
def _gather_saveables_for_checkpoint(self):
return {
"foo": lambda name: _VarSaveable(self, name),
"bar": lambda name: _VarSaveable(self, name)
}
root = MyTrackable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackableWithMultiSaveables()
with self.assertRaisesRegex(
ValueError,
"Trackable <.*?> expects checkpointed values but checkpoint does not "
"contain serialized tensors for node_id: 0."):
restore.restore_nodes(root_save_path, {0: root2})
root = MyTrackableWithSingleSaveable()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackableWithMultiSaveables()
with self.assertRaisesRegex(
ValueError,
"Size for saveable_objects for Trackable: 2 did not match the size for "
"serialized_tensors for checkpoint: 1."):
restore.restore_nodes(root_save_path, {0: root2})
root = MyTrackableWithMultiSaveables()
root_ckpt = trackable_utils.Checkpoint(root=root)
root_save_path = root_ckpt.save(
os.path.join(self.get_temp_dir(), "root_ckpt"))
root2 = MyTrackableWithSingleSaveable()
with self.assertRaisesRegex(
ValueError,
"Size for saveable_objects for Trackable: 1 did not match the size for "
"serialized_tensors for checkpoint: 2."):
restore.restore_nodes(root_save_path, {0: root2})
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/save_util.py
+347
View File
@@ -0,0 +1,347 @@
# Copyright 2022 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.
# ==============================================================================
"""Extracts tensors for checkpointing while updating a TrackableObjectGraph.
The tensors are extracted from `Trackable._serialize_to_tensors`.
"""
import collections
from typing import Any, Callable, List, Optional, Tuple, Mapping, Union, Dict
from tensorflow.core.protobuf import trackable_object_graph_pb2
from tensorflow.python.checkpoint import graph_view as graph_view_lib
from tensorflow.python.checkpoint import save_util_v1
from tensorflow.python.checkpoint import saveable_compat
from tensorflow.python.checkpoint import util
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import registration
from tensorflow.python.trackable import base
from tensorflow.python.trackable import python_state
from tensorflow.python.trackable import trackable_utils
from tensorflow.python.training.saving import saveable_object as saveable_object_lib
from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.types import core
from tensorflow.python.util import object_identity
# Attributes for each Trackable in the checkpointed object graph.
_TrackableData = collections.namedtuple(
"_TrackableData",
[
# A trackable in the root Trackable object graph.
"trackable",
# The index at which the Trackable appears in TrackableObjectGraph.nodes.
"node_id",
# The BFS-generated path from the root object / used to generate readable
# checkpoint keys.
"object_name",
# A list of ObjectReference for each child connected to this Trackable.
"children_proto",
# A list of SlotVariableReference to save to the object (only valid for
# Optimizer objects).
"slot_variable_proto",
# The object to save to checkpoint. Usually this is the same as `trackable`,
# but can differ when the caller wants to specify a different object to
# save. For example, when saving checkpoints asynchronously, variables are
# copied to the CPU. `object_to_save` is set as the copied variable.
"object_to_save",
],
)
def _split_trackables(
trackable_data: List[_TrackableData]
) -> Tuple[List[_TrackableData], List[_TrackableData],
Dict[str, List[_TrackableData]]]:
"""Splits Trackables into 3 categories (tensor/pystate/registered)."""
tensor_trackables = []
pystate_trackables = []
registered_trackables = collections.defaultdict(list)
for td in trackable_data:
saver_name = registration.get_registered_saver_name(td.object_to_save)
if isinstance(td.object_to_save, python_state.PythonState):
pystate_trackables.append(td)
elif saver_name:
registered_trackables[saver_name].append(td)
else:
tensor_trackables.append(td)
return tensor_trackables, pystate_trackables, registered_trackables
def _gather_trackable_data(
graph_view: graph_view_lib.ObjectGraphView,
object_map: Mapping[base.Trackable, base.Trackable]
) -> Tuple[List[_TrackableData], Dict[base.Trackable, int]]:
"""Returns a list of generated TrackableData based on the ObjectGraphView."""
trackable_objects, node_paths = graph_view.breadth_first_traversal()
object_names = object_identity.ObjectIdentityDictionary()
for obj, path in node_paths.items():
object_names[obj] = trackable_utils.object_path_to_string(path)
node_ids = object_identity.ObjectIdentityDictionary()
for node_id, node in enumerate(trackable_objects):
node_ids[node] = node_id
slot_variables = util.serialize_slot_variables(
trackable_objects=trackable_objects,
node_ids=node_ids,
object_names=object_names)
trackable_data = []
for trackable in trackable_objects:
children_proto = []
for child in graph_view.list_children(trackable):
children_proto.append(
trackable_object_graph_pb2.TrackableObjectGraph.TrackableObject
.ObjectReference(node_id=node_ids[child.ref],
local_name=child.name))
trackable_data.append(_TrackableData(
trackable,
node_id=node_ids[trackable],
object_name=object_names[trackable],
children_proto=children_proto,
slot_variable_proto=slot_variables.get(trackable, []),
object_to_save=util.get_mapped_trackable(trackable, object_map)))
return trackable_data, node_ids
def _fill_object_graph_proto(
trackable_data: List[_TrackableData]
) -> trackable_object_graph_pb2.TrackableObjectGraph:
"""Name non-slot `Trackable`s and add them to `object_graph_proto`."""
object_graph_proto = trackable_object_graph_pb2.TrackableObjectGraph()
for checkpoint_id, td in enumerate(trackable_data):
assert td.node_id == checkpoint_id
object_graph_proto.nodes.add(
slot_variables=td.slot_variable_proto,
children=td.children_proto)
return object_graph_proto
def _get_and_write_tensors_to_serialize(
tensor_trackables: List[_TrackableData],
node_ids: Dict[base.Trackable, int],
call_with_mapped_captures: Union[Callable[..., Any], None],
cache: Union[Dict[base.Trackable, any], None],
object_graph_proto: trackable_object_graph_pb2.TrackableObjectGraph
) -> Dict[base.Trackable, Any]:
"""Creates dictionary of tensors to checkpoint, and updates the proto."""
# Maps trackable to the a dictionary of tensors, which maps
# checkpoint key (-> slice_spec) -> tensor.
serialized_tensors = object_identity.ObjectIdentityDictionary()
for td in tensor_trackables:
if cache is not None and td.object_to_save in cache:
trackable, tensor_dict, object_proto = cache[td.object_to_save]
serialized_tensors[trackable] = tensor_dict
object_graph_proto.nodes[td.node_id].attributes.MergeFrom(object_proto)
continue
legacy_name = saveable_compat.get_saveable_name(td.object_to_save) or ""
if (not saveable_object_util.trackable_has_serialize_to_tensor(
td.object_to_save) or
legacy_name):
# Use the legacy code path for objects that are using SaveableObjects
# or the compat saveable name decorator.
trackable, tensor_dict = _get_tensors_from_legacy_saveable(
td, node_ids, call_with_mapped_captures, object_graph_proto)
else:
tensor_dict = _get_tensors_from_trackable(
td, call_with_mapped_captures, object_graph_proto)
trackable = td.object_to_save
serialized_tensors[trackable] = tensor_dict
if cache is not None and td.object_to_save not in cache:
cache[td.object_to_save] = (
trackable, tensor_dict,
object_graph_proto.nodes[td.node_id].attributes)
return serialized_tensors
def _get_tensors_from_legacy_saveable(
trackable_data: _TrackableData,
node_ids: Dict[base.Trackable, int],
call_with_mapped_captures: Callable[..., Any],
object_graph_proto: trackable_object_graph_pb2.TrackableObjectGraph
) -> Tuple[base.Trackable, Dict[str, Any]]:
"""Gets tensors to serialize from a Trackable with legacy SaveableObjects."""
# Call `save_util_v1` methods to create legacy SaveableObjects and update the
# proto.
object_names = object_identity.ObjectIdentityDictionary()
object_names[trackable_data.trackable] = trackable_data.object_name
object_map = object_identity.ObjectIdentityDictionary()
object_map[trackable_data.trackable] = trackable_data.object_to_save
checkpoint_factory_map, _ = save_util_v1.get_checkpoint_factories_and_keys(
object_names, object_map)
named_saveable_objects, _ = (
save_util_v1.generate_saveable_objects(
checkpoint_factory_map,
object_graph_proto,
node_ids,
object_map,
call_with_mapped_captures,
saveables_cache=None))
trackable = (
saveable_object_util.SaveableCompatibilityConverter(
trackable_data.object_to_save, named_saveable_objects))
return trackable, trackable._serialize_to_tensors() # pylint: disable=protected-access
def _get_tensors_from_trackable(
trackable_data: _TrackableData,
call_with_mapped_captures: Union[Callable[..., Any], None],
object_graph_proto: trackable_object_graph_pb2.TrackableObjectGraph
) -> Dict[str, Any]:
"""Gets tensors to serialize from a Trackable."""
trackable = trackable_data.object_to_save
save_fn = trackable._serialize_to_tensors # pylint: disable=protected-access
if (call_with_mapped_captures and
isinstance(save_fn, core.ConcreteFunction)):
ret_tensor_dict = call_with_mapped_captures(save_fn, [])
else:
ret_tensor_dict = save_fn()
# Create checkpoint keys for each entry in the returned tensor dict, and
# write each entry to the object proto.
tensor_dict = {}
for tensor_name, maybe_tensor in ret_tensor_dict.items():
local_name = trackable_utils.escape_local_name(tensor_name)
checkpoint_key = trackable_utils.checkpoint_key(trackable_data.object_name,
local_name)
tensor_dict[checkpoint_key] = maybe_tensor
# TODO(b/261786493): Delete this when DCheckpoint is removed.
if isinstance(maybe_tensor, saveable_object_lib.SaveSpec):
maybe_tensor.name = checkpoint_key
maybe_tensor.slice_spec = ""
if object_graph_proto is not None:
object_graph_proto.nodes[trackable_data.node_id].attributes.add(
name=local_name,
checkpoint_key=checkpoint_key,
full_name=util.get_full_name(trackable))
return tensor_dict
def _get_and_write_pystate_feed_additions(
pystate_trackables: List[_TrackableData],
cache: Union[Dict[base.Trackable, Any], None],
object_graph_proto=None
) -> Tuple[Dict[base.Trackable, Any], Dict[base.Trackable, Any]]:
"""Gets feed additions needed for checkpointing Python State."""
serialized_tensors = object_identity.ObjectIdentityDictionary()
# Maps tensor placeholders to python values.
feed_additions = {}
for td in pystate_trackables:
trackable = td.object_to_save
checkpoint_key = trackable_utils.checkpoint_key(td.object_name,
python_state.PYTHON_STATE)
if trackable in cache:
save_string = cache[td.object_to_save][python_state.PYTHON_STATE]
else:
with ops.device("/cpu:0"):
save_string = constant_op.constant("", dtype=dtypes.string)
cache[trackable] = {python_state.PYTHON_STATE: save_string}
with ops.init_scope():
value = trackable.serialize()
feed_additions[save_string] = value
serialized_tensors[trackable] = {checkpoint_key: save_string}
object_graph_proto.nodes[td.node_id].attributes.add(
name=python_state.PYTHON_STATE,
checkpoint_key=checkpoint_key,
full_name=util.get_full_name(trackable))
return serialized_tensors, feed_additions
def _get_and_write_registered_savers(
registered_trackables: Dict[str, List[_TrackableData]],
object_graph_proto: trackable_object_graph_pb2.TrackableObjectGraph
) -> Dict[str, Dict[str, base.Trackable]]:
"""Generates dictionary of registered savers and updates the proto."""
registered_savers = collections.defaultdict(dict)
for saver_name, trackables in registered_trackables.items():
for td in trackables:
registered_savers[saver_name][td.object_name] = td.object_to_save
object_proto = object_graph_proto.nodes[td.node_id]
object_proto.registered_saver.name = saver_name
object_proto.registered_saver.object_name = td.object_name
return registered_savers
def serialize_graph_view(
graph_view: graph_view_lib.ObjectGraphView,
object_map: Optional[Mapping[base.Trackable, base.Trackable]] = None,
call_with_mapped_captures: Optional[Callable[..., Any]] = None,
cache: Optional[Dict[base.Trackable, Any]] = None,
):
"""Gathers serialization objects, and creates a TrackableObjectGraph proto."""
# There are 3 types of checkpoint serialization types supported:
# 1. Trackables that override `Trackable._serialize_to_tensor()`.
# 2. PythonState: A special type of Trackable that serializes a Python string.
# 3. Registered Trackable Savers: For objects that need to define advanced
# checkpointing operations not supported by (1) or (2).
trackable_data, node_ids = _gather_trackable_data(graph_view, object_map)
tensor_trackables, pystate_trackables, registered_trackables = (
_split_trackables(trackable_data))
object_graph_proto = _fill_object_graph_proto(trackable_data)
serialized_tensors = _get_and_write_tensors_to_serialize(
tensor_trackables,
node_ids,
call_with_mapped_captures,
cache,
object_graph_proto)
registered_savers = _get_and_write_registered_savers(
registered_trackables, object_graph_proto)
# PythonState trackables must be treated differently depending on if the
# checkpoint is being saved in TF1 graph mode (`cache` exists) or
# eager mode (`cache` is None).
if cache is None:
# When the tensor cache is None, get the serialized tensors directly.
feed_additions = None
serialized_tensors.update(_get_and_write_tensors_to_serialize(
pystate_trackables,
node_ids,
call_with_mapped_captures,
cache,
object_graph_proto))
else:
# Python state is not automatically updated within a TF session so these
# values must be passed to sess.run(feed_additions=...).
new_serialized_tensors, feed_additions = (
_get_and_write_pystate_feed_additions(pystate_trackables,
cache,
object_graph_proto))
serialized_tensors.update(new_serialized_tensors)
# Gather all trackables that have checkpoint values or descendants with
# checkpoint values, and add that info to the proto.
util.add_checkpoint_values_check(object_graph_proto)
return (serialized_tensors, feed_additions, registered_savers,
object_graph_proto)
tensorflow/python/checkpoint/save_util_v1.py
+319
View File
@@ -0,0 +1,319 @@
# Copyright 2022 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.
# ==============================================================================
"""Extracts tensors for checkpointing while updating a TrackableObjectGraph.
This is labelled "v1" because the methods here use SaveableObject, which will
soon be deprecated.
"""
import collections
from tensorflow.core.protobuf import trackable_object_graph_pb2
from tensorflow.python.checkpoint import saveable_compat
from tensorflow.python.checkpoint import util
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import registration
from tensorflow.python.trackable import base
from tensorflow.python.trackable import python_state
from tensorflow.python.trackable import trackable_utils
from tensorflow.python.training.saving import saveable_object as saveable_object_lib
from tensorflow.python.training.saving import saveable_object_util
from tensorflow.python.util import object_identity
# Factory and related info used to build a SaveableObject that saves a Trackable
# to checkpoint.
_CheckpointFactoryData = collections.namedtuple(
"_CheckpointFactoryData", ["factory", "name", "checkpoint_key"])
def get_checkpoint_factories_and_keys(object_names, object_map=None):
"""Gets a map of saveable factories and corresponding checkpoint keys.
Args:
object_names: a dictionary that maps `Trackable` objects to auto-generated
string names.
object_map: a dictionary mapping `Trackable` to copied `Trackable` objects.
The copied objects are generated from `Trackable.
_export_to_saved_model_graph()` which copies the object into another
graph. Generally only resource objects (e.g. Variables, Tables) will be
in this map.
Returns:
A tuple of (
Dictionary mapping trackable -> list of _CheckpointFactoryData,
Dictionary mapping registered saver name -> {object name -> trackable})
"""
checkpoint_factory_map = object_identity.ObjectIdentityDictionary()
unmapped_registered_savers = collections.defaultdict(dict)
for trackable, object_name in object_names.items():
# object_to_save is only used to retrieve the saving functionality. For keys
# and other data, use the original `trackable`.
object_to_save = util.get_mapped_trackable(trackable, object_map)
saver_name = registration.get_registered_saver_name(object_to_save)
if saver_name:
# Add the original trackable instead of `object_to_save` to the returned
# dict because the original is needed for writing the object proto.
unmapped_registered_savers[saver_name][object_name] = trackable
else:
checkpoint_factory_map[trackable] = []
for name, saveable_factory in (
saveable_object_util.saveable_objects_from_trackable(
object_to_save).items()): # pylint: disable=protected-access
# Retrieve the legacy saveable name (for compatibility purposes during
# SaveableObject deprecation)
key_suffix = saveable_compat.get_saveable_name(object_to_save) or name
checkpoint_key = trackable_utils.checkpoint_key(object_name, key_suffix)
if not saveable_compat.force_checkpoint_conversion_enabled():
# Make sure the set the name as the legacy saveable name if there
# is one (only when checkpoint conversion is disabled)
name = key_suffix
checkpoint_factory_map[trackable].append(
_CheckpointFactoryData(
factory=saveable_factory,
name=name,
checkpoint_key=checkpoint_key))
return checkpoint_factory_map, unmapped_registered_savers
def _add_attributes_to_object_graph(trackable_objects, object_graph_proto,
node_ids, object_names, object_map,
call_with_mapped_captures, saveables_cache):
"""Create saveables/savers and corresponding protos in the object graph."""
# The loop below creates TrackableObject protos in the TrackableObjectGraph,
# which are filled in the `_add_attributes_to_object_graph_for_*` methods.
for checkpoint_id, (trackable, unused_object_proto) in enumerate(
zip(trackable_objects, object_graph_proto.nodes)):
assert node_ids[trackable] == checkpoint_id
checkpoint_factory_map, unmapped_registered_savers = (
get_checkpoint_factories_and_keys(object_names, object_map))
# Add attributes, which describe what values are saved in checkpoint for
# this trackable.
registered_savers = _add_attributes_to_object_graph_for_registered_savers(
unmapped_registered_savers, object_graph_proto, node_ids, object_map)
named_saveable_objects, feed_additions = (
generate_saveable_objects(checkpoint_factory_map, object_graph_proto,
node_ids, object_map, call_with_mapped_captures,
saveables_cache))
return named_saveable_objects, feed_additions, registered_savers
def _add_attributes_to_object_graph_for_registered_savers(
unmapped_registered_savers, object_graph_proto, node_ids, object_map):
"""Fills the object graph proto with data about the registered savers."""
registered_savers = collections.defaultdict(dict)
for saver_name, trackables in unmapped_registered_savers.items():
for object_name, trackable in trackables.items():
object_proto = object_graph_proto.nodes[node_ids[trackable]]
object_proto.registered_saver.name = saver_name
object_proto.registered_saver.object_name = object_name
object_to_save = util.get_mapped_trackable(trackable, object_map)
registered_savers[saver_name][object_name] = object_to_save
return registered_savers
def generate_saveable_objects(checkpoint_factory_map,
object_graph_proto=None,
node_ids=None,
object_map=None,
call_with_mapped_captures=None,
saveables_cache=None):
"""Create SaveableObjects and corresponding SerializedTensor protos."""
named_saveable_objects = []
if saveables_cache is None:
# No SaveableObject caching. Either we're executing eagerly, or building a
# static save which is specialized to the current Python state.
feed_additions = None
else:
# If we are caching SaveableObjects, we need to build up a feed_dict with
# functions computing volatile Python state to be saved with the
# checkpoint.
feed_additions = {}
for trackable, factory_data_list in checkpoint_factory_map.items():
fill_object_proto = object_graph_proto is not None and node_ids is not None
if fill_object_proto:
object_proto = object_graph_proto.nodes[node_ids[trackable]]
object_to_save = util.get_mapped_trackable(trackable, object_map)
if saveables_cache is not None:
cached_attributes = saveables_cache.setdefault(object_to_save, {})
else:
cached_attributes = None
for factory_data in factory_data_list:
name = factory_data.name
key = factory_data.checkpoint_key
saveable_factory = factory_data.factory
# See if we can skip saving this checkpoint key.
saveables = cached_attributes.get(name) if cached_attributes else None
if saveables is not None:
for saveable in saveables:
if key not in saveable.name:
# The checkpoint key for this SaveableObject is different. We
# need to re-create it.
saveables = None
del cached_attributes[name]
break
if saveables is None:
if callable(saveable_factory):
maybe_saveable = saveable_object_util.create_saveable_object(
name, key, saveable_factory, call_with_mapped_captures)
else:
maybe_saveable = saveable_factory
if isinstance(maybe_saveable, saveable_object_lib.SaveableObject):
saveables = (maybe_saveable,)
else:
saveables = tuple(
saveable_object_util.saveable_objects_for_op(
op=maybe_saveable, name=key))
for saveable in saveables:
if key not in saveable.name:
raise AssertionError(
f"The object {trackable} produced a SaveableObject with name "
f"'{saveable.name}' for attribute '{name}'. Expected a name"
f" containing '{key}'.")
if cached_attributes is not None:
cached_attributes[name] = saveables
if isinstance(object_to_save, python_state.PythonState):
assert len(saveables) == 1
saveable = saveables[0]
if feed_additions is None:
assert saveables_cache is None
# If we're not caching saveables, then we're either executing
# eagerly or building a static save/restore (e.g. for a
# SavedModel). In either case, we should embed the current Python
# state in the graph rather than relying on a feed dict.
saveables = (saveable.freeze(),)
else:
feed_additions.update(saveable.feed_dict_additions())
named_saveable_objects.extend(saveables)
# Update the object proto.
# For updated Trackables that override serialize_to_tensors, add an
# attribute for each tensor that is serialized.
# For Trackables that have SaveableObjects or a legacy saveable name,
# add a single attribute to the proto.
if not fill_object_proto:
continue
if (isinstance(saveables[0], saveable_object_util.TrackableSaveable) and
(saveable_compat.force_checkpoint_conversion_enabled() or
saveable_compat.get_saveable_name(object_to_save) is None)):
for local_name, local_key in (
saveables[0].get_proto_names_and_checkpoint_keys()):
object_proto.attributes.add(
name=local_name,
checkpoint_key=local_key,
full_name=util.get_full_name(object_to_save))
else:
object_proto.attributes.add(
name=name,
checkpoint_key=key,
full_name=util.get_full_name(object_to_save))
return named_saveable_objects, feed_additions
def _fill_object_graph_proto(graph_view,
trackable_objects,
node_ids,
slot_variables):
"""Name non-slot `Trackable`s and add them to `object_graph_proto`."""
object_graph_proto = trackable_object_graph_pb2.TrackableObjectGraph()
for checkpoint_id, trackable in enumerate(trackable_objects):
assert node_ids[trackable] == checkpoint_id
object_proto = object_graph_proto.nodes.add(
slot_variables=slot_variables.get(trackable, ())
)
for child in graph_view.list_children(trackable):
object_proto.children.add(
node_id=node_ids[child.ref],
local_name=child.name)
return object_graph_proto
def serialize_gathered_objects(graph_view,
object_map=None,
call_with_mapped_captures=None,
saveables_cache=None):
"""Create SaveableObjects and protos for gathered objects."""
trackable_objects, node_paths = graph_view.breadth_first_traversal()
object_names = object_identity.ObjectIdentityDictionary()
for obj, path in node_paths.items():
object_names[obj] = trackable_utils.object_path_to_string(path)
node_ids = object_identity.ObjectIdentityDictionary()
for node_id, node in enumerate(trackable_objects):
node_ids[node] = node_id
slot_variables = util.serialize_slot_variables(
trackable_objects=trackable_objects,
node_ids=node_ids,
object_names=object_names)
object_graph_proto = _fill_object_graph_proto(
graph_view=graph_view,
trackable_objects=trackable_objects,
node_ids=node_ids,
slot_variables=slot_variables)
named_saveable_objects, feed_additions, registered_savers = (
_add_attributes_to_object_graph(
trackable_objects=trackable_objects,
object_graph_proto=object_graph_proto,
node_ids=node_ids,
object_names=object_names,
object_map=object_map,
call_with_mapped_captures=call_with_mapped_captures,
saveables_cache=saveables_cache))
# Gather all trackables that have checkpoint values or descendants with
# checkpoint values, and add that info to the proto.
util.add_checkpoint_values_check(object_graph_proto)
return (named_saveable_objects, object_graph_proto, feed_additions,
registered_savers)
def serialize_object_graph_with_registered_savers(graph_view, saveables_cache):
"""Determine checkpoint keys for variables and build a serialized graph."""
return serialize_gathered_objects(graph_view, saveables_cache=saveables_cache)
def frozen_saveables_and_savers(graph_view,
object_map=None,
to_graph=None,
call_with_mapped_captures=None,
saveables_cache=None):
"""Generates SaveableObjects and registered savers in the frozen graph."""
if to_graph:
target_context = to_graph.as_default
else:
target_context = ops.NullContextmanager
with target_context():
named_saveable_objects, graph_proto, _, registered_savers = (
serialize_gathered_objects(graph_view, object_map,
call_with_mapped_captures, saveables_cache))
with ops.device("/cpu:0"):
object_graph_tensor = constant_op.constant(
graph_proto.SerializeToString(), dtype=dtypes.string)
named_saveable_objects.append(
base.NoRestoreSaveable(
tensor=object_graph_tensor, name=base.OBJECT_GRAPH_PROTO_KEY))
return named_saveable_objects, registered_savers
tensorflow/python/checkpoint/save_util_v1_test.py
+78
View File
@@ -0,0 +1,78 @@
# 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 the checkpoint/util.py."""
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.checkpoint import save_util_v1
from tensorflow.python.eager import test
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import registration
from tensorflow.python.trackable import autotrackable
from tensorflow.python.util import object_identity
class TrackableWithRegisteredSaver(autotrackable.AutoTrackable):
pass
registration.register_checkpoint_saver(
name="RegisteredSaver",
predicate=lambda x: isinstance(x, TrackableWithRegisteredSaver),
save_fn=lambda trackables, file_prefix: [],
restore_fn=lambda trackables, merged_prefix: None)
class SerializationTest(test.TestCase):
def test_serialize_gathered_objects(self):
root = autotrackable.AutoTrackable()
root.v = variables.Variable(1.0)
root.registered = TrackableWithRegisteredSaver()
named_saveable_objects, _, _, registered_savers = (
save_util_v1.serialize_gathered_objects(
graph_view.ObjectGraphView(root)))
self.assertLen(named_saveable_objects, 1)
self.assertIs(named_saveable_objects[0].op, root.v)
self.assertDictEqual(
{"Custom.RegisteredSaver": {"registered": root.registered}},
registered_savers)
def test_serialize_gathered_objects_with_map(self):
root = autotrackable.AutoTrackable()
root.v = variables.Variable(1.0)
root.registered = TrackableWithRegisteredSaver()
copy_of_registered = TrackableWithRegisteredSaver()
copy_of_v = variables.Variable(1.0)
object_map = object_identity.ObjectIdentityDictionary()
object_map[root.registered] = copy_of_registered
object_map[root.v] = copy_of_v
named_saveable_objects, _, _, registered_savers = (
save_util_v1.serialize_gathered_objects(
graph_view.ObjectGraphView(root), object_map))
self.assertLen(named_saveable_objects, 1)
self.assertIsNot(named_saveable_objects[0].op, root.v)
self.assertIs(named_saveable_objects[0].op, copy_of_v)
ret_value = registered_savers["Custom.RegisteredSaver"]["registered"]
self.assertIsNot(root.registered, ret_value)
self.assertIs(copy_of_registered, ret_value)
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/saveable_compat.py
+115
View File
@@ -0,0 +1,115 @@
# Copyright 2022 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.
# ==============================================================================
"""Checkpoint compatibility functions with SaveableObject.
Compatibility methods to ensure that checkpoints are saved with the same
metadata attributes before/after the SaveableObject deprecation.
"""
_LEGACY_SAVEABLE_NAME = "_LEGACY_SAVEABLE_NAME"
def legacy_saveable_name(name):
"""Decorator to set the local name to use in the Checkpoint.
Needed for migrating certain Trackables (see next paragraph) from the legacy
`_gather_saveables_for_checkpoint` to the new `_serialize_to_tensors`
function.
This decorator should be used if the SaveableObject generates tensors with
different names from the name that is passed to the factory.
Example migration:
*Before*
```
class MyTrackable(Trackable):
def _gather_saveables_for_checkpoint(self):
return {"key": _MySaveable}
class _MySaveable(SaveableObject):
def __init__(self, name):
specs = [
SaveSpec(tensor1, "", name + "-1")
SaveSpec(tensor2, "", name + "-2")
]
super().__init__(None, specs, name)
```
*After*
```
@legacy_saveable_name("key")
class MyTrackable(Trackable):
def _serialize_to_tensors(self):
return {"key-1": tensor1, "key-2": tensor2}
```
Args:
name: String name of the SaveableObject factory (the key returned in the
`_gather_saveables_for_checkpoint` function)
Returns:
A decorator.
"""
def decorator(cls_or_obj):
setattr(cls_or_obj, _LEGACY_SAVEABLE_NAME, name)
return cls_or_obj
return decorator
def get_saveable_name(cls_or_obj):
return getattr(cls_or_obj, _LEGACY_SAVEABLE_NAME, None)
_FORCE_CHECKPOINT_CONVERSION = False
def force_checkpoint_conversion(value=True):
"""Forces checkpoint to use the new implementation.
The new checkpoint implementation is changing the saved metadata slightly,
and therefore may break forward compatibility in newly saved checkpoints. This
means:
- Previous versions of TensorFlow may not be able to load new checkpoints.
- Backwards compatibility is unchanged: Old checkpoints can still be loaded.
TensorFlow guarantees 3 weeks of forward compatibility, so this flag will be
removed in the future weeks, after which checkpoint conversion will happen by
default.
**What happens when this flag is enabled?**
The checkpoint will be saved with different metadata, meaning that previous
versions of TensorFlow (<=2.10) will not be able to load this checkpoint.
Args:
value: Boolean value, whether or not to force checkpoint conversion to the
new implementation.
"""
# TODO(kathywu): Add definite date for flag removal.
global _FORCE_CHECKPOINT_CONVERSION
_FORCE_CHECKPOINT_CONVERSION = value
def force_checkpoint_conversion_enabled():
return _FORCE_CHECKPOINT_CONVERSION
class CheckpointConversionError(Exception):
pass
tensorflow/python/checkpoint/saveable_compat_test.py
+168
View File
@@ -0,0 +1,168 @@
# Copyright 2022 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 SaveableObject compatibility."""
import os
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import saveable_compat
from tensorflow.python.checkpoint.testdata import generate_checkpoint
from tensorflow.python.eager import test
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import base
from tensorflow.python.training import checkpoint_utils
from tensorflow.python.training.saving import saveable_object
_LEGACY_TABLE_CHECKPOINT_PATH = test.test_src_dir_path(
"python/checkpoint/testdata/table_legacy_saveable_object")
class SaveableCompatTest(test.TestCase):
def test_lookup_table_compatibility(self):
saveable_compat.force_checkpoint_conversion(False)
table_module = generate_checkpoint.TableModule()
ckpt = checkpoint.Checkpoint(table_module)
checkpoint_directory = self.get_temp_dir()
checkpoint_path = os.path.join(checkpoint_directory, "ckpt")
ckpt.write(checkpoint_path)
# Ensure that the checkpoint metadata and keys are the same.
legacy_metadata = checkpoint.object_metadata(_LEGACY_TABLE_CHECKPOINT_PATH)
metadata = checkpoint.object_metadata(checkpoint_path)
def _get_table_node(object_metadata):
for child in object_metadata.nodes[0].children:
if child.local_name == "lookup_table":
return object_metadata.nodes[child.node_id]
table_proto = _get_table_node(metadata)
legacy_table_proto = _get_table_node(legacy_metadata)
self.assertAllEqual(
[table_proto.attributes[0].name,
table_proto.attributes[0].checkpoint_key],
[legacy_table_proto.attributes[0].name,
legacy_table_proto.attributes[0].checkpoint_key])
legacy_reader = checkpoint_utils.load_checkpoint(
_LEGACY_TABLE_CHECKPOINT_PATH)
reader = checkpoint_utils.load_checkpoint(checkpoint_path)
self.assertEqual(
legacy_reader.get_variable_to_shape_map().keys(),
reader.get_variable_to_shape_map().keys())
# Ensure that previous checkpoint can be loaded into current table.
ckpt.read(_LEGACY_TABLE_CHECKPOINT_PATH).assert_consumed()
class TestForceCheckpointConversionFlag(test.TestCase):
def test_checkpoint(self):
saveable_compat.force_checkpoint_conversion()
table_module = generate_checkpoint.TableModule()
table_module.lookup_table.insert(3, 9)
ckpt = checkpoint.Checkpoint(table_module)
checkpoint_directory = self.get_temp_dir()
checkpoint_path = os.path.join(checkpoint_directory, "ckpt")
ckpt.write(checkpoint_path)
new_table_module = generate_checkpoint.TableModule()
self.assertEqual(-1, self.evaluate(new_table_module.lookup_table.lookup(3)))
new_ckpt = checkpoint.Checkpoint(new_table_module)
new_ckpt.read(checkpoint_path).assert_consumed()
self.assertEqual(9, self.evaluate(new_table_module.lookup_table.lookup(3)))
def test_backwards_compatibility(self):
saveable_compat.force_checkpoint_conversion()
table_module = generate_checkpoint.TableModule()
table_module.lookup_table.insert(3, 9)
self.assertEqual(9, self.evaluate(table_module.lookup_table.lookup(3)))
ckpt = checkpoint.Checkpoint(table_module)
ckpt.read(_LEGACY_TABLE_CHECKPOINT_PATH).assert_consumed()
self.assertEqual(-1, self.evaluate(table_module.lookup_table.lookup(3)))
self.assertEqual(4, self.evaluate(table_module.lookup_table.lookup(2)))
def test_forward_compatibility(self):
class _MultiSpecSaveable(saveable_object.SaveableObject):
def __init__(self, obj, name):
self.obj = obj
specs = [
saveable_object.SaveSpec(obj.a, "", name + "-a"),
saveable_object.SaveSpec(obj.b, "", name + "-b")]
super(_MultiSpecSaveable, self).__init__(None, specs, name)
def restore(self, restored_tensors, restored_shapes):
del restored_shapes # Unused.
self.obj.a.assign(restored_tensors[0])
self.obj.b.assign(restored_tensors[1])
class DeprecatedTrackable(base.Trackable):
def __init__(self):
self.a = variables.Variable(1.0)
self.b = variables.Variable(2.0)
def _gather_saveables_for_checkpoint(self):
return {"foo": lambda name: _MultiSpecSaveable(self, name)}
@saveable_compat.legacy_saveable_name("foo")
class NewTrackable(base.Trackable):
def __init__(self):
self.a = variables.Variable(3.0)
self.b = variables.Variable(4.0)
def _serialize_to_tensors(self):
return {"-a": self.a, "-b": self.b}
def _restore_from_tensors(self, restored_tensors):
return control_flow_ops.group(
self.a.assign(restored_tensors["-a"]),
self.b.assign(restored_tensors["-b"]))
new = NewTrackable()
# Test with the checkpoint conversion flag disabled (normal compatibility).
saveable_compat.force_checkpoint_conversion(False)
checkpoint_path = os.path.join(self.get_temp_dir(), "ckpt")
checkpoint.Checkpoint(new).write(checkpoint_path)
dep = DeprecatedTrackable()
checkpoint.Checkpoint(dep).read(checkpoint_path).assert_consumed()
self.assertEqual(3, self.evaluate(dep.a))
self.assertEqual(4, self.evaluate(dep.b))
# Now test with the checkpoint conversion flag enabled (forward compat).
# The deprecated object will try to load from the new checkpoint.
saveable_compat.force_checkpoint_conversion()
checkpoint_path = os.path.join(self.get_temp_dir(), "ckpt2")
checkpoint.Checkpoint(new).write(checkpoint_path)
dep = DeprecatedTrackable()
checkpoint.Checkpoint(dep).read(checkpoint_path).assert_consumed()
self.assertEqual(3, self.evaluate(dep.a))
self.assertEqual(4, self.evaluate(dep.b))
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/sharding/BUILD
+104
View File
@@ -0,0 +1,104 @@
# Description:
# Utilities for sharding object-based checkpoints.
load("//tensorflow:pytype.default.bzl", "pytype_strict_library")
load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
],
licenses = ["notice"],
)
pytype_strict_library(
name = "sharding_policies",
srcs = ["sharding_policies.py"],
deps = [
":sharding_util",
"//tensorflow/python/eager:context",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:string_ops",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/util:tf_export",
"@absl_py//absl/logging",
],
)
tf_py_strict_test(
name = "sharding_policies_test",
srcs = ["sharding_policies_test.py"],
tags = [
"notap",
"oss_excluded",
], # See b/455621783.
deps = [
":sharding_policies",
":sharding_util",
"//tensorflow/python/checkpoint",
"//tensorflow/python/checkpoint:checkpoint_options",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/eager:def_function",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:test_lib",
"//tensorflow/python/module",
"//tensorflow/python/ops:variables",
"//tensorflow/python/platform:gfile",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
"@absl_py//absl/logging",
],
)
pytype_strict_library(
name = "sharding_util",
srcs = ["sharding_util.py"],
deps = [
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/framework:tensor_shape",
"//tensorflow/python/framework:tensor_spec",
"//tensorflow/python/ops:variables",
"//tensorflow/python/trackable:base",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/util:tf_export",
],
)
tf_py_strict_test(
name = "sharding_util_test",
srcs = ["sharding_util_test.py"],
deps = [
":sharding_policies",
":sharding_util",
"//tensorflow/python/checkpoint",
"//tensorflow/python/checkpoint:graph_view",
"//tensorflow/python/eager:remote",
"//tensorflow/python/eager:test",
"//tensorflow/python/framework:device",
"//tensorflow/python/framework:dtypes",
"//tensorflow/python/framework:ops",
"//tensorflow/python/framework:tensor",
"//tensorflow/python/module",
"//tensorflow/python/ops:array_ops",
"//tensorflow/python/ops:math_ops",
"//tensorflow/python/ops:resource_variable_ops",
"//tensorflow/python/training:server_lib",
"//tensorflow/python/training/saving:saveable_object",
"//tensorflow/python/training/saving:saveable_object_util",
],
)
tensorflow/python/checkpoint/sharding/sharding_policies.py
+353
View File
@@ -0,0 +1,353 @@
# Copyright 2023 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.
# ==============================================================================
"""Checkpoint policies that determine how tensors are split into shards."""
import math
import operator
from typing import MutableSequence, Sequence
from absl import logging
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import context
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import string_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import base
from tensorflow.python.util import tf_export
@tf_export.tf_export("train.experimental.ShardByTaskPolicy")
class ShardByTaskPolicy(sharding_util.ShardingCallback):
"""Policy that splits tensors into shards based on their device spec task."""
@property
def description(self) -> str:
return "Split tensors into shards based on their device spec task."
def __call__(
self,
shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
"""Callback to split tensors into shards based on their device spec task.
Args:
shardable_tensors: A list of ShardableTensors.
Returns:
List of shard dicts containing tensors.
[ {checkpoint key: {slice_spec: tensor} } ]
"""
tensors_by_task = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
(tensors_by_task
.setdefault(checkpoint_key, {})[slice_spec]) = tensor
return [tensors_by_task]
_OffsetAndShape = tuple[Sequence[int], Sequence[int]]
@tf_export.tf_export("train.experimental.MaxShardSizePolicy")
class MaxShardSizePolicy(sharding_util.ShardingCallback):
"""Policy that splits tensors into shards with a max shard size.
Shards may exceed the max shard size if they contain 1. a single scalar/string
tensor that could not be sliced and exceeds the max shard size or 2. the
checkpoint object graph, whose size cannot be calculated when saving.
"""
class MaxShardSizePartitioner():
"""Partition tensors into shards with a max shard size."""
max_shard_size: int
_large_scalars: MutableSequence[sharding_util.Shard]
_tensors_by_shard: MutableSequence[sharding_util.Shard]
_shard_size_remaining: int
_checkpoint_key: str
_dtype: dtypes.DType
_device: device_lib.DeviceSpec
_root_tensor: tensor_lib.Tensor
_slice_spec: variables.Variable.SaveSliceInfo
_full_shape: tensor_shape.TensorShape
_root_shape: tensor_shape.TensorShape
_root_offset: Sequence[int]
_dtype_size: int
_working_tensor_offset: MutableSequence[float]
_working_tensor_shape: tensor_shape.TensorShape
def _get_next_partition(self) -> tuple[int, float]:
"""Gets tensor partition with size closest to shard_size_remaining.
Returns:
A tuple containing the axis and size of the next partition.
"""
rank = self._working_tensor_shape.rank
if rank is None or rank == 0:
return 0, math.inf
num_elems = self._working_tensor_shape.num_elements()
def num_partitions(axis: int) -> float:
axis_len = self._working_tensor_shape.dims[axis].value
slice_elems = num_elems // axis_len
bytes_per_slice = slice_elems * self._dtype_size
slices_per_shard = self._shard_size_remaining // bytes_per_slice
if slices_per_shard == 0:
return math.inf
return math.ceil(axis_len / slices_per_shard)
# Find axis with minimum partitions. (axis with maximum partition size)
# (max partition size is as close as possible to the shard_size_remaining)
min_parts = num_partitions(0)
min_axis = 0
for axis in range(1, rank):
parts_along_axis = num_partitions(axis)
part_size = num_elems * self._dtype_size / parts_along_axis
if (parts_along_axis < min_parts and
part_size <= self._shard_size_remaining):
min_axis, min_parts = axis, int(parts_along_axis)
return (min_axis,
math.ceil(int(self._working_tensor_shape[min_axis]) / min_parts))
def _add_partition(self, part_axis: int, part_size: float):
"""Adds the tensor partition to the shard, if possible.
Args:
part_axis: The axis of the partition.
part_size: The size of the partition.
Raises:
RuntimeError: When the slice size is larger than the remaining shard
size.
"""
# Add what we can to the current shard.
relative_offset = list(
map(operator.sub, self._working_tensor_offset, self._root_offset))
slice_shape = list(map(operator.sub, self._root_shape, relative_offset))
slice_shape[part_axis] = part_size
slice_size_in_bytes = int(math.prod(slice_shape)) * self._dtype_size
with ops.device(self._device):
tensor_slice = array_ops.slice(
self._root_tensor, begin=relative_offset, size=slice_shape)
slice_spec = variables.Variable.SaveSliceInfo(
full_name=self._checkpoint_key,
full_shape=self._full_shape,
var_offset=self._working_tensor_offset,
var_shape=slice_shape).spec.strip()
if slice_size_in_bytes > self.max_shard_size:
logging.warning("Tensor %s's minimum slice %s has size %s bytes and "
"cannot be partitioned into a shard of max shard size "
"%s bytes. It will be added as an individual shard "
"that exceeds the max shard size.",
self._checkpoint_key, slice_spec, slice_size_in_bytes,
self.max_shard_size)
self._large_scalars.append(
{self._checkpoint_key: {slice_spec: tensor_slice}})
elif slice_size_in_bytes > self._shard_size_remaining:
raise RuntimeError(
f"Slice size ({slice_size_in_bytes} bytes) is larger than the "
f"remaining shard size ({self._shard_size_remaining} bytes). This "
"should have been caught in MaxShardSizePolicy._add_partition().")
else:
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})[slice_spec]) = tensor_slice
self._shard_size_remaining -= slice_size_in_bytes
if self._shard_size_remaining == 0:
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
# Get remaining portion of tensor to add to the next shard(s).
self._working_tensor_offset[part_axis] += part_size
relative_offset[part_axis] += part_size
self._working_tensor_shape = tensor_shape.TensorShape(list(
map(operator.sub, self._root_shape, relative_offset)))
def get_shards(
self,
max_shard_size: int,
shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
"""Callback to split tensors into shards with a max shard size.
Args:
max_shard_size: The maximum size of a shard file in bytes.
shardable_tensors: A list of ShardableTensors.
Returns:
List of shard dicts containing tensors.
[ {checkpoint key: {slice_spec: tensor} } ]
"""
self.max_shard_size = max_shard_size
self._tensors_by_shard = [{}]
self._large_scalars = []
string_size_warning_printed = False
self._shard_size_remaining = self.max_shard_size
for shardable_tensor in shardable_tensors:
self._checkpoint_key = shardable_tensor.checkpoint_key
self._dtype = shardable_tensor.dtype
self._device = shardable_tensor.device
self._root_tensor = shardable_tensor.tensor
self._slice_spec = shardable_tensor.slice_spec
# If the tensor has already been sliced, make sure to keep track of its
# parent tensor's shape & offset. These will be used when creating slice
# specs later.
if self._slice_spec:
save_slice_info = variables.Variable.SaveSliceInfo.from_spec(
self._slice_spec)
self._full_shape = tensor_shape.TensorShape(
save_slice_info.full_shape)
self._root_shape = tensor_shape.TensorShape(save_slice_info.var_shape)
self._root_offset = save_slice_info.var_offset
else:
self._full_shape = self._root_shape = shardable_tensor.shape
self._root_offset = [0] * self._root_shape.rank
self._dtype_size = dtypes.as_dtype(self._dtype).size
total_size = self._root_shape.num_elements() * self._dtype_size # bytes
# Calculate string tensor sizes.
if self._checkpoint_key == base.OBJECT_GRAPH_PROTO_KEY:
# In graph mode, the object graph is populated using feed_additions
# when the session is run. So, we can't calculate the size here.
# Fortunately, the serialized object graph string will never be that
# big, so we just place it in the current shard without worrying about
# its size.
total_size = self._dtype_size = 0
elif self._dtype == dtypes.variant:
# Can't determine a variant's type, so can't calculate its size.
total_size = self._dtype_size = 0
elif (self._dtype == dtypes.string
and not context.executing_eagerly()
and ops.get_default_session() is None):
# TODO(b/326287351): Get string tensor size in tf.function.
total_size = self._dtype_size = 0
if not string_size_warning_printed:
logging.warning("The checkpoint sharding policy is being executed "
"in a tf.function. The size of the string/variant "
"constant cannot be obtained.")
string_size_warning_printed = True
elif self._dtype == dtypes.string:
with ops.device(self._device):
if not context.executing_eagerly():
self._root_tensor = ops.get_default_session().run(
self._root_tensor)
if self._root_shape.rank is None or self._root_shape.rank == 0:
sizes = [string_ops.string_length(self._root_tensor,
unit="BYTE")]
else:
sizes = [string_ops.string_length(elem, unit="BYTE")
for elem in self._root_tensor]
if context.executing_eagerly():
sizes = [size.numpy() for size in sizes]
else:
sizes = ops.get_default_session().run(sizes)
total_size = sum(sizes)
self._dtype_size = max(sizes)
if self._root_shape.rank is None or self._root_shape.rank == 0:
if total_size > self.max_shard_size:
logging.warning(
"Tensor %s is a %s scalar of size %s bytes and cannot be "
"partitioned into a shard of max shard size %s bytes. It will "
"be added as an individual shard that exceeds the max shard "
"size.", self._checkpoint_key, self._dtype, total_size,
self.max_shard_size)
self._large_scalars.append(
{self._checkpoint_key: {self._slice_spec: self._root_tensor}})
else:
if total_size > self._shard_size_remaining:
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})
[self._slice_spec]) = self._root_tensor
self._shard_size_remaining -= total_size
continue
# Partition tensor and add partitions to shards.
self._working_tensor_offset = self._root_offset[:]
self._working_tensor_shape = self._root_shape
working_tensor_size = total_size
while working_tensor_size > self._shard_size_remaining:
(part_axis, part_size) = self._get_next_partition()
if part_size == 0:
# Tensor partition couldn't fit in remaining shard space. Try again
# with the next full shard.
self._tensors_by_shard.append({})
self._shard_size_remaining = self.max_shard_size
else:
self._add_partition(part_axis, part_size)
working_tensor_size = (
int(math.prod(self._working_tensor_shape)) * self._dtype_size)
if self._working_tensor_shape.num_elements() > 0:
if self._working_tensor_offset and self._working_tensor_shape:
with ops.device(self._device):
working_tensor = array_ops.slice(
self._root_tensor,
begin=list(map(
operator.sub,
self._working_tensor_offset, self._root_offset)),
size=self._working_tensor_shape.as_list())
else:
working_tensor = self._root_tensor
remaining_tensor_slice_spec = variables.Variable.SaveSliceInfo(
full_name=self._checkpoint_key,
full_shape=self._full_shape,
var_offset=self._working_tensor_offset,
var_shape=self._working_tensor_shape).spec.strip()
(self._tensors_by_shard[-1]
.setdefault(self._checkpoint_key, {})
[remaining_tensor_slice_spec]) = working_tensor
self._shard_size_remaining -= working_tensor_size
shards = []
if self._tensors_by_shard[0]:
shards.extend(self._tensors_by_shard)
shards.extend(self._large_scalars)
return shards
def __init__(self, max_shard_size: int):
self.max_shard_size = max_shard_size
@property
def description(self) -> str:
return "Split tensors into shards with a max shard size."
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
return self.MaxShardSizePartitioner().get_shards(
self.max_shard_size, shardable_tensors)
tensorflow/python/checkpoint/sharding/sharding_policies_test.py
+810
View File
@@ -0,0 +1,810 @@
# Copyright 2023 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 checkpoint sharding policies."""
import random
import re
import string
from absl import logging
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import checkpoint_options
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.checkpoint.sharding import sharding_policies
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import def_function
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import test_util
from tensorflow.python.module import module
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import gfile
from tensorflow.python.training import server_lib
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
class ShardingPoliciesTest(test.TestCase):
def _get_shardable_tensors_by_task(self, root):
serialized_tensors, _, _, _ = (
checkpoint.TrackableSaver(graph_view.ObjectGraphView(root))
._gather_serialized_tensors(None))
shardable_tensors_by_task = {}
for obj, tensor_dict in serialized_tensors.items():
# Divide tensor_dict by device.
for checkpoint_key, tensor_slice_dict in tensor_dict.items():
if not isinstance(tensor_slice_dict, dict):
# Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
tensor_slice_dict = {"": tensor_slice_dict}
for slice_spec, tensor_save_spec in tensor_slice_dict.items():
if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
tensor_save_spec = saveable_object.SaveSpec(
tensor=tensor_save_spec,
slice_spec=slice_spec,
name=checkpoint_key,
dtype=tensor_save_spec.dtype,
device=tensor_save_spec.device)
save_spec_tensor = tensor_save_spec.tensor
device = (device_lib.DeviceSpec.from_string(tensor_save_spec.device)
if isinstance(tensor_save_spec.device, str)
else tensor_save_spec.device)
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(device.to_string()))
shardable_tensors_by_task.setdefault(task, []).append(
sharding_util.ShardableTensor(
_tensor_save_spec=tensor_save_spec,
tensor=save_spec_tensor,
dtype=tensor_save_spec.dtype,
device=device,
name=tensor_save_spec.name,
shape=save_spec_tensor.shape,
slice_spec=slice_spec,
checkpoint_key=checkpoint_key,
trackable=obj))
return shardable_tensors_by_task.values()
def test_ShardByTaskPolicy(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")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
with ops.device("/job:worker/task:2/cpu:0"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
root.v0 = v0
root.v1 = v1
root.v2 = v2
shardable_tensors = self._get_shardable_tensors_by_task(root)
callback = sharding_policies.ShardByTaskPolicy()
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertAllEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE"},
{"v1/.ATTRIBUTES/VARIABLE_VALUE"},
{"v2/.ATTRIBUTES/VARIABLE_VALUE"},
{"_CHECKPOINTABLE_OBJECT_GRAPH"}
])
self.assertEqual(
self.evaluate(shards[0]["v0/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v0.numpy())
self.assertEqual(
self.evaluate(shards[1]["v1/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v1.numpy())
self.assertEqual(
self.evaluate(shards[2]["v2/.ATTRIBUTES/VARIABLE_VALUE"][""]),
v2.numpy())
def test_CheckpointOption_ShardByTaskPolicy(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")
self.evaluate(v1.initializer)
with ops.device("/job:worker/task:2/cpu:0"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
self.evaluate(v2.initializer)
root.v0 = v0
root.v1 = v1
root.v2 = v2
tmp_dir = self.create_tempdir("ckpt")
ckpt = checkpoint.Checkpoint(root)
save_path = ckpt.save(
tmp_dir, options=checkpoint_options.CheckpointOptions(
experimental_sharding_callback=(
sharding_policies.ShardByTaskPolicy())))
self.assertLen(gfile.Glob(save_path + ".data*"), 4)
ckpt.restore(save_path)
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_1D(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([0.0, 1.0, 2.0, 3.0],
name="v0",
dtype=dtypes.float32)
v1 = resource_variable_ops.ResourceVariable([[4],
[5],
[6],
[7]],
name="v1",
dtype=dtypes.int32)
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
root.v0 = v0
root.v1 = v1
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
v1_name = "v1/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[4]),
var_offset=var_offset,
var_shape=var_shape)
class V1SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v1_name,
full_shape=tensor_shape.TensorShape(dims=[4, 1]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 4 bytes
# Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
# split into 4 shards.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=4)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[0][v0_name][slice_spec]), 0.0)
slice_spec = V0SaveSliceInfo(var_offset=[1], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[1][v0_name][slice_spec]), 1.0)
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[2][v0_name][slice_spec]), 2.0)
slice_spec = V0SaveSliceInfo(var_offset=[3], var_shape=[1]).spec
self.assertEqual(self.evaluate(shards[3][v0_name][slice_spec]), 3.0)
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[4][v1_name][slice_spec]), [4])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[5][v1_name][slice_spec]), [5])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[6][v1_name][slice_spec]), [6])
slice_spec = V1SaveSliceInfo(var_offset=[3, 0], var_shape=[1, 1]).spec
self.assertEqual(self.evaluate(shards[7][v1_name][slice_spec]), [7])
# max_shard_size: 8 bytes
# v0/v1 haven't changed, so they should now be split into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]])
# max_shard_size: 10 bytes
# 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
# into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]), [[4], [5]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[6], [7]])
# max_shard_size: 16 bytes
# 16 bytes the exact size of each variable, so they should get 1 shard each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=16)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0, 2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[4, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[4], [5], [6], [7]])
# max_shard_size: 18 bytes
# 18 bytes slightly larger than the size of each variable, but not large
# enough to fit another 4 byte element, so they should get 1 shard each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=18)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[4]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [0.0, 1.0, 2.0, 3.0])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0], var_shape=[4, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[4], [5], [6], [7]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_2D(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0, 1],
[2, 3],
[4, 5]],
name="v0")
v1 = resource_variable_ops.ResourceVariable([[[6.0], [7.0]],
[[8.0], [9.0]],
[[10.0], [11.0]]], name="v1")
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
root.v0 = v0
root.v1 = v1
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
v1_name = "v1/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[3, 2]),
var_offset=var_offset,
var_shape=var_shape)
class V1SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v1_name,
full_shape=tensor_shape.TensorShape(dims=[3, 2, 1]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 8 bytes
# Each element of v0/v1 is a 32 bit/4 byte value, so each variable should be
# split into 3 shards.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[4][v1_name][slice_spec]), [[[8.0], [9.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[5][v1_name][slice_spec]), [[[10.0], [11.0]]])
# max_shard_size: 10 bytes
# 10 bytes is an uneven boundary for 4 byte elements. v0/v1 should be split
# into 3 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[4][v1_name][slice_spec]), [[[8.0], [9.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[2, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[5][v1_name][slice_spec]), [[[10.0], [11.0]]])
# max_shard_size: 12 bytes
# 12 bytes is enough to fit 3 elements per variable in each shard.
# v0/v1 should be split into 2 shards each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=12)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE",},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[3, 1]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0], [2], [4]])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[3, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[1], [3], [5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[3, 1, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]),
[[[6.0]], [[8.0]], [[10.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[0, 1, 0], var_shape=[3, 1, 1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v1_name][slice_spec]),
[[[7.0]], [[9.0]], [[11.0]]])
# max_shard_size: 16 bytes
# Each variable should be split into 1.5 shards. The middle shard will
# contain elements from both variables.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=16)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE", "v1/.ATTRIBUTES/VARIABLE_VALUE"},
{"v1/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
# V0
slice_spec = V0SaveSliceInfo(var_offset=[0, 0], var_shape=[2, 2]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [[0, 1], [2, 3]])
slice_spec = V0SaveSliceInfo(var_offset=[2, 0], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [[4, 5]])
# V1
slice_spec = V1SaveSliceInfo(var_offset=[0, 0, 0], var_shape=[1, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v1_name][slice_spec]), [[[6.0], [7.0]]])
slice_spec = V1SaveSliceInfo(var_offset=[1, 0, 0], var_shape=[2, 2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v1_name][slice_spec]),
[[[8.0], [9.0]], [[10.0], [11.0]]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_Strings(self):
v_strings = [
"".join(random.choices(string.ascii_uppercase + string.digits, k=10))
for _ in range(4)]
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(v_strings, name="v0",
dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=v0_name,
full_shape=tensor_shape.TensorShape(dims=[4]),
var_offset=var_offset,
var_shape=var_shape)
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the v0 & v1 tensors with different max shard sizes.
# max_shard_size: 10 bytes
# Each string in v0 is 10 bytes, so there should be 1 string per shard.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=10)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE", "_CHECKPOINTABLE_OBJECT_GRAPH",}
])
slice_spec = V0SaveSliceInfo(var_offset=[0], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[0][v0_name][slice_spec]), [v_strings[0]])
slice_spec = V0SaveSliceInfo(var_offset=[1], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[1][v0_name][slice_spec]), [v_strings[1]])
slice_spec = V0SaveSliceInfo(var_offset=[2], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[2][v0_name][slice_spec]), [v_strings[2]])
slice_spec = V0SaveSliceInfo(var_offset=[3], var_shape=[1]).spec
self.assertAllEqual(
self.evaluate(shards[3][v0_name][slice_spec]), [v_strings[3]])
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_LargeScalar(self):
v_string = "".join(random.choices(
string.ascii_uppercase + string.digits, k=10)).encode("utf-8")
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(
v_string, name="v0", dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
v0_name = "v0/.ATTRIBUTES/VARIABLE_VALUE"
shardable_tensors = self._get_shardable_tensors_by_task(root)
# max_shard_size: 8 bytes
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"_CHECKPOINTABLE_OBJECT_GRAPH",},
{"v0/.ATTRIBUTES/VARIABLE_VALUE",}
])
tensor_val = (self.evaluate(shards[1][v0_name][""])
if ops.context.executing_eagerly()
else shards[1][v0_name][""])
self.assertEqual(tensor_val, v_string)
@test_util.run_in_graph_and_eager_modes
def test_CheckpointOption_MaxShardSizePolicy(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0, 1],
[2, 3],
[4, 5]],
name="v0")
v1 = resource_variable_ops.ResourceVariable([[[6.0], [7.0]],
[[8.0], [9.0]],
[[10.0], [11.0]]], name="v1")
v2 = resource_variable_ops.ResourceVariable("test_string", name="v1")
self.evaluate(v0.initializer)
self.evaluate(v1.initializer)
self.evaluate(v2.initializer)
root.v0 = v0
root.v1 = v1
root.v2 = v2
tmp_dir = self.create_tempdir("ckpt")
ckpt = checkpoint.Checkpoint(root)
save_path = ckpt.save(
tmp_dir, options=checkpoint_options.CheckpointOptions(
experimental_sharding_callback=(
sharding_policies.MaxShardSizePolicy(max_shard_size=10))))
# 8 files = 3 shards for v0, 3 for v1, 1 for v2, and 1 for the object graph
self.assertLen(gfile.Glob(save_path + ".data*"), 8)
ckpt.restore(save_path)
@test_util.run_in_graph_and_eager_modes
def test_MaxShardSizePolicy_PreSlicedTensor(self):
root = module.Module()
sliced_v0_name = "sliced_v0/.ATTRIBUTES/VARIABLE_VALUE"
class V0SaveSliceInfo(variables.Variable.SaveSliceInfo):
def __init__(self, var_offset, var_shape):
super().__init__(
full_name=sliced_v0_name,
full_shape=tensor_shape.TensorShape(dims=[2, 5]),
var_offset=var_offset,
var_shape=var_shape)
v0_slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[2, 3])
class ResourceVariableWithSliceSpec(resource_variable_ops.ResourceVariable):
def _serialize_to_tensors(self):
ckpt_key, tensor = list(super()._serialize_to_tensors().items())[0]
return {ckpt_key: {v0_slice_spec.spec: tensor}}
with ops.device("cpu:0"):
# full_v0 = [[0.0, 1.0, 2.0, 3.0, 4.0],
# [5.0, 6.0, 7.0, 8.0, 9.0]]
sliced_v0 = ResourceVariableWithSliceSpec([[1.0, 2.0, 3.0],
[6.0, 7.0, 8.0]],
name="sliced_v0",
dtype=dtypes.float32)
sliced_v0._set_save_slice_info(v0_slice_spec)
self.evaluate(sliced_v0.initializer)
root.sliced_v0 = sliced_v0
shardable_tensors = self._get_shardable_tensors_by_task(root)
# Test sharding the pre-sliced v0 tensor with different max shard sizes.
# max_shard_size: 8 bytes
# Each element of v0 is a 32 bit/4 byte value, so v0 should be split into 3
# shards containing 2 elements each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=8)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH",},
])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[2, 1]).spec
self.assertAllEqual(
self.evaluate(shards[0][sliced_v0_name][slice_spec]), [[1.0], [6.0]])
slice_spec = V0SaveSliceInfo(var_offset=[0, 2], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[1][sliced_v0_name][slice_spec]), [[2.0, 3.0]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 2], var_shape=[1, 2]).spec
self.assertAllEqual(
self.evaluate(shards[2][sliced_v0_name][slice_spec]), [[7.0, 8.0]])
# max_shard_size: 12 bytes
# Each element of v0 is a 32 bit/4 byte value, so v0 should be split into 2
# shards containing 3 elements each.
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=12)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
self.assertEqual(
[set(shard.keys()) for shard in shards],
[
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",},
{"sliced_v0/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH",},
])
slice_spec = V0SaveSliceInfo(var_offset=[0, 1], var_shape=[1, 3]).spec
self.assertAllEqual(
self.evaluate(shards[0][sliced_v0_name][slice_spec]), [[1.0, 2.0, 3.0]])
slice_spec = V0SaveSliceInfo(var_offset=[1, 1], var_shape=[1, 3]).spec
self.assertAllEqual(
self.evaluate(shards[1][sliced_v0_name][slice_spec]), [[6.0, 7.0, 8.0]])
def test_MaxShardSizePolicy_TFFunction(self):
v_string = "".join(random.choices(
string.ascii_uppercase + string.digits, k=10)).encode("utf-8")
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(
v_string, name="v0", dtype=dtypes.string)
self.evaluate(v0.initializer)
root.v0 = v0
shardable_tensors = self._get_shardable_tensors_by_task(root)
@def_function.function
def wrapped_policy(shardable_tensors):
callback = sharding_policies.MaxShardSizePolicy(max_shard_size=4)
shards = []
for tensors in shardable_tensors:
shards.extend(callback(tensors))
return shards
# TODO(b/326287351): Get string tensor size in tf.function.
# This test case should be changed when the bug is fixed/warning removed.
with self.assertLogs(level="WARNING") as log_output:
log_level = logging.get_verbosity()
logging.set_verbosity(logging.WARNING)
try:
wrapped_policy(shardable_tensors)
finally:
logging.set_verbosity(log_level)
output = log_output[0][0].message
self.assertTrue(
re.search("sharding policy is being executed in a tf.function", output))
if __name__ == "__main__":
ops.enable_eager_execution()
test.main()
tensorflow/python/checkpoint/sharding/sharding_util.py
+288
View File
@@ -0,0 +1,288 @@
# Copyright 2023 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.
# ==============================================================================
"""Data structures and utilities for checkpoint sharding."""
import abc
import dataclasses
import inspect
from typing import Hashable, MutableMapping, Sequence
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.framework import tensor_shape
from tensorflow.python.framework import tensor_spec
from tensorflow.python.ops import variables
from tensorflow.python.trackable import base
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.util import tf_export
TensorSlices = MutableMapping[tensor_spec.TensorSpec, tensor_lib.Tensor]
# A mapping from a checkpoint key (full tensor name) to the corresponding tensor
# slices of the full tensor. It represents the collection of tensors stored in a
# checkpoint shard data file.
Shard = MutableMapping[str, TensorSlices]
@tf_export.tf_export("train.experimental.ShardableTensor")
@dataclasses.dataclass(frozen=True)
class ShardableTensor:
"""Tensor wrapper containing data necessary for sharding.
The tensor representation used as inputs to pre-made and custom
`tf.train.experiemental.ShardingCallback`s, which can be specified using the
`experimental_sharding_callback` option in `tf.train.CheckpointOptions`.
"""
_tensor_save_spec: saveable_object.SaveSpec
tensor: tensor_lib.Tensor
dtype: dtypes.DType
device: device_lib.DeviceSpec
name: str
shape: tensor_shape.TensorShape
slice_spec: variables.Variable.SaveSliceInfo
checkpoint_key: str
trackable: base.Trackable
def __hash__(self) -> int:
return hash((self.name, self.dtype, str(self.device), self.checkpoint_key))
def __repr__(self) -> str:
return (f"\n{self.__class__.__name__}:\n"
f" _tensor_save_spec={self._tensor_save_spec!r}\n"
f" tensor={self.tensor!r}\n"
f" dtype={self.dtype!r}\n"
f" device={self.device!r}\n"
f" name={self.name!r}\n"
f" shape={self.shape!r}\n"
f" slice_spec={self.slice_spec!r}\n"
f" checkpoint_key={self.checkpoint_key!r}\n"
f" trackable={self.trackable!r}")
@tf_export.tf_export("train.experimental.ShardingCallback")
class ShardingCallback(abc.ABC):
"""Checkpoint sharding callback function, along with a text description.
A callback function wrapper that will be executed to determine how tensors
will be split into shards when the saver writes the checkpoint shards to disk.
The callback takes a list of `tf.train.experimental.ShardableTensor`s as input
(as well as any kwargs defined by the `tf.train.experimental.ShardingCallback`
subclass), and organizes the input tensors into different shards. Tensors are
first organized by device task (see `tf.DeviceSpec`), then the callback will
be called for each collection of tensors.
There are a few restrictions to keep in mind when creating a custom callback:
- Tensors must not be removed from the checkpoint.
- Tensors must not be reshaped.
- Tensor dtypes must not change.
- Tensors within a shard must belong to the same task.
Validation checks will be performed after the callback function is executed to
ensure these restrictions aren't violated.
Here's an example of a simple custom callback:
```
# Place all tensors in a single shard.
class AllInOnePolicy(tf.train.experimental.ShardingCallback):
@property
def description(self):
return "Place all tensors in a single shard."
def __call__(self, shardable_tensors):
tensors = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor_save_spec.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
tensors.set_default(checkpoint_key, {})[slice_spec] = tensor
return [tensors]
ckpt.save(
"path",
options=tf.train.CheckpointOptions(
experimental_sharding_callback=AllInOnePolicy()))
```
The `description` attribute is used to identify the callback and to aid
debugging during saving and restoration.
To take in kwargs, simply define the constructor and pass them in:
```
class ParameterPolicy(tf.train.experimental.ShardingCallback):
def __init__(self, custom_param):
self.custom_param = custom_param
...
ckpt.save(
"path",
options=tf.train.CheckpointOptions(
experimental_sharding_callback=ParameterPolicy(custom_param=...)))
```
"""
@property
@abc.abstractmethod
def description(self) -> str:
"""Returns a text description of the sharding policy."""
pass
@abc.abstractmethod
def __call__(
self, shardable_tensors: Sequence[ShardableTensor]
) -> Sequence[Shard]:
"""Returns a list of shards for the given shardable tensors."""
pass
def __hash__(self) -> int:
hash_val = hash(self.description)
# vars() only includes user-defined attributes.
for attr_name, attr_val in vars(self).items():
if not (inspect.ismethod(attr_val) or inspect.isfunction(attr_val)):
hash_val ^= hash(attr_name)
if isinstance(attr_val, Hashable):
hash_val ^= hash(attr_val)
return hash_val
def validate_shards(
shards: Sequence[Shard],
shardable_tensors: Sequence[ShardableTensor],
callback_description: str
) -> None:
"""Validates shards generated by the sharding_callback."""
unseen_tensor_dict = {}
for shardable_tensor in shardable_tensors:
unseen_tensor_dict.setdefault(
shardable_tensor.checkpoint_key, {}
)[shardable_tensor.slice_spec] = shardable_tensor.tensor
seen_tensor_set = set()
for shard_tensors in shards:
task_tensor = None
for checkpoint_key, tensor_slice_dict in shard_tensors.items():
for slice_spec, shard_tensor in tensor_slice_dict.items():
slice_spec = slice_spec.strip()
# Validate uniqueness.
if (checkpoint_key, slice_spec) in seen_tensor_set:
raise RuntimeError(
"After executing the checkpoint sharding callback, multiple "
"tensors with the same checkpoint key and slice spec were "
"found:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n")
# Validate no added tensors.
if checkpoint_key not in unseen_tensor_dict:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"not originally in the object graph was found in the "
"checkpoint shards:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n")
# Validate no shape change.
target_shape = unseen_tensor_dict[checkpoint_key][slice_spec].shape
if shard_tensor.shape != target_shape:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered shape:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_shape: {target_shape}\n"
f" new tensor_shape: {shard_tensor.shape}\n")
# Validate no dtype change.
target_dtype = unseen_tensor_dict[checkpoint_key][slice_spec].dtype
if shard_tensor.dtype != target_dtype:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered dtype:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_dtype: {target_dtype}\n"
f" new tensor_dtype: {shard_tensor.dtype}\n")
# Validate no task change.
target_task = device_lib.DeviceSpec.from_string(
unseen_tensor_dict[checkpoint_key][slice_spec].device).task
shard_tensor_task = device_lib.DeviceSpec.from_string(
shard_tensor.device).task
if shard_tensor_task != target_task:
raise RuntimeError(
"After executing the checkpoint sharding callback, a tensor "
"was found with an altered task:\n"
f" callback_description: {callback_description}\n"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" original tensor_task: {target_task}\n"
f" new tensor_task: {shard_tensor_task}\n")
# Validate tensors in shard have the same task.
if task_tensor is None:
task_tensor = ShardableTensor(
_tensor_save_spec=None,
tensor=None,
dtype=None,
device=shard_tensor.device,
name=None,
shape=None,
slice_spec=slice_spec,
checkpoint_key=checkpoint_key,
trackable=None)
else:
task1 = device_lib.DeviceSpec.from_string(task_tensor.device).task
task2 = device_lib.DeviceSpec.from_string(shard_tensor.device).task
if task1 is not None and task2 is not None and task1 != task2:
raise RuntimeError(
"After executing the checkpoint sharding callback, tensors "
"with different tasks were found in the same shard:\n"
f" callback_description: {callback_description}\n"
" tensor #1:"
f" checkpoint_key: {task_tensor.checkpoint_key}\n"
f" slice_spec: {task_tensor.slice_spec}\n"
f" task: {task1}\n"
" tensor #2:"
f" checkpoint_key: {checkpoint_key}\n"
f" slice_spec: {slice_spec}\n"
f" task: {task2}\n")
del unseen_tensor_dict[checkpoint_key][slice_spec]
if not unseen_tensor_dict[checkpoint_key]:
del unseen_tensor_dict[checkpoint_key]
seen_tensor_set.add((checkpoint_key, slice_spec))
# validate no tensor removal
if unseen_tensor_dict:
tensors_info = ""
for ckpt_key, slice_spec in unseen_tensor_dict.items():
tensors_info += " tensor:\n"
tensors_info += f" checkpoint_key: {ckpt_key}\n"
tensors_info += f" slice_spec: {slice_spec}\n"
raise RuntimeError(
"After executing the checkpoint sharding callback, tensors in the "
"object graph were not found in the checkpoint shards:\n"
f" callback_description: {callback_description}\n"
f"{tensors_info}")
tensorflow/python/checkpoint/sharding/sharding_util_test.py
+430
View File
@@ -0,0 +1,430 @@
# Copyright 2023 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 checkpoint sharding structures and utilities."""
from typing import Sequence
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import graph_view
from tensorflow.python.checkpoint.sharding import sharding_policies
from tensorflow.python.checkpoint.sharding import sharding_util
from tensorflow.python.eager import remote
from tensorflow.python.eager import test
from tensorflow.python.framework import device as device_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor as tensor_lib
from tensorflow.python.module import module
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.training import server_lib
from tensorflow.python.training.saving import saveable_object
from tensorflow.python.training.saving import saveable_object_util
class ShardingUtilTest(test.TestCase):
def _get_shardable_tensors_by_task(self, root):
serialized_tensors, _, _, _ = (
checkpoint.TrackableSaver(graph_view.ObjectGraphView(root))
._gather_serialized_tensors(None))
shardable_tensors_by_task = {}
for obj, tensor_dict in serialized_tensors.items():
for checkpoint_key, tensor_slice_dict in tensor_dict.items():
if not isinstance(tensor_slice_dict, dict):
# Make sure that maybe_tensor is structured as {slice_spec -> tensor}.
tensor_slice_dict = {"": tensor_slice_dict}
for slice_spec, tensor_save_spec in tensor_slice_dict.items():
if not isinstance(tensor_save_spec, saveable_object.SaveSpec):
tensor_save_spec = saveable_object.SaveSpec(
tensor=tensor_save_spec,
slice_spec=slice_spec,
name=checkpoint_key,
dtype=tensor_save_spec.dtype,
device=tensor_save_spec.device)
save_spec_tensor = tensor_save_spec.tensor
device = (device_lib.DeviceSpec.from_string(tensor_save_spec.device)
if isinstance(tensor_save_spec.device, str)
else tensor_save_spec.device)
task = device_lib.DeviceSpec.from_string(
saveable_object_util.set_cpu0(device.to_string()))
shardable_tensors_by_task.setdefault(task, []).append(
sharding_util.ShardableTensor(
_tensor_save_spec=tensor_save_spec,
tensor=save_spec_tensor,
dtype=tensor_save_spec.dtype,
device=device,
name=tensor_save_spec.name,
shape=save_spec_tensor.shape,
slice_spec=slice_spec.strip(),
checkpoint_key=checkpoint_key,
trackable=obj))
return shardable_tensors_by_task.values()
def test_hash_ShardingCallback(self):
class BlankCallback(sharding_util.ShardingCallback):
@property
def description(self):
return ""
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
pass
self.assertEqual(hash(BlankCallback()), hash(BlankCallback()))
class ValueCallback(sharding_util.ShardingCallback):
def __init__(self, val):
self.val = val
@property
def description(self):
return "value callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
pass
self.assertEqual(hash(ValueCallback(1)), hash(ValueCallback(1)))
self.assertNotEqual(hash(ValueCallback(1)), hash(ValueCallback(2)))
def test_validate_shards_correct(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
with ops.device("cpu:2"):
v2 = resource_variable_ops.ResourceVariable(2.0, name="v2")
root.v0 = v0
root.v1 = v1
root.v2 = v2
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = sharding_policies.ShardByTaskPolicy()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
self.assertEqual(
[list(shard.keys()) for shard in shards],
[[
"v0/.ATTRIBUTES/VARIABLE_VALUE",
"v1/.ATTRIBUTES/VARIABLE_VALUE",
"v2/.ATTRIBUTES/VARIABLE_VALUE",
"_CHECKPOINTABLE_OBJECT_GRAPH"
]])
self.assertEqual(
shards[0]["v0/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v0.numpy())
self.assertEqual(
shards[0]["v1/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v1.numpy())
self.assertEqual(
shards[0]["v2/.ATTRIBUTES/VARIABLE_VALUE"][""].numpy(),
v2.numpy())
def test_validate_shards_duplicate_tensor(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
with ops.device("cpu:1"):
v1 = resource_variable_ops.ResourceVariable(1.0, name="v1")
root.v0 = v0
root.v1 = v1
class DuplicateTensorCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "duplicate tensor callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
tensor = shardable_tensors[0].tensor
checkpoint_key = shardable_tensors[0].checkpoint_key
slice_spec = shardable_tensors[0].slice_spec
shards = [
{checkpoint_key: {slice_spec: tensor}},
{checkpoint_key: {slice_spec: tensor}}
]
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DuplicateTensorCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"multiple tensors with the same checkpoint "
"key and slice spec were found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_added_tensor(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class AddedTensorCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "added tensor callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
checkpoint_key = "ADDED_TENSOR_ABC123"
slice_spec = ""
tensor = tensor_lib.Tensor()
return [{checkpoint_key: {slice_spec: tensor}}]
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = AddedTensorCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor not originally in the object graph"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_shape_change(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable([[0.0, 1.0]], name="v0")
root.v0 = v0
class ShapeChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "shape change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
tensor = array_ops.transpose(tensor)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = ShapeChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered shape"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_dtype_change(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class DtypeChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "dtype change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
tensor = math_ops.cast(tensor, dtype=dtypes.int32)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DtypeChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered dtype"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_task_change(self):
servers = [server_lib.Server.create_local_server() for _ in range(2)]
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")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(0.0, name="v1")
root.v0 = v0
root.v1 = v1
class TaskChangeCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "task change callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shards = []
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
if checkpoint_key == "v0/.ATTRIBUTES/VARIABLE_VALUE":
with ops.device("/job:worker/task:1/cpu:0"):
tensor = array_ops.identity(tensor)
shards.append({checkpoint_key: {slice_spec: tensor}})
return shards
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = TaskChangeCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"a tensor was found with an altered task"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_different_tasks(self):
servers = [server_lib.Server.create_local_server() for _ in range(2)]
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")
with ops.device("/job:worker/task:1/cpu:0"):
v1 = resource_variable_ops.ResourceVariable(0.0, name="v1")
root.v0 = v0
root.v1 = v1
class DifferentTasksCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "different tasks callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
shard = {}
for shardable_tensor in shardable_tensors:
tensor = shardable_tensor.tensor
checkpoint_key = shardable_tensor.checkpoint_key
slice_spec = shardable_tensor.slice_spec
shard.setdefault(checkpoint_key, {})[slice_spec] = tensor
return [shard]
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = DifferentTasksCallback()
shards = sharding_callback(shardable_tensors_flat)
with self.assertRaisesRegex(RuntimeError,
"tensors with different tasks were found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
def test_validate_shards_tensor_removal(self):
root = module.Module()
with ops.device("cpu:0"):
v0 = resource_variable_ops.ResourceVariable(0.0, name="v0")
root.v0 = v0
class TensorRemovalCallback(sharding_util.ShardingCallback):
@property
def description(self):
return "tensor removal callback"
def __call__(
self, shardable_tensors: Sequence[sharding_util.ShardableTensor]
) -> Sequence[sharding_util.Shard]:
return []
shardable_tensors = self._get_shardable_tensors_by_task(root)
shardable_tensors_flat = []
for tensors in shardable_tensors:
shardable_tensors_flat.extend(tensors)
sharding_callback = TensorRemovalCallback()
shards = []
for tensors in shardable_tensors:
shards.extend(sharding_callback(tensors))
with self.assertRaisesRegex(RuntimeError,
"tensors in the object graph were not found"):
sharding_util.validate_shards(
shards, shardable_tensors_flat, sharding_callback.description)
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/tensor_callable.py
+41
View File
@@ -0,0 +1,41 @@
# Copyright 2022 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.
# ==============================================================================
"""`Callable` class used for checkpointing."""
from tensorflow.python.training.saving import saveable_object
class Callable(saveable_object.SaveSpec):
"""A callable that represents a Tensor that should be saved to checkpoint.
This can be returned from `_serialize_to_tensor` in place of a Tensor. The
callable will be executed on the specified device when the checkpoint is
about to be written.
Any class can use `Callable` for checkpointing, but for SavedModel export,
only resource-type variables* are supported.
* `resource_variable_ops.is_resource_variable(obj)` must return True.
"""
def __init__(self, tensor_callable, dtype, device):
"""Initializes a `Callable` object.
Args:
tensor_callable: A callable that takes no arguments and returns a Tensor.
dtype: Dtype of the tensor returned by the callable.
device: Device of the tensor returned by the callable.
"""
super().__init__(tensor_callable, None, None, dtype, device)
tensorflow/python/checkpoint/tensor_callable_test.py
+73
View File
@@ -0,0 +1,73 @@
# Copyright 2022 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 SaveableObject compatibility."""
import os
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.checkpoint import tensor_callable
from tensorflow.python.eager import test
from tensorflow.python.ops import variables
from tensorflow.python.saved_model import save as saved_model_save
from tensorflow.python.trackable import base
class IncrementWhenSave(base.Trackable):
def __init__(self):
self.read_counter = variables.Variable(0)
def _serialize_to_tensors(self):
def _get_and_increment_counter():
value = self.read_counter.read_value()
self.read_counter.assign_add(1)
return value
return {
"read_counter":
tensor_callable.Callable(_get_and_increment_counter,
self.read_counter.dtype,
self.read_counter.device)
}
def _restore_from_tensors(self, restored_tensors):
self.read_counter.assign(restored_tensors["read_counter"])
class CallableTest(test.TestCase):
def test_callable(self):
trackable = IncrementWhenSave()
ckpt = checkpoint.Checkpoint(attr=trackable)
prefix = os.path.join(self.get_temp_dir(), "ckpt")
save_path = ckpt.save(prefix)
self.assertEqual(1, self.evaluate(trackable.read_counter))
ckpt.save(prefix)
self.assertEqual(2, self.evaluate(trackable.read_counter))
ckpt.restore(save_path)
self.assertEqual(0, self.evaluate(trackable.read_counter))
def test_callable_saved_model_compatibility(self):
trackable = IncrementWhenSave()
trackable.read_counter.assign(15)
save_path = os.path.join(self.get_temp_dir(), "saved_model")
with self.assertRaisesRegex(NotImplementedError, "returns a Callable"):
saved_model_save.save(trackable, save_path)
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/testdata/generate_checkpoint.py
Vendored
+77
View File
@@ -0,0 +1,77 @@
# Copyright 2022 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.
# ==============================================================================
"""Standalone utility to generate some test saved models."""
from absl import app
from tensorflow.python.checkpoint import checkpoint
from tensorflow.python.compat import v2_compat
from tensorflow.python.framework import dtypes
from tensorflow.python.module import module
from tensorflow.python.ops import lookup_ops
from tensorflow.python.ops import variables
class TableModule(module.Module):
"""Three vars (one in a sub-module) and compute method."""
def __init__(self):
default_value = -1
empty_key = 0
deleted_key = -1
self.lookup_table = lookup_ops.DenseHashTable(
dtypes.int64,
dtypes.int64,
default_value=default_value,
empty_key=empty_key,
deleted_key=deleted_key,
name="t1",
initial_num_buckets=32)
self.lookup_table.insert(1, 1)
self.lookup_table.insert(2, 4)
class VariableModule(module.Module):
def __init__(self):
self.v = variables.Variable([1., 2., 3.])
self.w = variables.Variable([4., 5.])
MODULE_CTORS = {
"TableModule": TableModule,
"VariableModule": VariableModule,
}
def main(args):
if len(args) != 3:
print("Expected: {export_path} {ModuleName}")
print("Allowed ModuleNames:", MODULE_CTORS.keys())
return 1
_, export_path, module_name = args
module_ctor = MODULE_CTORS.get(module_name)
if not module_ctor:
print("Expected ModuleName to be one of:", MODULE_CTORS.keys())
return 2
tf_module = module_ctor()
ckpt = checkpoint.Checkpoint(tf_module)
ckpt.write(export_path)
if __name__ == "__main__":
v2_compat.enable_v2_behavior()
app.run(main)
tensorflow/python/checkpoint/testdata/table_legacy_saveable_object.data-00000-of-00001
Vendored
BIN
View File
Binary file not shown.
tensorflow/python/checkpoint/testdata/table_legacy_saveable_object.index
Vendored
BIN
View File
Binary file not shown.
tensorflow/python/checkpoint/trackable_view.py
+117
View File
@@ -0,0 +1,117 @@
"""Manages a Trackable object graph."""
# 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.
# ==============================================================================
import collections
import weakref
from tensorflow.python.trackable import base
from tensorflow.python.trackable import converter
from tensorflow.python.util import object_identity
from tensorflow.python.util.tf_export import tf_export
@tf_export("train.TrackableView", v1=[])
class TrackableView(object):
"""Gathers and serializes a trackable view.
Example usage:
>>> class SimpleModule(tf.Module):
... def __init__(self, name=None):
... super().__init__(name=name)
... self.a_var = tf.Variable(5.0)
... self.b_var = tf.Variable(4.0)
... self.vars = [tf.Variable(1.0), tf.Variable(2.0)]
>>> root = SimpleModule(name="root")
>>> root.leaf = SimpleModule(name="leaf")
>>> trackable_view = tf.train.TrackableView(root)
Pass root to tf.train.TrackableView.children() to get the dictionary of all
children directly linked to root by name.
>>> trackable_view_children = trackable_view.children(root)
>>> for item in trackable_view_children.items():
... print(item)
('a_var', <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=5.0>)
('b_var', <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=4.0>)
('vars', ListWrapper([<tf.Variable 'Variable:0' shape=() dtype=float32,
numpy=1.0>, <tf.Variable 'Variable:0' shape=() dtype=float32, numpy=2.0>]))
('leaf', ...)
"""
def __init__(self, root):
"""Configure the trackable view.
Args:
root: A `Trackable` object whose variables (including the variables of
dependencies, recursively) should be saved. May be a weak reference.
"""
# TrackableView should never contain a strong reference to root, since it
# may result in a cycle:
# root -> deferred dependencies -> CheckpointPosition
# -> CheckpointRestoreCoordinator -> TrackableView -> root
self._root_ref = (root if isinstance(root, weakref.ref)
else weakref.ref(root))
@classmethod
def children(cls, obj, save_type=base.SaveType.CHECKPOINT, **kwargs):
"""Returns all child trackables attached to obj.
Args:
obj: A `Trackable` object.
save_type: A string, can be 'savedmodel' or 'checkpoint'.
**kwargs: kwargs to use when retrieving the object's children.
Returns:
Dictionary of all children attached to the object with name to trackable.
"""
# pylint: disable=protected-access
obj._maybe_initialize_trackable()
children = {}
for name, ref in obj._trackable_children(save_type, **kwargs).items():
ref = converter.convert_to_trackable(ref, parent=obj)
children[name] = ref
return children
@property
def root(self):
if isinstance(self._root_ref, weakref.ref):
derefed = self._root_ref()
assert derefed is not None
return derefed
else:
return self._root_ref
def descendants(self):
"""Returns a list of all nodes from self.root using a breadth first traversal."""
return self._descendants_with_paths()[0]
def _descendants_with_paths(self):
"""Returns a list of all nodes and its paths from self.root using a breadth first traversal."""
bfs_sorted = []
to_visit = collections.deque([self.root])
node_paths = object_identity.ObjectIdentityDictionary()
node_paths[self.root] = ()
while to_visit:
current_trackable = to_visit.popleft()
bfs_sorted.append(current_trackable)
for name, dependency in self.children(current_trackable).items():
if dependency not in node_paths:
node_paths[dependency] = (
node_paths[current_trackable] +
(base.TrackableReference(name, dependency),))
to_visit.append(dependency)
return bfs_sorted, node_paths
tensorflow/python/checkpoint/trackable_view_test.py
+44
View File
@@ -0,0 +1,44 @@
# 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 the trackable view."""
from tensorflow.python.checkpoint import trackable_view
from tensorflow.python.eager import test
from tensorflow.python.trackable import base
class TrackableViewTest(test.TestCase):
def test_children(self):
root = base.Trackable()
leaf = base.Trackable()
root._track_trackable(leaf, name="leaf")
(current_name,
current_dependency), = trackable_view.TrackableView.children(root).items()
self.assertIs(leaf, current_dependency)
self.assertEqual("leaf", current_name)
def test_descendants(self):
root = base.Trackable()
leaf = base.Trackable()
root._track_trackable(leaf, name="leaf")
descendants = trackable_view.TrackableView(root).descendants()
self.assertIs(2, len(descendants))
self.assertIs(root, descendants[0])
self.assertIs(leaf, descendants[1])
if __name__ == "__main__":
test.main()
tensorflow/python/checkpoint/util.py
+183
View File
@@ -0,0 +1,183 @@
# Copyright 2022 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.
# ==============================================================================
"""Utilities for extracting and writing checkpoint info`."""
from tensorflow.core.protobuf import trackable_object_graph_pb2
from tensorflow.python.ops import resource_variable_ops
from tensorflow.python.ops import variables
from tensorflow.python.trackable import trackable_utils
from tensorflow.python.util import object_identity
def serialize_slot_variables(trackable_objects, node_ids, object_names):
"""Gather and name slot variables."""
non_slot_objects = list(trackable_objects)
slot_variables = object_identity.ObjectIdentityDictionary()
for trackable in non_slot_objects:
# TODO(b/110718070): Fix Keras imports.
# Note: dir() is used rather than hasattr() here to avoid triggering
# custom __getattr__ code, see b/152031870 for context.
if "get_slot_names" in dir(trackable):
slot_names = trackable.get_slot_names()
for slot_name in slot_names:
for original_variable_node_id, original_variable in enumerate(
non_slot_objects):
try:
slot_variable = trackable.get_slot(original_variable, slot_name)
except (AttributeError, KeyError):
slot_variable = None
if slot_variable is None:
continue
slot_variable._maybe_initialize_trackable() # pylint: disable=protected-access
if slot_variable._trackable_children(): # pylint: disable=protected-access
# TODO(allenl): Gather dependencies of slot variables.
raise NotImplementedError(
"Currently only variables with no dependencies can be saved as "
"slot variables. File a feature request if this limitation "
"bothers you.")
if slot_variable in node_ids:
raise NotImplementedError(
"A slot variable was re-used as a dependency of a Trackable "
f"object: {slot_variable}. This is not currently allowed. "
"File a feature request if this limitation bothers you.")
checkpoint_name = trackable_utils.slot_variable_key(
variable_path=object_names[original_variable],
optimizer_path=object_names[trackable],
slot_name=slot_name)
object_names[slot_variable] = checkpoint_name
slot_variable_node_id = len(trackable_objects)
node_ids[slot_variable] = slot_variable_node_id
trackable_objects.append(slot_variable)
slot_variable_proto = (
trackable_object_graph_pb2.TrackableObjectGraph.TrackableObject
.SlotVariableReference(
slot_name=slot_name,
original_variable_node_id=original_variable_node_id,
slot_variable_node_id=slot_variable_node_id))
slot_variables.setdefault(trackable, []).append(slot_variable_proto)
return slot_variables
def get_mapped_trackable(trackable, object_map):
"""Returns the mapped trackable if possible, otherwise returns trackable."""
if object_map is None:
return trackable
else:
return object_map.get(trackable, trackable)
def get_full_name(var):
"""Gets the full name of variable for name-based checkpoint compatibility."""
# pylint: disable=protected-access
if (not (isinstance(var, variables.Variable) or
# Some objects do not subclass Variable but still act as one.
resource_variable_ops.is_resource_variable(var))):
return ""
if getattr(var, "_save_slice_info", None) is not None:
# Use getattr because `var._save_slice_info` may be set as `None`.
return var._save_slice_info.full_name
else:
return var._shared_name
# pylint: enable=protected-access
def add_checkpoint_values_check(object_graph_proto):
"""Determines which objects have checkpoint values and save this to the proto.
Args:
object_graph_proto: A `TrackableObjectGraph` proto.
"""
# Trackable -> set of all trackables that depend on it (the "parents").
# If a trackable has checkpoint values, then all of the parents can be
# marked as having checkpoint values.
parents = {}
checkpointed_trackables = object_identity.ObjectIdentitySet()
# First pass: build dictionary of parent objects and initial set of
# checkpointed trackables.
checkpointed_trackables = set()
for node_id, object_proto in enumerate(object_graph_proto.nodes):
if (object_proto.attributes or object_proto.slot_variables or
object_proto.HasField("registered_saver")):
checkpointed_trackables.add(node_id)
for child_proto in object_proto.children:
child = child_proto.node_id
if child not in parents:
parents[child] = set()
parents[child].add(node_id)
# Second pass: add all connected parents to set of checkpointed trackables.
to_visit = set()
to_visit.update(checkpointed_trackables)
while to_visit:
trackable = to_visit.pop()
if trackable not in parents:
# Some trackables may not have parents (e.g. slot variables).
continue
current_parents = parents.pop(trackable)
checkpointed_trackables.update(current_parents)
for parent in current_parents:
if parent in parents:
to_visit.add(parent)
for node_id, object_proto in enumerate(object_graph_proto.nodes):
object_proto.has_checkpoint_values.value = bool(
node_id in checkpointed_trackables)
def objects_ids_and_slot_variables_and_paths(graph_view,
skip_slot_variables=False):
"""Traverse the object graph and list all accessible objects.
Looks for `Trackable` objects which are dependencies of
`root_trackable`. Includes slot variables only if the variable they are
slotting for and the optimizer are dependencies of `root_trackable`
(i.e. if they would be saved with a checkpoint).
Args:
graph_view: A GraphView object.
skip_slot_variables: If True does not return trackables for slot variable.
Default False.
Returns:
A tuple of (trackable objects, paths from root for each object,
object -> node id, slot variables, object_names)
"""
trackable_objects, node_paths = graph_view.breadth_first_traversal()
object_names = object_identity.ObjectIdentityDictionary()
for obj, path in node_paths.items():
object_names[obj] = trackable_utils.object_path_to_string(path)
node_ids = object_identity.ObjectIdentityDictionary()
for node_id, node in enumerate(trackable_objects):
node_ids[node] = node_id
if skip_slot_variables:
slot_variables = object_identity.ObjectIdentityDictionary()
else:
slot_variables = serialize_slot_variables(
trackable_objects=trackable_objects,
node_ids=node_ids,
object_names=object_names,
)
return (trackable_objects, node_paths, node_ids, slot_variables, object_names)
def list_objects(graph_view, skip_slot_variables=False):
"""Traverse the object graph and list all accessible objects."""
trackable_objects = objects_ids_and_slot_variables_and_paths(
graph_view, skip_slot_variables
)[0]
return trackable_objects