chore: import upstream snapshot with attribution

tensorflow/lite/tools/strip_buffers/BUILD

@@ -0,0 +1,67 @@
+# Tools to strip out eligible buffers (tensor data) from a TFLite flatbuffer, and load them with
+# random values at runtime.
+
+load("@rules_cc//cc:cc_binary.bzl", "cc_binary")
+load("@rules_cc//cc:cc_library.bzl", "cc_library")
+load("//tensorflow/lite:build_def.bzl", "tflite_copts", "tflite_linkopts")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:LICENSE"],
+    default_visibility = [
+        "//visibility:public",
+    ],
+    licenses = ["notice"],
+)
+
+cc_library(
+    name = "stripping_lib",
+    srcs = ["stripping_lib.cc"],
+    hdrs = ["stripping_lib.h"],
+    copts = tflite_copts(),
+    deps = [
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite:string",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/schema:schema_fbs",
+        "@com_google_absl//absl/container:flat_hash_set",
+        "@com_google_absl//absl/log",
+        "@com_google_absl//absl/strings",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
+cc_binary(
+    name = "strip_buffers_from_fb",
+    srcs = ["strip_buffers_from_fb.cc"],
+    copts = tflite_copts(),
+    linkopts = tflite_linkopts(),
+    deps = [
+        ":stripping_lib",
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/tools:command_line_flags",
+        "@com_google_absl//absl/log",
+        "@flatbuffers//:runtime_cc",
+    ],
+)
+
+cc_binary(
+    name = "reconstitute_buffers_into_fb",
+    srcs = ["reconstitute_buffers_into_fb.cc"],
+    copts = tflite_copts(),
+    linkopts = tflite_linkopts(),
+    deps = [
+        ":stripping_lib",
+        "//tensorflow/core:tflite_portable_logging",
+        "//tensorflow/lite:framework",
+        "//tensorflow/lite/core:framework",
+        "//tensorflow/lite/core/c:common",
+        "//tensorflow/lite/tools:command_line_flags",
+        "@com_google_absl//absl/log",
+        "@flatbuffers//:runtime_cc",
+    ],
+)

tensorflow/lite/tools/strip_buffers/README.md

@@ -0,0 +1,57 @@
+# TFLite Buffer-Stripping Tool/Library
+
+**NOTE: This is an advanced tool used to reduce bandwidth usage in Neural
+Architecture Search applications. Use with caution.**
+
+The tools in this directory make it easier to distribute TFLite models to
+multiple devices over networks with the sole aim of benchmarking *latency*
+performance. The intended workflow is as follows:
+
+*   The stripping tool empties eligible constants from a TFLite flatbuffer to
+    reduce its size.
+*   This lean model can be easily transported to devices over a network.
+*   The reconstitution tool on the device takes in a flatbuffer in memory, and
+    fills in the appropriate buffers with random data.
+
+As an example, see the before/after sizes for MobileNetV1:
+
+*   Float: 16.9MB -> 12KB
+*   Quantized: 4.3MB -> 17.6 KB
+
+**NOTE: This tool only supports single subgraphs for now.**
+
+There are two tools in this directory:
+
+## 1. Stripping buffers out of TFLite flatbuffers
+
+This tool takes in an input `flatbuffer`, and strips out (or 'empties') the
+buffers (constant data) for tensors that follow the following guidelines:
+
+*   Are either of: Float32, Int32, UInt8, Int8
+*   If Int32, the tensor should have a min of 10 elements
+
+The second rule above protects us from invalidating constant data that cannot be
+randomised (for example, Reshape 'shape' input).
+
+To run the associated script:
+
+```
+bazel run -c opt tensorflow/lite/tools/strip_buffers:strip_buffers_from_fb -- --input_flatbuffer=/input/path.tflite --output_flatbuffer=/output/path.tflite
+```
+
+## 2. Stripping buffers out of TFLite flatbuffers
+
+The idea here is to reconstitute the lean flatbuffer `Model` generared in the
+above step, by filling in random data whereever necessary.
+
+The prototype script can be called as:
+
+```
+bazel run -c opt tensorflow/lite/tools/strip_buffers:reconstitute_buffers_into_fb -- --input_flatbuffer=/input/path.tflite --output_flatbuffer=/output/path.tflite
+```
+
+## C++ Library
+
+Both the above tools are present as `stripping_lib` in this directory, which
+mutate the flatbuffer(s) in-memory. This ensures we can do the above two steps
+without touching the filesystem again.

tensorflow/lite/tools/strip_buffers/reconstitute_buffers_into_fb.cc

@@ -0,0 +1,82 @@
+/* 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.
+==============================================================================*/
+
+// Binary to test strip_buffers/reconstitution.h.
+#include <fstream>  // NOLINT
+#include <string>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/tools/command_line_flags.h"
+#include "tensorflow/lite/tools/strip_buffers/stripping_lib.h"
+
+#define TFLITE_SCHEMA_VERSION 3
+
+namespace tflite {
+
+using ::flatbuffers::FlatBufferBuilder;
+
+constexpr char kInputFlatbufferFlag[] = "input_flatbuffer";
+constexpr char kOutputFlatbufferFlag[] = "output_flatbuffer";
+
+int Main(int argc, char* argv[]) {
+  // Command Line Flags.
+  std::string input_flatbuffer_path;
+  std::string output_flatbuffer_path;
+
+  std::vector<Flag> flag_list = {
+      tflite::Flag::CreateFlag(kInputFlatbufferFlag, &input_flatbuffer_path,
+                               "Path to input TFLite flatbuffer."),
+      tflite::Flag::CreateFlag(kOutputFlatbufferFlag, &output_flatbuffer_path,
+                               "Path to output TFLite flatbuffer."),
+  };
+  Flags::Parse(&argc, const_cast<const char**>(argv), flag_list);
+
+  // Read in input flatbuffer.
+  auto input_model =
+      FlatBufferModel::BuildFromFile(input_flatbuffer_path.c_str());
+
+  if (!FlatbufferHasStrippedWeights(input_model->GetModel())) {
+    LOG(ERROR) << "The weights are already available in the input model!";
+    return 0;
+  }
+  // Reconstitute flatbuffer with appropriate random constant tensors
+  FlatBufferBuilder builder(/*initial_size=*/10240);
+  if (ReconstituteConstantTensorsIntoFlatbuffer(input_model->GetModel(),
+                                                &builder) != kTfLiteOk) {
+    return 0;
+  }
+
+  LOG(INFO) << "Flatbuffer size (KB) BEFORE: "
+            << input_model->allocation()->bytes() / 1000.0;
+  LOG(INFO) << "Flatbuffer size (KB) AFTER: " << builder.GetSize() / 1000.0;
+
+  // Write the output model to file.
+  std::string output_model_content(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+  std::ofstream output_file_stream(output_flatbuffer_path);
+  output_file_stream << output_model_content;
+  output_file_stream.close();
+
+  return 0;
+}
+}  // namespace tflite
+
+int main(int argc, char* argv[]) { return tflite::Main(argc, argv); }

tensorflow/lite/tools/strip_buffers/strip_buffers_from_fb.cc

@@ -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.
+==============================================================================*/
+
+// Binary to test strip_buffers/reconstitution.h.
+#include <fstream>  // NOLINT
+#include <string>
+#include <vector>
+
+#include "absl/log/log.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/tools/command_line_flags.h"
+#include "tensorflow/lite/tools/strip_buffers/stripping_lib.h"
+
+#define TFLITE_SCHEMA_VERSION 3
+
+namespace tflite {
+
+using ::flatbuffers::FlatBufferBuilder;
+
+constexpr char kInputFlatbufferFlag[] = "input_flatbuffer";
+constexpr char kOutputFlatbufferFlag[] = "output_flatbuffer";
+
+int Main(int argc, char* argv[]) {
+  // Command Line Flags.
+  std::string input_flatbuffer_path;
+  std::string output_flatbuffer_path;
+
+  std::vector<Flag> flag_list = {
+      tflite::Flag::CreateFlag(kInputFlatbufferFlag, &input_flatbuffer_path,
+                               "Path to input TFLite flatbuffer."),
+      tflite::Flag::CreateFlag(kOutputFlatbufferFlag, &output_flatbuffer_path,
+                               "Path to output TFLite flatbuffer."),
+  };
+  Flags::Parse(&argc, const_cast<const char**>(argv), flag_list);
+
+  // Read in input flatbuffer.
+  auto input_model =
+      FlatBufferModel::BuildFromFile(input_flatbuffer_path.c_str());
+
+  // Strip applicable constants from the flatbuffer.
+  FlatBufferBuilder builder(/*initial_size=*/10240);
+  if (StripWeightsFromFlatbuffer(input_model->GetModel(), &builder) !=
+      kTfLiteOk) {
+    return 0;
+  }
+
+  LOG(INFO) << "Flatbuffer size (KB) BEFORE: "
+            << input_model->allocation()->bytes() / 1000.0;
+  LOG(INFO) << "Flatbuffer size (KB) AFTER: " << builder.GetSize() / 1000.0;
+
+  // Write the output model to file.
+  std::string output_model_content(
+      reinterpret_cast<const char*>(builder.GetBufferPointer()),
+      builder.GetSize());
+  std::ofstream output_file_stream(output_flatbuffer_path);
+  output_file_stream << output_model_content;
+  output_file_stream.close();
+
+  return 0;
+}
+}  // namespace tflite
+
+int main(int argc, char* argv[]) { return tflite::Main(argc, argv); }

tensorflow/lite/tools/strip_buffers/stripping_lib.cc

@@ -0,0 +1,361 @@
+/* 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/tools/strip_buffers/stripping_lib.h"
+
+#include <stdint.h>
+
+#include <limits>
+#include <memory>
+#include <random>
+#include <vector>
+
+#include "absl/container/flat_hash_set.h"
+#include "absl/log/log.h"
+#include "absl/strings/string_view.h"
+#include "flatbuffers/buffer.h"  // from @flatbuffers
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/model_builder.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+#include "tensorflow/lite/string_type.h"
+
+#define TFLITE_SCHEMA_VERSION 3
+
+namespace tflite {
+
+using ::flatbuffers::FlatBufferBuilder;
+using ::flatbuffers::Offset;
+
+// Parameters for a simple Gaussian distribution to generate values roughly in
+// [0, 1).
+constexpr float kGaussianFloatMean = 0.5;
+constexpr float kGaussianStdDev = 1.0 / 3;
+
+template <typename Type, typename TypeT>
+void CopyToOffsetVector(FlatBufferBuilder* builder, const Type* data,
+                        std::vector<Offset<Type>>& vec) {
+  std::unique_ptr<TypeT> unpacked(data->UnPack());
+  flatbuffers::Offset<Type> offset = Type::Pack(*builder, unpacked.get());
+  vec.push_back(offset);
+}
+
+int GetNumElements(const std::vector<int>& dims) {
+  int num_elements = 1;
+  for (int i = 0; i < dims.size(); i++) {
+    num_elements *= dims[i];
+  }
+  return num_elements;
+}
+
+// TODO(b/141023954): Reconcile this with the function in
+// inference_profiler_stage.
+template <typename T>
+void GenerateRandomGaussianData(int64_t num_elements, float min, float max,
+                                std::vector<T>* data) {
+  data->clear();
+  data->reserve(num_elements);
+
+  static std::normal_distribution<double> distribution(kGaussianFloatMean,
+                                                       kGaussianStdDev);
+  static std::default_random_engine generator;
+  for (int i = 0; i < num_elements; ++i) {
+    auto rand_n = distribution(generator);
+    while (rand_n < 0 || rand_n >= 1) {
+      rand_n = distribution(generator);
+    }
+    auto rand_float = min + (max - min) * static_cast<float>(rand_n);
+    data->push_back(static_cast<T>(rand_float));
+  }
+}
+
+TfLiteStatus ConvertTensorType(TensorType tensor_type, TfLiteType* type) {
+  *type = kTfLiteNoType;
+  switch (tensor_type) {
+    case TensorType_FLOAT32:
+      *type = kTfLiteFloat32;
+      break;
+    case TensorType_INT32:
+      *type = kTfLiteInt32;
+      break;
+    case TensorType_UINT32:
+      *type = kTfLiteUInt32;
+      break;
+    case TensorType_UINT8:
+      *type = kTfLiteUInt8;
+      break;
+    case TensorType_INT8:
+      *type = kTfLiteInt8;
+      break;
+    default:
+      break;
+  }
+  if (*type == kTfLiteNoType) {
+    VLOG(0) << "Unsupported data type %d in tensor: " << tensor_type;
+    return kTfLiteError;
+  }
+  return kTfLiteOk;
+}
+
+TfLiteStatus StripWeightsFromFlatbuffer(
+    const Model* input_model,
+    flatbuffers::FlatBufferBuilder* new_model_builder) {
+  // TODO(b/141023954): Generalize to N subgraphs.
+  if (input_model->subgraphs()->size() != 1) {
+    VLOG(0) << "Only 1 subgraph supported for now: "
+            << input_model->subgraphs()->size();
+    return kTfLiteError;
+  }
+
+  // Data structures for output flatbuffer.
+  std::vector<Offset<SubGraph>> output_subgraphs;
+  std::vector<Offset<OperatorCode>> output_opcodes;
+  std::vector<Offset<Buffer>> output_buffers;
+
+  const SubGraph* input_subgraph = (*input_model->subgraphs())[0];
+  std::unique_ptr<SubGraphT> mutable_subgraph(input_subgraph->UnPack());
+
+  // For constant tensors that meet requirements:
+  // 1. Set the buffer-id to something larger than total number of buffers.
+  // This indicates to reconstitute_weights_into_fb that random data should be
+  // generated for them.
+  // 2. Mark that buffer for not getting carried into the output flatbuffer.
+  absl::flat_hash_set<uint32_t> erased_buffers;
+  const int num_buffers = input_model->buffers()->size();
+  int i = 0;
+  for (auto& tensor : mutable_subgraph->tensors) {
+    // We don't support Int32 tensors because they could contain
+    // non-randomisable information like Reshape dims.
+    if (tensor->type == TensorType_INT32 &&
+        GetNumElements(tensor->shape) < 10) {
+      // Int32 tensors of elements < 10 could be non-randomisable: for example,
+      // 'shape' input to a Reshape op.
+      continue;
+    } else if (tensor->type != TensorType_INT32 &&
+               tensor->type != TensorType_FLOAT32 &&
+               tensor->type != TensorType_UINT8 &&
+               tensor->type != TensorType_INT8) {
+      continue;
+    }
+
+    if (auto* buffer = (*input_model->buffers())[tensor->buffer]) {
+      if (auto* array = buffer->data()) {
+        VLOG(1) << "Tensor " << i
+                << " is constant, with buffer = " << tensor->buffer;
+        // Set tensor buffer to a high value (num_buffers * 2) & put an empty
+        // buffer in place of the original one.
+        erased_buffers.insert(tensor->buffer);
+        tensor->buffer = num_buffers * 2;
+      }
+    }
+    ++i;
+  }
+
+  i = 0;
+  for (const Buffer* buffer : *(input_model->buffers())) {
+    if (erased_buffers.find(i) == erased_buffers.end()) {
+      // If buffer is not to be erased, insert it into the output flatbuffer to
+      // retain data.
+      CopyToOffsetVector<Buffer, BufferT>(new_model_builder, buffer,
+                                          output_buffers);
+    } else {
+      output_buffers.push_back(CreateBuffer(*new_model_builder));
+    }
+    ++i;
+  }
+
+  flatbuffers::Offset<SubGraph> output_subgraph =
+      SubGraph::Pack(*new_model_builder, mutable_subgraph.get());
+  output_subgraphs.push_back(output_subgraph);
+
+  // Write all ops as they are.
+  for (const OperatorCode* opcode : *(input_model->operator_codes())) {
+    CopyToOffsetVector<OperatorCode, OperatorCodeT>(new_model_builder, opcode,
+                                                    output_opcodes);
+  }
+
+  // Generate output model.
+  auto description =
+      new_model_builder->CreateString("Generated by strip_buffers_from_fb");
+  auto new_model_offset =
+      CreateModel(*new_model_builder, TFLITE_SCHEMA_VERSION,
+                  new_model_builder->CreateVector(output_opcodes),
+                  new_model_builder->CreateVector(output_subgraphs),
+                  description, new_model_builder->CreateVector(output_buffers),
+                  /* metadata_buffer */ 0, /* metadatas */ 0);
+  FinishModelBuffer(*new_model_builder, new_model_offset);
+
+  return kTfLiteOk;
+}
+
+string StripWeightsFromFlatbuffer(const absl::string_view input_flatbuffer) {
+  auto input_model = FlatBufferModel::VerifyAndBuildFromBuffer(
+      input_flatbuffer.data(), input_flatbuffer.size());
+  if (!input_model) {
+    return string();
+  }
+
+  FlatBufferBuilder builder(/*initial_size=*/10240);
+  if (StripWeightsFromFlatbuffer(input_model->GetModel(), &builder) !=
+      kTfLiteOk) {
+    return string();
+  }
+
+  return string(reinterpret_cast<const char*>(builder.GetBufferPointer()),
+                builder.GetSize());
+}
+
+bool FlatbufferHasStrippedWeights(const Model* input_model) {
+  if (input_model->subgraphs()->size() != 1) {
+    VLOG(0) << "Only 1 subgraph supported for now";
+    return false;
+  }
+  const SubGraph* input_subgraph = (*input_model->subgraphs())[0];
+  std::unique_ptr<SubGraphT> mutable_subgraph(input_subgraph->UnPack());
+
+  // For all tensors that have buffer > num_buffers + 1 (set to be so in
+  // strip_buffers_from_fb), create a buffer with random data & assign to them.
+  // For others, just copy over the original buffer from source model.
+  const int num_buffers = input_model->buffers()->size();
+  for (auto& tensor : mutable_subgraph->tensors) {
+    if (tensor->buffer > num_buffers + 1) {
+      return true;
+    }
+  }
+  return false;
+}
+
+TfLiteStatus ReconstituteConstantTensorsIntoFlatbuffer(
+    const Model* input_model,
+    flatbuffers::FlatBufferBuilder* new_model_builder) {
+  // TODO(b/141023954): Generalize to N subgraphs.
+  if (input_model->subgraphs()->size() != 1) {
+    VLOG(0) << "Only 1 subgraph supported for now";
+    return kTfLiteError;
+  }
+  const SubGraph* input_subgraph = (*input_model->subgraphs())[0];
+  std::unique_ptr<SubGraphT> mutable_subgraph(input_subgraph->UnPack());
+
+  // Containers for output flatbuffer.
+  std::vector<Offset<::tflite::SubGraph>> output_subgraphs;
+  std::vector<Offset<::tflite::OperatorCode>> output_opcodes;
+  std::vector<Offset<::tflite::Buffer>> output_buffers;
+
+  // An empty buffer, needed as a TFLite convention.
+  output_buffers.push_back(CreateBuffer(*new_model_builder));
+
+  // For all tensors that have buffer > num_buffers + 1 (set to be so in
+  // strip_buffers_from_fb), create a buffer with random data & assign to them.
+  // For others, just copy over the original buffer from source model.
+  const int num_buffers = input_model->buffers()->size();
+  for (auto& tensor : mutable_subgraph->tensors) {
+    int buffer_id = output_buffers.size();
+    if (tensor->buffer > num_buffers + 1) {
+      // Num tensor elements.
+      int num_elements = GetNumElements(tensor->shape);
+      // Tensor type.
+      TfLiteType type;
+      if (ConvertTensorType(tensor->type, &type) != kTfLiteOk) {
+        return kTfLiteError;
+      }
+      // Generate buffer random data.
+      // Use different min/max bounds per tensor-type to ensure that random data
+      // 'appears' similar to typical values.
+      if (type == kTfLiteUInt8) {
+        std::vector<uint8_t> data;
+        GenerateRandomGaussianData(num_elements,
+                                   std::numeric_limits<uint8_t>::min(),
+                                   std::numeric_limits<uint8_t>::max(), &data);
+        auto data_buffer = new_model_builder->CreateVector(
+            reinterpret_cast<const uint8_t*>(data.data()),
+            sizeof(uint8_t) * data.size());
+        output_buffers.push_back(CreateBuffer(*new_model_builder, data_buffer));
+      } else if (type == kTfLiteInt8) {
+        std::vector<int8_t> data;
+        GenerateRandomGaussianData(num_elements,
+                                   std::numeric_limits<int8_t>::min(),
+                                   std::numeric_limits<int8_t>::max(), &data);
+        auto data_buffer = new_model_builder->CreateVector(
+            reinterpret_cast<const uint8_t*>(data.data()),
+            sizeof(int8_t) * data.size());
+        output_buffers.push_back(CreateBuffer(*new_model_builder, data_buffer));
+      } else if (type == kTfLiteFloat32) {
+        std::vector<float> data;
+        GenerateRandomGaussianData(num_elements, -1, 1, &data);
+        auto data_buffer = new_model_builder->CreateVector(
+            reinterpret_cast<const uint8_t*>(data.data()),
+            sizeof(float) * data.size());
+        output_buffers.push_back(CreateBuffer(*new_model_builder, data_buffer));
+      } else if (type == kTfLiteInt32) {
+        std::vector<int32_t> data;
+        GenerateRandomGaussianData(num_elements, 10, 10, &data);
+        auto data_buffer = new_model_builder->CreateVector(
+            reinterpret_cast<const uint8_t*>(data.data()),
+            sizeof(int32_t) * data.size());
+        output_buffers.push_back(CreateBuffer(*new_model_builder, data_buffer));
+      }
+    } else {
+      // For intermediate tensors, create a new buffer & assign the id to them.
+      // output_buffers.push_back(CreateBuffer(*new_model_builder));
+      CopyToOffsetVector<Buffer, BufferT>(
+          new_model_builder, input_model->buffers()->Get(tensor->buffer),
+          output_buffers);
+    }
+    tensor->buffer = buffer_id;
+  }
+
+  for (const ::tflite::OperatorCode* opcode :
+       *(input_model->operator_codes())) {
+    CopyToOffsetVector<::tflite::OperatorCode, ::tflite::OperatorCodeT>(
+        new_model_builder, opcode, output_opcodes);
+  }
+
+  flatbuffers::Offset<::tflite::SubGraph> output_subgraph =
+      ::tflite::SubGraph::Pack(*new_model_builder, mutable_subgraph.get());
+  output_subgraphs.push_back(output_subgraph);
+
+  auto description = new_model_builder->CreateString(
+      "Generated by TFLite strip_buffers/reconstitution");
+  auto new_model_offset =
+      CreateModel(*new_model_builder, TFLITE_SCHEMA_VERSION,
+                  new_model_builder->CreateVector(output_opcodes),
+                  new_model_builder->CreateVector(output_subgraphs),
+                  description, new_model_builder->CreateVector(output_buffers),
+                  /* metadata_buffer */ 0, /* metadatas */ 0);
+  FinishModelBuffer(*new_model_builder, new_model_offset);
+
+  return kTfLiteOk;
+}
+
+string ReconstituteConstantTensorsIntoFlatbuffer(
+    const absl::string_view input_flatbuffer) {
+  auto input_model = FlatBufferModel::VerifyAndBuildFromBuffer(
+      input_flatbuffer.data(), input_flatbuffer.size());
+  if (!input_model) {
+    return string();
+  }
+
+  FlatBufferBuilder builder(/*initial_size=*/10240);
+  if (ReconstituteConstantTensorsIntoFlatbuffer(input_model->GetModel(),
+                                                &builder) != kTfLiteOk) {
+    return string();
+  }
+
+  return string(reinterpret_cast<const char*>(builder.GetBufferPointer()),
+                builder.GetSize());
+}
+
+}  // namespace tflite

