tensor_bundle: tighten Lookup() semantics, add LookupDtypeShape().

Zongheng Yang · tensorflower-gardener · commit cae3713cd4a2 · 2016-10-07T12:33:03.000-07:00
This fixes a subtle allocation issue in CheckpointReader.
Change: 135503615
diff --git a/tensorflow/c/checkpoint_reader.cc b/tensorflow/c/checkpoint_reader.cc
@@ -81,9 +81,14 @@ void CheckpointReader::GetTensor(
   if (reader_ != nullptr) {
     status = reader_->GetTensor(name, out_tensor);
   } else {
-    std::unique_ptr<Tensor> tensor(new Tensor);
-    status = v2_reader_->Lookup(name, tensor.get());
-    if (status.ok()) std::swap(*out_tensor, tensor);
+    tensorflow::DataType dtype;
+    tensorflow::TensorShape shape;
+    status = v2_reader_->LookupDtypeAndShape(name, &dtype, &shape);
+    if (status.ok()) {
+      out_tensor->reset(new Tensor(dtype, shape));
+      status = v2_reader_->Lookup(name, out_tensor->get());
+      if (!status.ok()) out_tensor->reset();
+    }
   }
   if (!status.ok()) {
     Set_TF_Status_from_Status(out_status, status);
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle.cc
@@ -745,11 +745,10 @@ Status BundleReader::GetSliceValue(StringPiece full_tensor_key,
     // allocate_temp()?  Note that without major refactorings to Saver, it's
     // hard for the caller of the tensor bundle module to allocate these
     // precisely-shaped scratch storage.
-    // TODO(zongheng): implement an important optimization: if the stored slice
-    // is a subset of the to-restore slice, directly read the stored slice into
-    // the latter's already-allocated backing buffer.
 
     // Optimization for the common case: stored slice == to-restore slice.
+    // TODO(zongheng): also include the case where "slice_spec" is full ("-"),
+    // and "stored_slice" is logically full but contains actual extents.
     if (stored_slice == slice_spec) {
       VLOG(1) << "Optimized for common case: directly copying into "
                  "pre-allocated buffer; spec: "
@@ -800,14 +799,20 @@ bool BundleReader::Contains(StringPiece key) {
   return Valid() && (this->key() == key);
 }
 
-Status BundleReader::LookupTensorShape(StringPiece key, TensorShape* shape) {
+Status BundleReader::LookupDtypeAndShape(StringPiece key, DataType* dtype,
+                                         TensorShape* shape) {
   BundleEntryProto entry;
   TF_RETURN_IF_ERROR(GetBundleEntryProto(key, &entry));
-
+  *dtype = entry.dtype();
   *shape = TensorShape(entry.shape());
   return Status::OK();
 }
 
+Status BundleReader::LookupTensorShape(StringPiece key, TensorShape* shape) {
+  DataType ignored;
+  return LookupDtypeAndShape(key, &ignored, shape);
+}
+
 string BundleReader::DebugString() {
   // Format used below emulates that of TensorSliceReader::DebugString().
   string shape_str;
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle.h b/tensorflow/core/util/tensor_bundle/tensor_bundle.h
@@ -179,6 +179,11 @@ class BundleReader {
   // REQUIRES: status().ok()
   bool Contains(StringPiece key);
 
+  // Looks up the dtype and the shape of the tensor keyed by "key".
+  // REQUIRES: status().ok()
+  Status LookupDtypeAndShape(StringPiece key, DataType* dtype,
+                             TensorShape* shape) TF_MUST_USE_RESULT;
+
   // Looks up the shape of the tensor keyed by "key".
   // Clears "shape" if not found.
   // REQUIRES: status().ok()
@@ -188,14 +193,9 @@ class BundleReader {
   // Looks up the tensor keyed by "key".  If "key" refers to a partitioned
   // tensor, attempts to look up the full contents using all stored slices.
   //
-  // Out-tensor "val" can be either empty or initialized with a non-empty shape:
-  //
-  // * If empty, this function allocates an exactly-sized Tensor to hold the
-  //   contents found in this bundle.
-  //
-  // * If non-empty, caller is responsible for making sure "val" has the same
-  //   shape as the corresponding contents. This function directly uses the
-  //   buffer without extra allocation.
+  // Caller must make sure "val" has the same shape and dtype as the
+  // corresponding contents, so that its buffer can be filled without needing
+  // extra allocation.  These can be queried via "LookupDtypeAndShape()".
   //
   // On error, "val" may contain nonsense data.  Returns a NotFound error if
   // tensor keyed by "key" does not exist in this bundle.
diff --git a/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc b/tensorflow/core/util/tensor_bundle/tensor_bundle_test.cc
@@ -52,9 +52,11 @@ void Expect(BundleReader* reader, const string& key,
             const Tensor& expected_val) {
   // Tests for Contains().
   EXPECT_TRUE(reader->Contains(key));
-  // Tests for LookupTensorShape().
+  // Tests for LookupDtypeAndShape().
+  DataType dtype;
   TensorShape shape;
-  TF_ASSERT_OK(reader->LookupTensorShape(key, &shape));
+  TF_ASSERT_OK(reader->LookupDtypeAndShape(key, &dtype, &shape));
+  EXPECT_EQ(expected_val.dtype(), dtype);
   EXPECT_EQ(expected_val.shape(), shape);
   // Tests for Lookup(), checking tensor contents.
   Tensor val(expected_val.dtype(), shape);
