Extract the model loading logic as a ModelLoader class.

Also move the Validator::CheckModel() functionality to this new ModelLoader class.

PiperOrigin-RevId: 460166696

Author: hehhjiang

tensorflow/lite/experimental/acceleration/mini_benchmark/BUILD

@@ -454,9 +454,11 @@ cc_library(
     deps = [
         ":call",
         ":decode_jpeg",
+        ":model_loader",
         ":status_codes",
         "//tensorflow/lite:framework",
         "//tensorflow/lite:minimal_logging",
+        "//tensorflow/lite/c:common",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/experimental/acceleration/configuration:configuration_fbs",
         "//tensorflow/lite/experimental/acceleration/configuration:delegate_registry",
@@ -471,15 +473,19 @@ cc_library(
     hdrs = ["validator_runner.h"],
     deps = [
         ":fb_storage",
+        ":model_loader",
         ":runner",
         ":status_codes",
         ":validator",
+        "@flatbuffers",
         "//tensorflow/lite:minimal_logging",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/experimental/acceleration/configuration:configuration_fbs",
-        "//tensorflow/lite/nnapi/sl:nnapi_support_library",
+        # For NNAPI support library, the headears and source files are defined
+        # as two separate targets. We need to include both targets for NNAPI to
+        # be invoked.
+        "//tensorflow/lite/nnapi/sl:nnapi_support_library",  # buildcleaner: keep
         "//tensorflow/lite/nnapi/sl:nnapi_support_library_headers",
-        "@flatbuffers",
     ],
 )
 
@@ -493,17 +499,20 @@ cc_library(
     srcs = ["validator_runner_entrypoint.cc"],
     deps = [
         ":fb_storage",
-        ":runner",
+        ":model_loader",
         ":set_big_core_affinity_h",
         ":status_codes",
         ":validator",
         ":validator_runner",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers",
         "//tensorflow/lite/core/api",
         "//tensorflow/lite/experimental/acceleration/configuration:configuration_fbs",
-        "//tensorflow/lite/nnapi/sl:nnapi_support_library",
+        # For NNAPI support library, the headears and source files are defined
+        # as two separate targets. We need to include both targets for NNAPI to
+        # be invoked.
+        "//tensorflow/lite/nnapi/sl:nnapi_support_library",  # buildcleaner: keep
         "//tensorflow/lite/nnapi/sl:nnapi_support_library_headers",
-        "@com_google_absl//absl/strings",
-        "@flatbuffers",
     ],
 )
 
@@ -709,6 +718,32 @@ cc_test(
     ],
 )
 
+cc_library(
+    name = "model_loader",
+    srcs = ["model_loader.cc"],
+    hdrs = ["model_loader.h"],
+    deps = [
+        ":status_codes",
+        "//tensorflow/lite:allocation",
+        "//tensorflow/lite:model_builder",
+        "//tensorflow/lite:stderr_reporter",
+        "@com_google_absl//absl/strings",
+    ],
+)
+
+cc_test(
+    name = "model_loader_test",
+    srcs = ["model_loader_test.cc"],
+    deps = [
+        ":embedded_mobilenet_model",
+        ":mini_benchmark_test_helper",
+        ":model_loader",
+        ":status_codes",
+        "@com_google_absl//absl/strings:str_format",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 #
 # Test targets for separate process.
 # Unit tests using cc_test and turned into Android tests with tflite_portable_test_suite().
@@ -770,11 +805,13 @@ cc_binary(
     linkshared = True,
     deps = [
         ":fb_storage",
+        ":model_loader",
         ":runner",
         ":status_codes",
         ":set_big_core_affinity",
         ":validator",
         ":validator_runner",
+        "@com_google_absl//absl/strings",
         "@flatbuffers",
         "//tensorflow/lite/experimental/acceleration/configuration:nnapi_plugin",
         "//tensorflow/lite/experimental/acceleration/configuration:configuration_fbs",
@@ -844,6 +881,7 @@ cc_test(
         ":embedded_mobilenet_validation_model",
         ":embedded_mobilenet_model",
         ":mini_benchmark_test_helper",
+        ":model_loader",
         ":status_codes",
         ":validator",
         "@com_google_googletest//:gtest_main",

tensorflow/lite/experimental/acceleration/mini_benchmark/model_loader.cc

@@ -0,0 +1,77 @@
+/* Copyright 2021 The TensorFlow Authors. 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/lite/experimental/acceleration/mini_benchmark/model_loader.h"
+
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/strings/numbers.h"
+#include "absl/strings/str_split.h"
+#include "tensorflow/lite/allocation.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/stderr_reporter.h"
+
+namespace tflite {
+namespace acceleration {
+
+std::unique_ptr<ModelLoader> ModelLoader::CreateFromFdOrPath(
+    absl::string_view fd_or_path) {
+  if (!absl::StartsWith(fd_or_path, "fd:")) {
+    return std::make_unique<ModelLoader>(fd_or_path);
+  }
+
+  std::vector<std::string> parts = absl::StrSplit(fd_or_path, ':');
+  int model_fd;
+  size_t model_offset, model_size;
+  if (parts.size() != 4 || !absl::SimpleAtoi(parts[1], &model_fd) ||
+      !absl::SimpleAtoi(parts[2], &model_offset) ||
+      !absl::SimpleAtoi(parts[3], &model_size)) {
+    return nullptr;
+  }
+  return std::make_unique<ModelLoader>(model_fd, model_offset, model_size);
+}
+
+MinibenchmarkStatus ModelLoader::Init() {
+  if (model_) {
+    // Already done.
+    return kMinibenchmarkSuccess;
+  }
+  if (model_path_.empty() && model_fd_ <= 0) {
+    return kMinibenchmarkPreconditionNotMet;
+  }
+  if (!model_path_.empty()) {
+    model_ = FlatBufferModel::VerifyAndBuildFromFile(model_path_.c_str());
+  } else if (MMAPAllocation::IsSupported()) {
+    auto allocation = std::make_unique<MMAPAllocation>(
+        model_fd_, model_offset_, model_size_, tflite::DefaultErrorReporter());
+    if (!allocation->valid()) {
+      return kMinibenchmarkModelReadFailed;
+    }
+    model_ =
+        FlatBufferModel::VerifyAndBuildFromAllocation(std::move(allocation));
+  } else {
+    return kMinibenchmarkUnsupportedPlatform;
+  }
+  if (!model_) {
+    return kMinibenchmarkModelBuildFailed;
+  }
+  return kMinibenchmarkSuccess;
+}
+
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_loader.h

@@ -0,0 +1,78 @@
+/* Copyright 2021 The TensorFlow Authors. 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.
+==============================================================================*/
+#ifndef TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_LOADER_H_
+#define TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_LOADER_H_
+
+#include <stddef.h>
+#include <unistd.h>
+
+#include <memory>
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+#include "tensorflow/lite/model_builder.h"
+
+namespace tflite {
+namespace acceleration {
+
+// Class to load the Model.
+class ModelLoader {
+ public:
+  // Create the model loader from a model_path or a file descriptor. File
+  // descriptor path must be in the format of
+  // "fd:%model_fd%:%model_offset%:%model_size%". Return nullptr if the path
+  // starts with "fd:" but cannot be parsed with the given format.
+  static std::unique_ptr<ModelLoader> CreateFromFdOrPath(
+      absl::string_view fd_or_path);
+
+  // Create the model loader from model_path.
+  explicit ModelLoader(absl::string_view model_path)
+      : model_path_(model_path) {}
+
+#ifndef _WIN32
+  // Create the model loader from file descriptor. The model_fd only has to be
+  // valid for the duration of the constructor (it's dup'ed inside). This
+  // constructor is not available on Windows.
+  ModelLoader(int model_fd, size_t model_offset, size_t model_size)
+      : model_fd_(dup(model_fd)),
+        model_offset_(model_offset),
+        model_size_(model_size) {}
+#endif  // !_WIN32
+
+  ~ModelLoader() {
+    if (model_fd_ >= 0) {
+      close(model_fd_);
+    }
+  }
+
+  // Return whether the model is loaded successfully.
+  MinibenchmarkStatus Init();
+
+  const FlatBufferModel* GetModel() const { return model_.get(); }
+
+ private:
+  const std::string model_path_;
+  const int model_fd_ = -1;
+  const size_t model_offset_ = 0;
+  const size_t model_size_ = 0;
+  std::unique_ptr<FlatBufferModel> model_;
+};
+
+}  // namespace acceleration
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_EXPERIMENTAL_ACCELERATION_MINI_BENCHMARK_MODEL_LOADER_H_

tensorflow/lite/experimental/acceleration/mini_benchmark/model_loader_test.cc

@@ -0,0 +1,91 @@
+/* Copyright 2021 The TensorFlow Authors. 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/lite/experimental/acceleration/mini_benchmark/model_loader.h"
+
+#include <fcntl.h>
+#include <sys/stat.h>
+
+#include <memory>
+#include <string>
+
+#include <gmock/gmock.h>
+#include <gtest/gtest.h>
+#include "absl/strings/str_format.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/embedded_mobilenet_model.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/mini_benchmark_test_helper.h"
+#include "tensorflow/lite/experimental/acceleration/mini_benchmark/status_codes.h"
+
+namespace tflite {
+namespace acceleration {
+namespace {
+
+class ModelLoaderTest : public ::testing::Test {
+ protected:
+  void SetUp() override {
+    model_path_ = MiniBenchmarkTestHelper::DumpToTempFile(
+        "mobilenet_quant.tflite",
+        g_tflite_acceleration_embedded_mobilenet_model,
+        g_tflite_acceleration_embedded_mobilenet_model_len);
+  }
+  std::string model_path_;
+};
+
+TEST_F(ModelLoaderTest, CreateFromModelPath) {
+  std::unique_ptr<ModelLoader> model_loader =
+      ModelLoader::CreateFromFdOrPath(model_path_);
+  ASSERT_NE(model_loader, nullptr);
+  EXPECT_THAT(model_loader->Init(), kMinibenchmarkSuccess);
+}
+
+TEST_F(ModelLoaderTest, CreateFromFdPath) {
+  int fd = open(model_path_.c_str(), O_RDONLY);
+  ASSERT_GE(fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(fd, &stat_buf), 0);
+  auto model_loader = std::make_unique<ModelLoader>(fd, 0, stat_buf.st_size);
+  close(fd);
+
+  ASSERT_NE(model_loader, nullptr);
+  EXPECT_THAT(model_loader->Init(), kMinibenchmarkSuccess);
+}
+
+TEST_F(ModelLoaderTest, CreateFromFdOrModelPath) {
+  int fd = open(model_path_.c_str(), O_RDONLY);
+  ASSERT_GE(fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(fd, &stat_buf), 0);
+  std::string path = absl::StrFormat("fd:%d:%zu:%zu", fd, 0, stat_buf.st_size);
+  auto model_loader = ModelLoader::CreateFromFdOrPath(path);
+  close(fd);
+
+  ASSERT_NE(model_loader, nullptr);
+  EXPECT_THAT(model_loader->Init(), kMinibenchmarkSuccess);
+}
+
+TEST_F(ModelLoaderTest, InvalidFdPath) {
+  int fd = open(model_path_.c_str(), O_RDONLY);
+  ASSERT_GE(fd, 0);
+  struct stat stat_buf = {0};
+  ASSERT_EQ(fstat(fd, &stat_buf), 0);
+  std::string path = absl::StrFormat("fd:%d:%zu", fd, 0);
+  auto model_loader = ModelLoader::CreateFromFdOrPath(path);
+  close(fd);
+
+  EXPECT_EQ(model_loader, nullptr);
+}
+
+}  // namespace
+}  // namespace acceleration
+}  // namespace tflite

tensorflow/lite/experimental/acceleration/mini_benchmark/model_validation_test.cc

@@ -105,8 +105,9 @@ class LocalizerValidationRegressionTest : public ::testing::Test {
     ASSERT_GE(fd, 0);
     struct stat stat_buf = {0};
     ASSERT_EQ(fstat(fd, &stat_buf), 0);
-    auto validator =
-        std::make_unique<Validator>(fd, 0, stat_buf.st_size, settings);
+    auto validator = std::make_unique<Validator>(
+        std::make_unique<ModelLoader>(fd, /*offset=*/0, stat_buf.st_size),
+        settings);
     close(fd);
 
     Validator::Results results;