tensorflow/lite/tools/strip_buffers/stripping_lib.h

@@ -0,0 +1,64 @@
+/* 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_TOOLS_STRIP_BUFFERS_STRIPPING_LIB_H_
+#define TENSORFLOW_LITE_TOOLS_STRIP_BUFFERS_STRIPPING_LIB_H_
+
+#include <string>
+
+#include "absl/strings/string_view.h"
+#include "flatbuffers/flatbuffer_builder.h"  // from @flatbuffers
+#include "tensorflow/lite/core/c/common.h"
+#include "tensorflow/lite/schema/schema_generated.h"
+
+namespace tflite {
+
+// Strips eligible buffers from Flatbuffer, to generate a 'leaner' model.
+// Buffers for tensors that satisfy the following constraints are stripped out:
+// 1. Are either of: Float32, Int32, UInt8, Int8
+// 2. If Int32, the tensor should have a min of 10 elements
+// NOTE: This only supports a single Subgraph for now.
+TfLiteStatus StripWeightsFromFlatbuffer(
+    const Model* input_model,
+    flatbuffers::FlatBufferBuilder* new_model_builder);
+
+// The same function as above, but takes in the input model flatbuffer in a
+// string and returns the stripped version in a string.
+// Returns empty string on error.
+// NOTE: This only supports a single Subgraph for now.
+std::string StripWeightsFromFlatbuffer(
+    const absl::string_view input_flatbuffer);
+
+// Generates buffers with random data, for tensors that were mutated using
+// strip_buffers_from_fb.
+// The modified flatbuffer is built into new_model_builder.
+// NOTE: This only supports a single Subgraph for now.
+TfLiteStatus ReconstituteConstantTensorsIntoFlatbuffer(
+    const Model* input_model,
+    flatbuffers::FlatBufferBuilder* new_model_builder);
+
+// The same function as above but takes in the input model flatbuffer in a
+// string and returns the reconstituded version in a string.
+// Returns empty string on error.
+// NOTE: This only supports a single Subgraph for now.
+std::string ReconstituteConstantTensorsIntoFlatbuffer(
+    const absl::string_view input_flatbuffer);
+
+// Return true if the input model has been stripped before.
+bool FlatbufferHasStrippedWeights(const Model* input_model);
+
+}  // namespace tflite
+
+#endif  // TENSORFLOW_LITE_TOOLS_STRIP_BUFFERS_STRIPPING_LIB_H_