Start moving Checkpointable utilities toward core

Doesn't add to the public API yet, just shifts code around. Changes:
  - A tiny bit of renaming (to avoid having _Checkpoint and Checkpoint in the same file)
  - Removed the garbage collection decorator from a few tests due to the uuid4() garbage issue (apparently core tests get run on Python 2.7.9?)
  - Renamed "Object" to "CheckpointableObject" in the proto, since core protos have Java bindings and apparently Java had something else in mind for the keyword "Object" :)
but otherwise this is a pure move.

After this CL I'll propose adding tf.train.Checkpoint to the API (currently tf.contrib.eager.Checkpoint), move the utilities that are still in contrib/eager to their own contrib directory (there will be a few more misc. utilities for inspecting checkpoints and managing dependencies), get tf.train.Saver to read object-based checkpoints for compatibility, and work on Model.save_weights/load_weights.

PiperOrigin-RevId: 192646890

1 files changed, 55 insertions, 0 deletions

diff --git a/tensorflow/core/protobuf/checkpointable_object_graph.proto b/tensorflow/core/protobuf/checkpointable_object_graph.proto
new file mode 100644
index 0000000000..651f692f6d
--- /dev/null
+++ b/tensorflow/core/protobuf/checkpointable_object_graph.proto
@@ -0,0 +1,55 @@
+syntax = "proto3";
+
+option cc_enable_arenas = true;
+
+package tensorflow;
+
+// A TensorBundle addition which saves extra information about the objects which
+// own variables, allowing for more robust checkpoint loading into modified
+// programs.
+
+message CheckpointableObjectGraph {
+  message CheckpointableObject {
+    message ObjectReference {
+      // An index into `CheckpointableObjectGraph.nodes`, indicating the object
+      // being referenced.
+      int32 node_id = 1;
+      // A user-provided name for the edge.
+      string local_name = 2;
+    }
+
+    message SerializedTensor {
+      // A name for the Tensor. Simple variables have only one
+      // `SerializedTensor` named "VARIABLE_VALUE" by convention. This value may
+      // be restored on object creation as an optimization.
+      string name = 1;
+      // The full name of the variable/tensor, if applicable. Used to allow
+      // name-based loading of checkpoints which were saved using an
+      // object-based API. Should match the checkpoint key which would have been
+      // assigned by tf.train.Saver.
+      string full_name = 2;
+      // The generated name of the Tensor in the checkpoint.
+      string checkpoint_key = 3;
+    }
+
+    message SlotVariableReference {
+      // An index into `CheckpointableObjectGraph.nodes`, indicating the
+      // variable object this slot was created for.
+      int32 original_variable_node_id = 1;
+      // The name of the slot (e.g. "m"/"v").
+      string slot_name = 2;
+      // An index into `CheckpointableObjectGraph.nodes`, indicating the
+      // `Object` with the value of the slot variable.
+      int32 slot_variable_node_id = 3;
+    }
+
+    // Objects which this object depends on.
+    repeated ObjectReference children = 1;
+    // Serialized data specific to this object.
+    repeated SerializedTensor attributes = 2;
+    // Slot variables owned by this object.
+    repeated SlotVariableReference slot_variables = 3;
+  }
+
+  repeated CheckpointableObject nodes = 1;
+}