Added a format for saving an inference graph that can be memmapped and an utility to convert a freezed graph into this format.

Change: 120128412

1 files changed, 150 insertions, 0 deletions

diff --git a/tensorflow/core/util/memmapped_file_system_test.cc b/tensorflow/core/util/memmapped_file_system_test.cc
new file mode 100644
index 0000000000..67a06824cd
--- /dev/null
+++ b/tensorflow/core/util/memmapped_file_system_test.cc
@@ -0,0 +1,150 @@
+/* Copyright 2016 Google Inc. 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.
+==============================================================================*/
+#include "tensorflow/core/util/memmapped_file_system.h"
+
+#include "tensorflow/core/framework/tensor_testutil.h"
+#include "tensorflow/core/graph/graph_def_builder.h"
+#include "tensorflow/core/lib/core/status_test_util.h"
+#include "tensorflow/core/lib/io/path.h"
+#include "tensorflow/core/platform/test.h"
+#include "tensorflow/core/util/memmapped_file_system_writer.h"
+
+namespace tensorflow {
+
+namespace {
+
+// Names of files in memmapped environment.
+constexpr char kTensor1FileName[] = "memmapped_package://t1";
+constexpr char kTensor2FileName[] = "memmapped_package://t2";
+constexpr char kProtoFileName[] = "memmapped_package://b";
+constexpr int kTestGraphDefVersion = 666;
+
+Status CreateMemmappedFileSystemFile(const string& filename, bool corrupted,
+                                     Tensor* test_tensor) {
+  Env* env = Env::Default();
+  MemmappedFileSystemWriter writer;
+  TF_RETURN_IF_ERROR(writer.InitializeToFile(env, filename));
+
+  // Try to write a tensor and proto.
+  test::FillFn<float>(test_tensor,
+                      [](int i) { return static_cast<float>(i * i); });
+
+  TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor1FileName));
+
+  // Create a proto with some fields.
+  GraphDef graph_def;
+  graph_def.set_version(kTestGraphDefVersion);
+  TF_RETURN_IF_ERROR(writer.SaveProtobuf(graph_def, kProtoFileName));
+
+  // Save a tensor after the proto to check that alignment works.
+  test::FillFn<float>(test_tensor,
+                      [](int i) { return static_cast<float>(i * i * i); });
+  TF_RETURN_IF_ERROR(writer.SaveTensor(*test_tensor, kTensor2FileName));
+
+  if (!corrupted) {
+    // Flush and close the file.
+    TF_RETURN_IF_ERROR(writer.FlushAndClose());
+  }
+  return Status::OK();
+}
+
+TEST(MemmappedFileSystemTest, SimpleTest) {
+  const TensorShape test_tensor_shape = {10, 200};
+  Tensor test_tensor(DT_FLOAT, test_tensor_shape);
+  const string dir = testing::TmpDir();
+  const string filename = io::JoinPath(dir, "memmapped_env_test");
+  TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, false, &test_tensor));
+
+  // Check that we can memmap the created file.
+  MemmappedEnv memmapped_env(Env::Default());
+  TF_ASSERT_OK(memmapped_env.InitializeFromFile(filename));
+  // Try to load a proto from the file.
+  GraphDef test_graph_def;
+  TF_EXPECT_OK(
+      ReadBinaryProto(&memmapped_env, kProtoFileName, &test_graph_def));
+  EXPECT_EQ(kTestGraphDefVersion, test_graph_def.version());
+  // Check that we can correctly get a tensor memory.
+  ReadOnlyMemoryRegion* memory_region;
+  TF_ASSERT_OK(memmapped_env.NewReadOnlyMemoryRegionFromFile(kTensor2FileName,
+                                                             &memory_region));
+  std::unique_ptr<ReadOnlyMemoryRegion> mem_region_ptr(memory_region);
+  // The memory region can be bigger but not less than Tensor size.
+  ASSERT_GE(memory_region->length(), test_tensor.TotalBytes());
+  EXPECT_EQ(test_tensor.tensor_data(),
+            StringPiece(static_cast<const char*>(memory_region->data()),
+                        test_tensor.TotalBytes()));
+  // Check that GetFileSize works.
+  uint64 file_size = 0;
+  TF_ASSERT_OK(memmapped_env.GetFileSize(kTensor2FileName, &file_size));
+  EXPECT_EQ(test_tensor.TotalBytes(), file_size);
+
+  // Check that if file not found correct error message returned.
+  EXPECT_EQ(
+      error::NOT_FOUND,
+      memmapped_env.NewReadOnlyMemoryRegionFromFile("bla-bla", &memory_region)
+          .code());
+
+  // Check FileExists.
+  EXPECT_TRUE(memmapped_env.FileExists(kTensor2FileName));
+  EXPECT_FALSE(memmapped_env.FileExists("bla-bla-bla"));
+}
+
+TEST(MemmappedFileSystemTest, NotInitalized) {
+  MemmappedEnv memmapped_env(Env::Default());
+  ReadOnlyMemoryRegion* memory_region;
+  EXPECT_EQ(
+      error::FAILED_PRECONDITION,
+      memmapped_env
+          .NewReadOnlyMemoryRegionFromFile(kTensor1FileName, &memory_region)
+          .code());
+  RandomAccessFile* file;
+  EXPECT_EQ(error::FAILED_PRECONDITION,
+            memmapped_env.NewRandomAccessFile(kProtoFileName, &file).code());
+}
+
+TEST(MemmappedFileSystemTest, Corrupted) {
+  // Create a corrupted file (it is not closed it properly).
+  const TensorShape test_tensor_shape = {100, 200};
+  Tensor test_tensor(DT_FLOAT, test_tensor_shape);
+  const string dir = testing::TmpDir();
+  const string filename = io::JoinPath(dir, "memmapped_env_corrupted_test");
+  TF_ASSERT_OK(CreateMemmappedFileSystemFile(filename, true, &test_tensor));
+  MemmappedFileSystem memmapped_env;
+  ASSERT_NE(memmapped_env.InitializeFromFile(Env::Default(), filename),
+            Status::OK());
+}
+
+TEST(MemmappedFileSystemTest, ProxyToDefault) {
+  MemmappedEnv memmapped_env(Env::Default());
+  const string dir = testing::TmpDir();
+  const string filename = io::JoinPath(dir, "test_file");
+  // Check that we can create write and read ordinary file.
+  WritableFile* writable_file;
+  TF_ASSERT_OK(memmapped_env.NewAppendableFile(filename, &writable_file));
+  std::unique_ptr<WritableFile> writable_file_ptr(writable_file);
+  const string test_string = "bla-bla-bla";
+  TF_ASSERT_OK(writable_file->Append(test_string));
+  TF_ASSERT_OK(writable_file->Close());
+  uint64 file_length = 0;
+  TF_EXPECT_OK(memmapped_env.GetFileSize(filename, &file_length));
+  EXPECT_EQ(test_string.length(), file_length);
+  RandomAccessFile* random_access_file;
+  TF_ASSERT_OK(
+      memmapped_env.NewRandomAccessFile(filename, &random_access_file));
+  delete random_access_file;
+}
+
+}  // namespace
+}  // namespace tensorflow