• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Description:
2#   Low-level utilities for reading and writing checkpoints.
3
4load("//tensorflow:tensorflow.bzl", "cuda_py_test")
5
6package(
7    default_visibility = [
8        "//tensorflow:internal",
9    ],
10    licenses = ["notice"],
11)
12
13py_library(
14    name = "checkpoint_options",
15    srcs = ["checkpoint_options.py"],
16    srcs_version = "PY3",
17    deps = [
18        "//tensorflow/python/util:tf_export",
19    ],
20)
21
22py_library(
23    name = "functional_saver",
24    srcs = ["functional_saver.py"],
25    srcs_version = "PY3",
26    deps = [
27        ":checkpoint_options",
28        ":saveable_hook",
29        ":saveable_object",
30        ":saveable_object_util",
31        "//tensorflow/python/eager:def_function",
32    ],
33)
34
35cuda_py_test(
36    name = "functional_saver_test",
37    size = "medium",
38    srcs = [
39        "functional_saver_test.py",
40    ],
41    deps = [
42        ":checkpoint_options",
43        ":functional_saver",
44        ":saveable_hook",
45        "//tensorflow/python/eager:remote",
46        "//tensorflow/python/eager:test",
47    ],
48)
49
50py_library(
51    name = "saveable_object",
52    srcs = ["saveable_object.py"],
53    srcs_version = "PY3",
54)
55
56py_library(
57    name = "saveable_hook",
58    srcs = ["saveable_hook.py"],
59    srcs_version = "PY3",
60    deps = [
61        "//tensorflow/python:constant_op",
62        "//tensorflow/python/training/tracking:base",
63    ],
64)
65
66py_library(
67    name = "saveable_object_util",
68    srcs = ["saveable_object_util.py"],
69    srcs_version = "PY3",
70    deps = [
71        "//tensorflow/python:resource_variable_ops",
72        "//tensorflow/python:variables",
73        "//tensorflow/python/training/tracking:base",
74        "@six_archive//:six",
75    ],
76)
77