From bbf70db4993618bcabc01cf2f02cd9eb089d5dca Mon Sep 17 00:00:00 2001 From: Ryan Kuester Date: Thu, 10 Jul 2025 15:32:57 -0500 Subject: [PATCH] feat(compression): add SpecBuilder for programmatic compression specs (#3133) Add a fluent builder API for creating compression specifications without writing YAML strings. This is useful in scripts and Jupyter notebooks. Example usage: spec = (compression.SpecBuilder() .add_tensor(subgraph=0, tensor=2) .with_lut(index_bitwidth=4) .build()) BUG=#3125 Co-authored-by: suleshahid <110432064+suleshahid@users.noreply.github.com> --- python/tflite_micro/postinstall_check.py | 4 + tensorflow/lite/micro/compression/BUILD | 20 ++++ tensorflow/lite/micro/compression/__init__.py | 3 +- .../lite/micro/compression/spec_builder.py | 101 ++++++++++++++++ .../micro/compression/spec_builder_test.py | 112 ++++++++++++++++++ 5 files changed, 239 insertions(+), 1 deletion(-) create mode 100644 tensorflow/lite/micro/compression/spec_builder.py create mode 100644 tensorflow/lite/micro/compression/spec_builder_test.py diff --git a/python/tflite_micro/postinstall_check.py b/python/tflite_micro/postinstall_check.py index 93ff9e12..d739d151 100644 --- a/python/tflite_micro/postinstall_check.py +++ b/python/tflite_micro/postinstall_check.py @@ -61,6 +61,10 @@ def compression_test(): # with compressible tensors, but we verify the function is importable assert callable(compression.compress) + # Test availability of the SpecBuilder + _ = (compression.SpecBuilder().add_tensor( + subgraph=0, tensor=0).with_lut(index_bitwidth=4).build()) + return True diff --git a/tensorflow/lite/micro/compression/BUILD b/tensorflow/lite/micro/compression/BUILD index 9f6e8afd..d506c308 100644 --- a/tensorflow/lite/micro/compression/BUILD +++ b/tensorflow/lite/micro/compression/BUILD @@ -23,6 +23,7 @@ py_library( deps = [ ":compress_lib", ":spec", + ":spec_builder", ], ) @@ -189,6 +190,25 @@ py_test( ], ) +py_library( + name = "spec_builder", + srcs = ["spec_builder.py"], + deps = [ + ":spec", + ], +) + +py_test( + name = "spec_builder_test", + size = "small", + srcs = ["spec_builder_test.py"], + deps = [ + ":spec", + ":spec_builder", + requirement("tensorflow"), + ], +) + py_library( name = "test_models", srcs = ["test_models.py"], diff --git a/tensorflow/lite/micro/compression/__init__.py b/tensorflow/lite/micro/compression/__init__.py index 11c635a5..9bdfe6d0 100644 --- a/tensorflow/lite/micro/compression/__init__.py +++ b/tensorflow/lite/micro/compression/__init__.py @@ -22,5 +22,6 @@ from .compress import compress from .spec import parse_yaml +from .spec_builder import SpecBuilder -__all__ = ["compress", "parse_yaml"] +__all__ = ["compress", "parse_yaml", "SpecBuilder"] diff --git a/tensorflow/lite/micro/compression/spec_builder.py b/tensorflow/lite/micro/compression/spec_builder.py new file mode 100644 index 00000000..f62deac4 --- /dev/null +++ b/tensorflow/lite/micro/compression/spec_builder.py @@ -0,0 +1,101 @@ +# Copyright 2025 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. +# +"""Builder pattern for creating compression specifications programmatically. + +This module provides a fluent API for building compression specs without +needing to write YAML strings. + +Example usage: + from tflite_micro.compression import SpecBuilder + + spec = (SpecBuilder() + .add_tensor(subgraph=0, tensor=2) + .with_lut(index_bitwidth=4) + .add_tensor(subgraph=0, tensor=4) + .with_lut(index_bitwidth=2) + .build()) +""" + +from typing import List, Optional +from . import spec + + +class TensorBuilder: + """Builder for individual tensor compression specifications.""" + + def __init__(self, subgraph: int, tensor: int, + parent_builder: 'SpecBuilder'): + self.subgraph = subgraph + self.tensor = tensor + self.compression_methods: List[spec.CompressionMethod] = [] + self._parent = parent_builder + + def with_lut(self, index_bitwidth: int) -> 'SpecBuilder': + """Add LUT compression to this tensor. + + Args: + index_bitwidth: Number of bits for the LUT index (e.g., 4 for 16 values) + + Returns: + The parent SpecBuilder for method chaining + """ + self.compression_methods.append( + spec.LookUpTableCompression(index_bitwidth=index_bitwidth)) + return self._parent + + def _build(self) -> spec.Tensor: + """Build the Tensor specification object.""" + return spec.Tensor(subgraph=self.subgraph, + tensor=self.tensor, + compression=self.compression_methods) + + +class SpecBuilder: + """Fluent builder for compression specifications.""" + + def __init__(self): + self._tensor_builders: List[TensorBuilder] = [] + self._current_tensor: Optional[TensorBuilder] = None + + def add_tensor(self, subgraph: int, tensor: int) -> TensorBuilder: + """Add a tensor to be compressed. + + Args: + subgraph: The subgraph index containing the tensor + tensor: The tensor index within the subgraph + + Returns: + A TensorBuilder for configuring compression methods + """ + # Finalize any current tensor + if self._current_tensor is not None: + self._tensor_builders.append(self._current_tensor) + + # Create new tensor builder + self._current_tensor = TensorBuilder(subgraph, tensor, self) + return self._current_tensor + + def build(self) -> List[spec.Tensor]: + """Build the final compression specification. + + Returns: + A list of Tensor specifications ready for use with compress() + """ + # Make sure to include the last tensor if there is one + if self._current_tensor is not None: + self._tensor_builders.append(self._current_tensor) + self._current_tensor = None + + return [tb._build() for tb in self._tensor_builders] diff --git a/tensorflow/lite/micro/compression/spec_builder_test.py b/tensorflow/lite/micro/compression/spec_builder_test.py new file mode 100644 index 00000000..7c434203 --- /dev/null +++ b/tensorflow/lite/micro/compression/spec_builder_test.py @@ -0,0 +1,112 @@ +# Copyright 2025 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. +# +"""Tests for the compression spec builder.""" + +import tensorflow as tf + +from tflite_micro.tensorflow.lite.micro.compression import spec +from tflite_micro.tensorflow.lite.micro.compression import spec_builder + + +class SpecBuilderTest(tf.test.TestCase): + + def test_basic_builder_pattern(self): + """Test basic fluent builder usage.""" + result = (spec_builder.SpecBuilder().add_tensor( + subgraph=0, tensor=2).with_lut(index_bitwidth=4).add_tensor( + subgraph=0, tensor=4).with_lut(index_bitwidth=2).build()) + + self.assertEqual(len(result), 2) + + # Check first tensor + self.assertEqual(result[0].subgraph, 0) + self.assertEqual(result[0].tensor, 2) + self.assertEqual(len(result[0].compression), 1) + self.assertIsInstance(result[0].compression[0], + spec.LookUpTableCompression) + self.assertEqual(result[0].compression[0].index_bitwidth, 4) + + # Check second tensor + self.assertEqual(result[1].subgraph, 0) + self.assertEqual(result[1].tensor, 4) + self.assertEqual(len(result[1].compression), 1) + self.assertIsInstance(result[1].compression[0], + spec.LookUpTableCompression) + self.assertEqual(result[1].compression[0].index_bitwidth, 2) + + def test_non_chained_usage(self): + """Test using builder without method chaining.""" + builder = spec_builder.SpecBuilder() + builder.add_tensor(0, 2).with_lut(4) + builder.add_tensor(0, 4).with_lut(2) + result = builder.build() + + self.assertEqual(len(result), 2) + self.assertEqual(result[0].tensor, 2) + self.assertEqual(result[0].compression[0].index_bitwidth, 4) + self.assertEqual(result[1].tensor, 4) + self.assertEqual(result[1].compression[0].index_bitwidth, 2) + + def test_empty_spec(self): + """Test building an empty spec.""" + result = spec_builder.SpecBuilder().build() + self.assertEqual(len(result), 0) + + def test_single_tensor(self): + """Test building a spec with just one tensor.""" + result = (spec_builder.SpecBuilder().add_tensor( + subgraph=2, tensor=42).with_lut(index_bitwidth=16).build()) + + self.assertEqual(len(result), 1) + self.assertEqual(result[0].subgraph, 2) + self.assertEqual(result[0].tensor, 42) + self.assertEqual(result[0].compression[0].index_bitwidth, 16) + + def test_tensor_without_compression(self): + """Test that tensors can be added without compression methods.""" + builder = spec_builder.SpecBuilder() + # Add tensor but don't call with_lut + builder.add_tensor(0, 1) + builder.add_tensor(0, 2).with_lut(4) + result = builder.build() + + self.assertEqual(len(result), 2) + self.assertEqual(result[0].tensor, 1) + self.assertEqual(len(result[0].compression), 0) + self.assertEqual(result[1].tensor, 2) + self.assertEqual(len(result[1].compression), 1) + + def test_builder_produces_same_type_as_parse_yaml(self): + """Test that builder produces same data structure as parse_yaml.""" + # Build using the builder + built_spec = (spec_builder.SpecBuilder().add_tensor( + subgraph=0, tensor=42).with_lut(index_bitwidth=4).add_tensor( + subgraph=0, tensor=55).with_lut(index_bitwidth=2).build()) + + # Parse the example YAML from spec.py + parsed_spec = spec.parse_yaml(spec.EXAMPLE_YAML_SPEC) + + # They should be equivalent + self.assertEqual(len(built_spec), len(parsed_spec)) + for built, parsed in zip(built_spec, parsed_spec): + self.assertEqual(built.subgraph, parsed.subgraph) + self.assertEqual(built.tensor, parsed.tensor) + self.assertEqual(len(built.compression), len(parsed.compression)) + self.assertEqual(built.compression[0].index_bitwidth, + parsed.compression[0].index_bitwidth) + + +if __name__ == "__main__": + tf.test.main()