Draft PR: testing

tensorflow/lite/micro/compression.h

@@ -0,0 +1,70 @@
+/* Copyright 2024 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_MICRO_MICRO_COMPRESSION_H_
+#define TENSORFLOW_LITE_MICRO_MICRO_COMPRESSION_H_
+
+#ifdef USE_TFLM_COMPRESSION
+
+#include "tensorflow/lite/c/common.h"
+
+namespace tflite {
+
+//
+// Compressed tensors
+//
+
+static constexpr const char* kCompressionMetadataString = "TFLM_COMPRESSION";
+
+enum class CompressionScheme : uint8_t {
+  kBinQuant,
+};
+
+// TODO(ddavis-2015): pack struct
+struct BinQuantData {
+  static constexpr size_t kMaxBitWidth = 7;
+  static constexpr size_t kMaxValueTableChannelStride = 128;
+
+  const void* value_table;             // Pointer into FlatBuffer Values.
+  uint8_t value_table_channel_stride;  // elements per channel
+  uint8_t compressed_bit_width : 3;    // 1 to 7 bits
+  bool is_per_channel_quantized : 1;   // tensor is per-channel quantized
+  bool use_alternate_axis : 1;         // shape default channel:
+                                       // 0 = first, 1 = last
+  uint8_t reserved : 3;
+};
+
+union CompressionData {
+  BinQuantData bin_quant;
+};
+
+// TODO(ddavis-2015): pack struct
+struct CompressionTensorData {
+  CompressionScheme scheme;
+  CompressionData data;
+};
+
+// TODO(ddavis-2015): pack struct
+struct CompressedTensorList {
+  // Sparsely populated array with the same number of elements as there are
+  // tensors in the Subgraph. An alternative would include a tensor index in
+  // the struct for each and walk the list on look up. This could be slow.
+  CompressionTensorData** tensors;
+};
+
+}  // namespace tflite
+
+#endif  // USE_TFLM_COMPRESSION
+#endif  // TENSORFLOW_LITE_MICRO_MICRO_COMPRESSION_H_

tensorflow/lite/micro/compression/BUILD

@@ -0,0 +1,94 @@
+load("@flatbuffers//:build_defs.bzl", "flatbuffer_cc_library", "flatbuffer_py_library")
+load("@rules_python//python:defs.bzl", "py_test")
+load("@tflm_pip_deps//:requirements.bzl", "requirement")
+
+package(
+    default_visibility = [
+        "//visibility:public",
+    ],
+)
+
+flatbuffer_cc_library(
+    name = "metadata_flatbuffer_cc",
+    srcs = ["metadata.fbs"],
+)
+
+flatbuffer_py_library(
+    name = "original_flatbuffer_py",
+    srcs = ["original.fbs"],
+)
+
+flatbuffer_py_library(
+    name = "metadata_flatbuffer_py",
+    srcs = ["metadata.fbs"],
+)
+
+cc_test(
+    name = "metadata_test_cc",
+    srcs = ["metadata_test.cc"],
+    deps = [
+        "metadata_flatbuffer_cc",
+        "//tensorflow/lite/micro:hexdump",
+        "@flatbuffers//:runtime_cc",
+    ],
+    size = "small",
+)
+
+py_binary(
+    name = "compress",
+    srcs = ["compress.py"],
+    deps = [
+        "@absl_py//absl:app",
+        "@absl_py//absl/flags",
+        "@absl_py//absl/logging",
+        "@flatbuffers//:runtime_py",
+        "metadata_flatbuffer_py",
+        "//tensorflow/lite/python:schema_py",
+        requirement("bitarray"),
+        requirement("numpy"),
+        requirement("scikit-learn"),
+    ],
+)
+
+py_binary(
+    name = "view",
+    srcs = [
+        "view.py",
+    ],
+    deps = [
+        "metadata_flatbuffer_py",
+        "//tensorflow/lite/python:schema_py",
+    ],
+)
+
+py_test(
+    name = "metadata_test_py",
+    main = "metadata_test.py",
+    srcs = ["metadata_test.py"],
+    deps = [
+        "metadata_flatbuffer_py",
+        "@flatbuffers//:runtime_py",
+        requirement("hexdump"),
+    ],
+    size = "small",
+)
+
+py_test(
+    name = "original_test_py",
+    main = "original_test.py",
+    srcs = ["original_test.py"],
+    deps = [
+        "original_flatbuffer_py",
+        "@flatbuffers//:runtime_py",
+        requirement("hexdump"),
+    ],
+    size = "small",
+)
+
+genrule(
+    name = "hello_world_int8.compressed",
+    srcs = ["//tensorflow/lite/micro/examples/hello_world/models:hello_world_int8.tflite"],
+    outs = ["hello_world_int8.compressed.tflite"],
+    cmd = "$(location :compress) --input_model_path $< --output_model_path $@",
+    tools = [":compress"],
+)