wehub-resource-sync
607 lines
19 KiB
Python
607 lines
19 KiB
Python
# 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",
|
|
],
|
|
)
|