diff --git a/python/tflite_micro/BUILD b/python/tflite_micro/BUILD index b358fd12adc..812cf7092fd 100644 --- a/python/tflite_micro/BUILD +++ b/python/tflite_micro/BUILD @@ -125,7 +125,7 @@ py_test( ":runtime", requirement("numpy"), requirement("tensorflow"), - "//tensorflow/lite/micro/compression", + "//tensorflow/lite/micro/compression:model_editor", ], ) diff --git a/python/tflite_micro/test_compression_unsupported.py b/python/tflite_micro/test_compression_unsupported.py index 3692dd0a43a..01c598374ce 100644 --- a/python/tflite_micro/test_compression_unsupported.py +++ b/python/tflite_micro/test_compression_unsupported.py @@ -12,84 +12,84 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -"""Test compression metadata detection when compression is disabled.""" +"""Test legacy compression metadata detection when compression is disabled.""" import os import numpy as np import tensorflow as tf from tflite_micro.python.tflite_micro import runtime -from tflite_micro.tensorflow.lite.micro import compression +from tflite_micro.tensorflow.lite.micro.compression import model_editor -class CompressionDetectionTest(tf.test.TestCase): - """Test compression metadata detection when compression is disabled.""" +def _create_test_model(): + """Create a simple quantized model for testing.""" + model = tf.keras.Sequential([ + tf.keras.layers.Dense(10, input_shape=(5, ), activation='relu'), + tf.keras.layers.Dense(5, activation='softmax') + ]) + model.compile(optimizer='adam', loss='sparse_categorical_crossentropy') - def _create_test_model(self): - """Create a simple quantized model for testing.""" - model = tf.keras.Sequential([ - tf.keras.layers.Dense(10, input_shape=(5, ), activation='relu'), - tf.keras.layers.Dense(5, activation='softmax') - ]) - model.compile(optimizer='adam', loss='sparse_categorical_crossentropy') + converter = tf.lite.TFLiteConverter.from_keras_model(model) + converter.optimizations = [tf.lite.Optimize.DEFAULT] - # Convert to quantized TFLite - converter = tf.lite.TFLiteConverter.from_keras_model(model) - converter.optimizations = [tf.lite.Optimize.DEFAULT] + def representative_dataset(): + for _ in range(10): + yield [np.random.randn(1, 5).astype(np.float32)] - def representative_dataset(): - for _ in range(10): - yield [np.random.randn(1, 5).astype(np.float32)] + converter.representative_dataset = representative_dataset + converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] + converter.inference_input_type = tf.uint8 + converter.inference_output_type = tf.uint8 - converter.representative_dataset = representative_dataset - converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] - converter.inference_input_type = tf.uint8 - converter.inference_output_type = tf.uint8 + tflite_model = converter.convert() + return bytes(tflite_model) if isinstance(tflite_model, + bytearray) else tflite_model - tflite_model = converter.convert() - return bytes(tflite_model) if isinstance(tflite_model, - bytearray) else tflite_model + +def _inject_compression_metadata(model_data): + """Inject raw COMPRESSION_METADATA into a model's flatbuffer metadata. + + This simulates a legacy-compressed model (one that uses the + COMPRESSION_METADATA metadata entry and kernel-level decompression) without + going through compress(), which now produces DECODE-based output. + """ + model = model_editor.read(model_data) + model.metadata["COMPRESSION_METADATA"] = b"\x00" + return bytes(model.build()) + + +class LegacyCompressionDetectionTest(tf.test.TestCase): + """Test that legacy COMPRESSION_METADATA is rejected without the flag.""" def test_regular_model_loads_successfully(self): """Non-compressed models should load without issues.""" - model_data = self._create_test_model() + model_data = _create_test_model() interpreter = runtime.Interpreter.from_bytes(model_data) self.assertIsNotNone(interpreter) - def test_compressed_model_raises_runtime_error(self): - """Compressed models should raise RuntimeError when compression is disabled.""" - # Create and compress a model - model_data = self._create_test_model() + def test_legacy_compressed_model_raises_runtime_error(self): + """Models with COMPRESSION_METADATA should raise RuntimeError.""" + model_data = _create_test_model() + legacy_model = _inject_compression_metadata(model_data) - spec = (compression.SpecBuilder().add_tensor( - subgraph=0, tensor=1).with_lut(index_bitwidth=4).build()) - - compressed_model = compression.compress(model_data, spec) - if isinstance(compressed_model, bytearray): - compressed_model = bytes(compressed_model) - - # Should raise RuntimeError with self.assertRaises(RuntimeError): - runtime.Interpreter.from_bytes(compressed_model) - - def test_can_load_regular_after_compressed_failure(self): - """Verify we can still load regular models after compressed model fails.""" - model_data = self._create_test_model() + runtime.Interpreter.from_bytes(legacy_model) - # First try compressed model (should fail) - spec = (compression.SpecBuilder().add_tensor( - subgraph=0, tensor=1).with_lut(index_bitwidth=4).build()) - compressed_model = compression.compress(model_data, spec) + def test_can_load_regular_after_legacy_failure(self): + """Verify regular models still load after a legacy-compressed failure.""" + model_data = _create_test_model() + legacy_model = _inject_compression_metadata(model_data) with self.assertRaises(RuntimeError): - runtime.Interpreter.from_bytes(bytes(compressed_model)) + runtime.Interpreter.from_bytes(legacy_model) - # Then load regular model (should succeed) interpreter = runtime.Interpreter.from_bytes(model_data) self.assertIsNotNone(interpreter) if __name__ == '__main__': - # Set TF environment variables to suppress warnings + # Suppress TF C++ info/debug logs (0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR) os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' + # Disable oneDNN to avoid non-deterministic floating point results os.environ['TF_ENABLE_ONEDNN_OPTS'] = '0' tf.test.main() diff --git a/tensorflow/lite/micro/compression/BUILD b/tensorflow/lite/micro/compression/BUILD index 375a42d7a49..a9fe9fa36de 100644 --- a/tensorflow/lite/micro/compression/BUILD +++ b/tensorflow/lite/micro/compression/BUILD @@ -123,14 +123,15 @@ py_library( "compress.py", ], deps = [ - ":metadata_py", + ":compressor", + ":decode_insert", + ":huffman", + ":lut", ":model_editor", + ":pruning", ":spec", "//tensorflow/lite/micro/tools:tflite_flatbuffer_align", requirement("absl_py"), - requirement("flatbuffers"), - requirement("bitarray"), - requirement("numpy"), ], ) @@ -159,11 +160,53 @@ py_test( target_compatible_with = INCOMPATIBLE_WITH_WINDOWS, deps = [ ":compress", - ":metadata_py", + ":compressor", + ":decode_insert", ":model_editor", ":spec", "//tensorflow/lite/python:schema_py", - requirement("bitarray"), + requirement("numpy"), + ], +) + +tflm_py_test( + name = "compression_integration_test", + size = "small", + srcs = ["compression_integration_test.py"], + tags = [ + "noasan", + "nomsan", + "noubsan", + ], + target_compatible_with = INCOMPATIBLE_WITH_WINDOWS, + deps = [ + ":compress_lib", + ":decode_insert", + ":model_editor", + ":spec", + "//python/tflite_micro:runtime", + "//tensorflow/lite/python:schema_py", + requirement("numpy"), + ], +) + +tflm_py_test( + name = "proprietary_integration_test", + size = "small", + srcs = ["proprietary_integration_test.py"], + tags = [ + "manual", + "noasan", + "nomsan", + "noubsan", + ], + target_compatible_with = INCOMPATIBLE_WITH_WINDOWS, + deps = [ + ":compress_lib", + ":model_editor", + ":spec", + "//python/tflite_micro:runtime", + "//tensorflow/lite/python:schema_py", requirement("numpy"), ], ) @@ -326,6 +369,36 @@ tflm_py_library( ], ) +tflm_py_library( + name = "decode_insert", + srcs = ["decode_insert.py"], + deps = [ + ":compressor", + ":model_editor", + "//tensorflow/lite/python:schema_py", + ], +) + +tflm_py_test( + name = "decode_insert_test", + size = "small", + srcs = ["decode_insert_test.py"], + tags = [ + "noasan", + "nomsan", + "noubsan", + ], + deps = [ + ":compressor", + ":decode", + ":decode_insert", + ":lut", + ":model_editor", + "//tensorflow/lite/python:schema_py", + requirement("numpy"), + ], +) + tflm_py_binary( name = "view", srcs = [ diff --git a/tensorflow/lite/micro/compression/compress.py b/tensorflow/lite/micro/compression/compress.py index b6d5aef4435..96b55d94fd7 100644 --- a/tensorflow/lite/micro/compression/compress.py +++ b/tensorflow/lite/micro/compression/compress.py @@ -16,22 +16,22 @@ See USAGE. """ -import bitarray -import bitarray.util -from dataclasses import dataclass, field import os import sys import tempfile -from typing import ByteString, Iterable, Optional +import warnings +from typing import ByteString, Iterable, Type import absl.app import absl.flags -import flatbuffers -import numpy as np +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import decode_insert +from tflite_micro.tensorflow.lite.micro.compression import huffman +from tflite_micro.tensorflow.lite.micro.compression import lut from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.micro.compression import pruning from tflite_micro.tensorflow.lite.micro.compression import spec -from tflite_micro.tensorflow.lite.micro.compression import metadata_py_generated as schema from tflite_micro.tensorflow.lite.micro.tools import tflite_flatbuffer_align_wrapper USAGE = f"""\ @@ -49,221 +49,48 @@ {spec.EXAMPLE_YAML_SPEC} --- -The only compression method currently implemented is "lut", i.e., -Look-Up-Table. This method requires the tensor in the input model to have a -small number of unique values, fewer than or equal to 2**index_bitwidth. LUT -compression collects these values into a lookup table, and rewrites the tensor -as bitwidth-wide integer indices into that lookup table. Presumably, the input -model has been trained or preprocessed in a way that the tensor values -are binned into a meaningful, limited set. -""" - -# A compressed model augments the usual .tflite flatbuffer with a flatbuffer of -# its own containing compression metadata, stored at the buffer index stored at -# the following key in the .tflite flatbuffer's metadata map. -TFLITE_METADATA_KEY = "COMPRESSION_METADATA" - - -class CompressionError(Exception): - """Raised when compression fails for the reason documented in the message.""" - - def __init__(self, message, wrapped_exception=None): - super().__init__(f"{message}: {str(wrapped_exception)}") - self.original_exception = wrapped_exception - - -class _MetadataBuilder: - """Builder for the compression metadata flatbuffer.""" - - def __init__(self): - self._metadata = schema.MetadataT() - self._metadata.subgraphs = [] - - def compile(self) -> bytearray: - """Packs the metadata into a binary array and returns it. - """ - builder = flatbuffers.Builder(1 * 2**10) - root = self._metadata.Pack(builder) - builder.Finish(root) - return builder.Output() - - def subgraph(self, index: int): - """Return subgraph at index, adding subgraphs if necessary. - """ - while len(self._metadata.subgraphs) <= index: - self._add_subgraph() - return self._metadata.subgraphs[index] - - def add_lut_tensor(self, subgraph_id: int): - """Add LUT tensor to the given subgraph and return it. - """ - tensor = schema.LutTensorT() - self.subgraph(subgraph_id).lutTensors.append(tensor) - return tensor - - def _add_subgraph(self): - subgraph = schema.SubgraphT() - subgraph.lutTensors = [] - self._metadata.subgraphs.append(subgraph) - return subgraph - - -@dataclass -class _LutCompressedArray: - compression_axis: Optional[int] = None - lookup_tables: list[np.ndarray] = field(default_factory=list) - indices: np.ndarray = field(default_factory=lambda: np.array([])) - - @property - def index_bitwidth(self) -> int: - """Returns the number of bits required to encode the indices.""" - if self.indices is None: - raise ValueError - - max_index = int(np.max(self.indices)) - return max_index.bit_length() or 1 - - -def _lut_compress_array(tensor: np.ndarray, - axis: Optional[int]) -> _LutCompressedArray: - """Compresses the given tensor using lookup tables. - - Args: - tensor (np.ndarray): The tensor to be compressed. - - axis (Optional[int]): The axis along which to compress the tensor. If an - axis is given, a lookup table is created for each slice along the - axis. If axis is None, a single lookup table is used for the entire - tensor. - - Compressing a tensor with a lookup table per slice along a - particular axis is analogous to quantizing a tensor with different - quantization parameters per slice along a particular axis (dimension). - - Returns: - _LutCompressedArray: An object containing the compressed tensor data, - including the lookup tables and indices. - """ - compressed = _LutCompressedArray() - compressed.compression_axis = axis - - if axis is None: - # Compute unique values and indices for the entire tensor - values, indices = np.unique(tensor, return_inverse=True) - compressed.lookup_tables.append(values) - compressed.indices = indices.reshape(tensor.shape) - else: - # Iterate over slices along the compression axis - slice_indices = [] - for slice in np.moveaxis(tensor, axis, 0): - values, indices = np.unique(slice, return_inverse=True) - compressed.lookup_tables.append(values) - indices = indices.reshape(slice.shape) - slice_indices.append(indices) +Supported compression methods: - # Reconstruct a tensor of indices from the slices - stacked = np.stack(slice_indices, axis=0) - compressed.indices = np.moveaxis(stacked, 0, axis) + lut: Look-Up-Table compression. Requires the tensor to have a small number of + unique values, fewer than or equal to 2**index_bitwidth. LUT compression + collects these values into a lookup table, and rewrites the tensor as + bitwidth-wide integer indices into that lookup table. - return compressed + huffman: Huffman compression using Xtensa-format decode tables. (Not yet + implemented.) + pruning: Pruning (sparsity) compression for sparse tensors. (Not yet + implemented.) -def _check_lut_compression(compression) -> spec.LookUpTableCompression: - if len(compression) != 1: - raise CompressionError("Each tensor must have exactly one compression") - if not isinstance(compression[0], spec.LookUpTableCompression): - raise CompressionError('Only "lut" compression may be specified') - - return compression[0] - - -def _identify_compression_axis(tensor: model_editor.Tensor) -> Optional[int]: - """Determines the axis along which to compress. - - The axis along which to compress is inferred from the tensor's quantization - parameters. - - Returns: - The axis along which to compress, or None to indicate one value table for - the entire tensor. - - Raises: - CompressionError: If the axis cannot be determined. - """ - q = tensor.quantization - if q is not None: - # model_editor wraps quantization, access scales/axis from wrapper - scales = q.scales if isinstance(q.scales, list) else [q.scales] - quantization_channels = len(scales) - - if quantization_channels == 1: - # Use one value table for the entire tensor - return None - - if q.axis is not None and q.axis < len(tensor.shape): - if quantization_channels == tensor.shape[q.axis]: - return q.axis - - raise CompressionError( - f"Invalid or no quanitzation parameters from which to " - f"infer the axis along which tensor should be compressed.") - - -def _check_bitwidth(compressed: int, specified: int, spec: spec.Tensor): - """Applies business logic regarding specified bitwidth. - - It is an error if the bitwidth required to compress a tensor exceeds the - specified bitwith, and a warning if the tensor can be compressed in less than - the specified bitwidth. The latter is allowed, and is not an error, to permit - testing with larger bitwidths without re-binning a model. - """ - if compressed > specified: - raise CompressionError( - f"index_bitwidth too small: {compressed} bits needed to " - f"enumerate unique values in tensor specified in {spec}") - elif compressed < specified: - print( - f"warning: index_bitwidth too large: only {compressed} " - f"bits needed to enumerate unique values in tensor specified in {spec}", - file=sys.stderr) - - -def _pack_indices(indices: np.ndarray, bitwidth: int) -> bytes: - """Packs indices into a bytearray using bitwidth-sized fields. - """ - endianness = "big" - bits = bitarray.bitarray(endian=endianness) - for i in indices.ravel(): - bits.extend( - bitarray.util.int2ba(int(i), length=bitwidth, endian=endianness)) - return bits.tobytes() - +Compressed models use DECODE operators to decompress tensors at runtime. +""" -def _pack_lookup_tables(tables: list[np.ndarray], table_len: int) -> bytearray: - """Packs the value tables of a LutCompressedArray. +# Plugin dispatch table: maps CompressionMethod subclasses to compressor instances +_COMPRESSORS: dict[Type[spec.CompressionMethod], compressor.Compressor] = { + spec.LookUpTableCompression: lut.LutCompressor(), + spec.HuffmanCompression: huffman.HuffmanCompressor(), + spec.PruningCompression: pruning.PruningCompressor(), +} - Pack the value tables of a LutCompressedArray into a bytes object in the - format writable to a value_table buffer in the .tflite flatbuffer. The - tables are concatenated. - """ - buffer = bytearray() - for t in tables: - padding_needed = table_len - len(t) - padded = np.pad(t, (0, padding_needed), mode='constant', constant_values=0) - buffer.extend(padded.tobytes()) - return buffer +def _get_compressor(method: spec.CompressionMethod) -> compressor.Compressor: + """Get the compressor plugin for a given compression method.""" + compressor_instance = _COMPRESSORS.get(type(method)) + if compressor_instance is None: + raise compressor.CompressionError( + f"No compressor registered for {type(method).__name__}") + return compressor_instance def _apply_flatbuffer_alignment(model_bytes: bytearray) -> bytearray: """Applies proper FlatBuffer alignment to a model. - + The Python flatbuffers library doesn't respect `force_align` schema attributes, so we use the C++ wrapper which properly handles alignment requirements. - + Args: model_bytes: The model flatbuffer to align - + Returns: The properly aligned model flatbuffer """ @@ -295,45 +122,63 @@ def _apply_flatbuffer_alignment(model_bytes: bytearray) -> bytearray: def compress(model_in: ByteString, specs: Iterable[spec.Tensor]) -> bytearray: """Compresses a model .tflite flatbuffer. + Compresses tensors according to the given specs and inserts DECODE operators + to decompress them at runtime. + Args: model_in: the original, uncompressed .tflite flatbuffer specs: an iterable of compression specs, see module spec.py Returns: - A compressed flatbuffer. + A compressed flatbuffer with DECODE operators inserted. """ + specs = list(specs) + if not specs: + raise compressor.CompressionError( + "Compression spec is empty; no tensors to compress") + model = model_editor.read(model_in) - metadata = _MetadataBuilder() + compression_results: dict[tuple[int, int], compressor.CompressionResult] = {} - for spec in specs: + for tensor_spec in specs: try: - tensor = model.subgraphs[spec.subgraph].tensors[spec.tensor] - lut_compression = _check_lut_compression(spec.compression) - spec_bitwidth = lut_compression.index_bitwidth - axis = _identify_compression_axis(tensor) - compressed = _lut_compress_array(tensor.array, axis) - _check_bitwidth(compressed.index_bitwidth, spec_bitwidth, spec) - - # overwrite tensor data with indices - tensor.buffer.data = _pack_indices(compressed.indices, spec_bitwidth) - - # write value buffer - value_buffer_data = _pack_lookup_tables(compressed.lookup_tables, - 2**spec_bitwidth) - value_buffer = model_editor.Buffer(data=value_buffer_data) - model.buffers.append(value_buffer) # Auto-sets value_buffer.index - - # add compression metadata for tensor - lut_tensor = metadata.add_lut_tensor(subgraph_id=spec.subgraph) - lut_tensor.tensor = spec.tensor - lut_tensor.valueBuffer = value_buffer.index - lut_tensor.indexBitwidth = spec_bitwidth - + tensor = model.subgraphs[tensor_spec.subgraph].tensors[ + tensor_spec.tensor] + + # Currently only one compression method per tensor + if len(tensor_spec.compression) != 1: + raise compressor.CompressionError( + "Each tensor must have exactly one compression method") + + method = tensor_spec.compression[0] + plugin = _get_compressor(method) + original_size = len(tensor.buffer.data) if tensor.buffer.data else 0 + result = plugin.compress(tensor, method) + + compressed_size = len(result.encoded_data) + len(result.ancillary_data) + if compressed_size > original_size: + warnings.warn( + f"Compression of tensor {tensor.name!r} (subgraph " + f"{tensor_spec.subgraph}, tensor {tensor_spec.tensor}) resulted " + f"in expansion: {original_size} bytes -> {compressed_size} bytes " + f"(encoded: {len(result.encoded_data)}, " + f"ancillary: {len(result.ancillary_data)})", + stacklevel=2) + + # Replace tensor data with encoded data + tensor.buffer.data = result.encoded_data + + # Store result for DECODE insertion + compression_results[(tensor_spec.subgraph, tensor_spec.tensor)] = result + + except compressor.CompressionError: + raise except Exception as e: - raise CompressionError(f"error compressing {spec}") from e + raise compressor.CompressionError( + f"error compressing {tensor_spec}") from e - # add compression metadata to model - model.metadata[TFLITE_METADATA_KEY] = metadata.compile() + # Insert DECODE operators into the graph + decode_insert.insert_decode_operators(model, compression_results) # Build the model and apply proper alignment unaligned_model = model.build() diff --git a/tensorflow/lite/micro/compression/compress_test.py b/tensorflow/lite/micro/compression/compress_test.py index 81bbdab3293..6ee80f200d5 100644 --- a/tensorflow/lite/micro/compression/compress_test.py +++ b/tensorflow/lite/micro/compression/compress_test.py @@ -11,164 +11,21 @@ # 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. +"""Integration tests for the compression system.""" + +import warnings -import bitarray -import bitarray.util import numpy as np import unittest from tflite_micro.tensorflow.lite.micro.compression import compress -from tflite_micro.tensorflow.lite.micro.compression import metadata_py_generated as schema +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import decode_insert from tflite_micro.tensorflow.lite.micro.compression import model_editor from tflite_micro.tensorflow.lite.micro.compression import spec from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite -class TestPackIndices(unittest.TestCase): - - def test_basic_case(self): - indices = np.array([1, 2, 3]) - bitwidth = 4 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b0001_0010, 0b0011_0000]) - self.assertEqual(result, expected_bytes) - - def test_single_element(self): - indices = np.array([10]) - bitwidth = 8 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b0000_1010]) - self.assertEqual(result, expected_bytes) - - def test_different_bitwidth(self): - indices = np.array([1, 2, 3]) - bitwidth = 8 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b0000_0001, 0b0000_0010, 0b0000_0011]) - self.assertEqual(result, expected_bytes) - - def test_large_numbers(self): - indices = np.array([255, 128, 64]) - bitwidth = 8 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b1111_1111, 0b1000_0000, 0b0100_0000]) - self.assertEqual(result, expected_bytes) - - def test_multidimensional_array(self): - indices = np.array([[1, 2], [3, 4]]) - bitwidth = 4 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b0001_0010, 0b0011_0100]) - self.assertEqual(result, expected_bytes) - - def test_zero_bitwidth(self): - indices = np.array([0, 1, 2]) - bitwidth = 0 - with self.assertRaises(ValueError): - compress._pack_indices(indices, bitwidth) - - def test_empty_array(self): - indices = np.array([]) - bitwidth = 4 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = b"" - self.assertEqual(result, expected_bytes) - - def test_bitwidth_1(self): - indices = np.array([1, 0, 1, 1, 0, 1]) - bitwidth = 1 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b101101_00]) - self.assertEqual(result, expected_bytes) - - def test_bitwidth_2(self): - indices = np.array([1, 2, 3, 0]) - bitwidth = 2 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b01_10_11_00]) - self.assertEqual(result, expected_bytes) - - def test_bitwidth_3(self): - indices = np.array([1, 3, 5, 7]) - bitwidth = 3 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b001_011_10, 0b1_111_0000]) - self.assertEqual(result, expected_bytes) - - def test_bitwidth_5(self): - indices = np.array([1, 2, 16, 31]) - bitwidth = 5 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes([0b00001_000, 0b10_10000_1, 0b1111_0000]) - self.assertEqual(result, expected_bytes) - - def test_bitwidth_7(self): - indices = np.array([1, 64, 127, 32]) - bitwidth = 7 - result = compress._pack_indices(indices, bitwidth) - expected_bytes = bytes( - [0b0000001_1, 0b000000_11, 0b11111_010, 0b0000_0000]) - self.assertEqual(result, expected_bytes) - - -class TestPackLookupTables(unittest.TestCase): - - def test_int16_positive(self): - tables = [np.array([0x1234, 0x5678], dtype=' tuple[int, bitarray.bitarray, np.ndarray]: - """Helper: extracts the compressed tensor parts for a given spec. - - Returns: - bitwidth - indices - values - """ - subgraph_obj = self.compressed.subgraphs[subgraph] - tensor_obj = subgraph_obj.tensors[tensor] - lut_tensors = self.metadata.subgraphs[subgraph_obj.index].lutTensors - lut_tensor = next(t for t in lut_tensors if t.tensor == tensor_obj.index) - bitwidth = lut_tensor.indexBitwidth - - indices = bitarray.bitarray(buffer=tensor_obj.buffer.data, endian="big") - n_indices = np.prod(tensor_obj.shape) - indices = indices[:n_indices * bitwidth] # trim possible padding - - value_buffer = self.compressed.buffers[lut_tensor.valueBuffer] - values = np.frombuffer(value_buffer.data, dtype=tensor_obj.numpy_dtype) - - return bitwidth, indices, values - - def _make_indices(self, s: str) -> bitarray.bitarray: - """Helper: makes indices from "01" strings for use as expected values.""" - return bitarray.bitarray(s, endian="big") - - def test_compressed_uint8(self): - bitwidth, indices, values = self._get_compressed(subgraph=0, tensor=0) - self.assertEqual(bitwidth, 4) - - # yapf: disable - expected_indices = self._make_indices(""" - 0000 0001 0010 0011 - 0100 0101 0110 0111 - 1000 1001 1010 1011 - 1100 1101 1110 1111 - """) - # yapf: enable - self.assertEqual(indices, expected_indices) - - expected_values = np.array(range(16), dtype=" [FC1 with weights1] -> output1 + input2 -> [FC2 with weights2] -> intermediate -> [FC3 with weights1] -> output2 + + weights1 is shared between FC1 and FC3. weights2 is used only by FC2, which + runs between the two consumers of weights1. + """ + # 4 unique values per tensor for 2-bit LUT compression. Small values avoid + # saturation in chained layers. Different row sums produce varied outputs. + weights1_data = np.array([ + [-1, 0, 0, 1], + [-1, 0, 1, 1], + [-1, 1, 1, 1], + [0, 1, 1, 1], + ], + dtype=np.int8) + weights1 = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=weights1_data, + name="weights1", + quantization=model_editor.Quantization(scales=1.0, zero_points=0), + ) + + weights2_data = np.array([ + [1, 1, 1, 1], + [1, 1, 2, 2], + [1, 2, 2, 3], + [2, 2, 3, 3], + ], + dtype=np.int8) + weights2 = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=weights2_data, + name="weights2", + quantization=model_editor.Quantization(scales=1.0, zero_points=0), + ) + + # All tensors need matching quantization for FULLY_CONNECTED + quant = model_editor.Quantization(scales=1.0, zero_points=0) + + input1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input1", + quantization=quant, + ) + input2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input2", + quantization=quant, + ) + output1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output1", + quantization=quant, + ) + intermediate = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="intermediate", + quantization=quant, + ) + output2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output2", + quantization=quant, + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights1, weights2], + inputs=[input1, input2], + outputs=[output1, output2], + operators=[ + # FC1: uses weights1 + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input1, weights1], + outputs=[output1], + ), + # FC2: uses weights2 (runs between FC1 and FC3) + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input2, weights2], + outputs=[intermediate], + ), + # FC3: uses weights1 (second consumer, after DECODE(weights2)) + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[intermediate, weights1], + outputs=[output2], + ), + ], + ) + ]) + return model.build() + + +class AltDecompressionMemoryTest(unittest.TestCase): + """Tests for alternate decompression memory with shared compressed tensors. + + These tests verify correct behavior when compressed tensors are shared + between multiple operators and alternate decompression memory is enabled. + """ + + def test_shared_compressed_tensor_with_alt_memory(self): + """Verify correct results when a shared compressed tensor is used with alt + decompression memory. + + This test uses a graph where a compressed tensor (weights1) is consumed by + two operators (FC1 and FC3), with an intervening DECODE of a different + compressed tensor (weights2) between them. + + The interpreter's alternate decompression memory has a limitation: each + DECODE's Prepare resets the allocation offset to zero. This means all + DECODE outputs are allocated at the same address, so they overwrite each + other. A DECODE output can only be used until the next DECODE runs. + + To work around this limitation, the DECODE insertion code inserts a + separate DECODE immediately before each consumer of a compressed tensor, + rather than sharing one DECODE output among all consumers. + """ + flatbuffer = _build_shared_weights_model() + + specs = [ + spec.Tensor( + subgraph=0, + tensor=0, # weights1 + compression=[spec.LookUpTableCompression(index_bitwidth=2)], + ), + spec.Tensor( + subgraph=0, + tensor=1, # weights2 + compression=[spec.LookUpTableCompression(index_bitwidth=2)], + ), + ] + + compressed_fb = compress.compress(flatbuffer, specs) + + # Run without alt decompression memory (baseline) + interp_no_alt = runtime.Interpreter.from_bytes(bytes(compressed_fb)) + + # Run with alt decompression memory + interp_with_alt = runtime.Interpreter.from_bytes( + bytes(compressed_fb), + alt_decompression_memory_size=256, + ) + + test_input1 = np.array([[1, 1, 1, 1]], dtype=np.int8) + test_input2 = np.array([[1, 1, 1, 1]], dtype=np.int8) + + interp_no_alt.set_input(test_input1, 0) + interp_no_alt.set_input(test_input2, 1) + interp_no_alt.invoke() + expected1 = interp_no_alt.get_output(0) + expected2 = interp_no_alt.get_output(1) + + interp_with_alt.set_input(test_input1, 0) + interp_with_alt.set_input(test_input2, 1) + interp_with_alt.invoke() + actual1 = interp_with_alt.get_output(0) + actual2 = interp_with_alt.get_output(1) + + np.testing.assert_array_equal( + expected1, actual1, "Output 1 mismatch with alt decompression memory") + np.testing.assert_array_equal( + expected2, actual2, "Output 2 mismatch with alt decompression memory") + + +class HuffmanCompressionTest(unittest.TestCase): + """Integration tests for Huffman compression.""" + + @unittest.skip("Huffman compression not yet implemented") + def test_huffman_compressed_model_matches_uncompressed(self): + """Huffman-compressed model produces same outputs as uncompressed.""" + pass + + @unittest.skip("Huffman compression not yet implemented") + def test_huffman_decode_operators_present(self): + """DECODE operators are inserted for Huffman-compressed tensors.""" + pass + + @unittest.skip("Huffman compression not yet implemented") + def test_huffman_compressed_model_is_smaller(self): + """Huffman-compressed model is smaller than original.""" + pass + + +class PruningCompressionTest(unittest.TestCase): + """Integration tests for pruning compression.""" + + @unittest.skip("Pruning compression not yet implemented") + def test_pruning_compressed_model_matches_uncompressed(self): + """Pruning-compressed model produces same outputs as uncompressed.""" + pass + + @unittest.skip("Pruning compression not yet implemented") + def test_pruning_decode_operators_present(self): + """DECODE operators are inserted for pruning-compressed tensors.""" + pass + + @unittest.skip("Pruning compression not yet implemented") + def test_pruning_compressed_model_is_smaller(self): + """Pruning-compressed model is smaller than original.""" + pass + + +if __name__ == "__main__": + # Suppress TF C++ info/debug logs (0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR) + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" + # Disable oneDNN to avoid non-deterministic floating point results + os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" + unittest.main() diff --git a/tensorflow/lite/micro/compression/decode_insert.py b/tensorflow/lite/micro/compression/decode_insert.py new file mode 100644 index 00000000000..fa91896e538 --- /dev/null +++ b/tensorflow/lite/micro/compression/decode_insert.py @@ -0,0 +1,280 @@ +# Copyright 2026 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. +"""DECODE operator insertion into TFLite model graphs. + +This module inserts DECODE operators into a compressed model. DECODE operators +transform encoded tensors (with their paired ancillary data tensors) into +tensors ready for use by downstream operators. + +The DECODE operator is registered as a custom operator named "TFLM_DECODE". +Each DECODE output requires two inputs: the encoded tensor and the ancillary +data tensor (containing the DCM header and decode-type-specific data). +""" + +import warnings +from collections import defaultdict +from dataclasses import dataclass +from typing import Optional + +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + +# Custom operator name for DECODE +DECODE_CUSTOM_OP_NAME = "TFLM_DECODE" + + +@dataclass +class _CompressedTensorInfo: + """Information about a compressed tensor for DECODE insertion.""" + subgraph_idx: int + tensor_idx: int + tensor: model_editor.Tensor + encoded_data: bytes + ancillary_data: bytes + consumers: list[model_editor.Operator] + + +def _find_tensor_consumers( + subgraph: model_editor.Subgraph, + tensor: model_editor.Tensor, +) -> list[model_editor.Operator]: + """Find all operators in subgraph that use tensor as an input.""" + consumers = [] + for op in subgraph.operators: + if tensor in op.inputs: + consumers.append(op) + return consumers + + +def _create_ancillary_tensor( + ancillary_data: bytes, + original_tensor: model_editor.Tensor, +) -> model_editor.Tensor: + """Create an ancillary data tensor for a compressed tensor. + + Args: + ancillary_data: The complete ancillary data (DCM + type-specific data). + original_tensor: The original tensor being decoded, for naming. + + Returns: + A new Tensor containing the ancillary data. + """ + name = None + if original_tensor.name: + name = f"{original_tensor.name}_ancillary" + + return model_editor.Tensor( + shape=(len(ancillary_data), ), + dtype=tflite.TensorType.UINT8, + data=ancillary_data, + name=name, + ) + + +def _create_output_tensor( + original_tensor: model_editor.Tensor, ) -> model_editor.Tensor: + """Create the output tensor for a DECODE operator. + + The output tensor has the same shape, dtype, and quantization as the + original tensor would have when decoded. It has no data---the DECODE + operator produces its values at runtime. + + Args: + original_tensor: The original tensor being decoded. + + Returns: + A new Tensor for the DECODE output. + """ + name = None + if original_tensor.name: + name = f"{original_tensor.name}_decoded" + + return model_editor.Tensor( + shape=original_tensor.shape, + dtype=original_tensor.dtype, + quantization=original_tensor.quantization, + name=name, + ) + + +def _rewire_consumers( + consumers: list[model_editor.Operator], + old_tensor: model_editor.Tensor, + new_tensor: model_editor.Tensor, +) -> None: + """Replace old_tensor with new_tensor in all consumer inputs.""" + for consumer in consumers: + consumer.inputs = [ + new_tensor if t is old_tensor else t for t in consumer.inputs + ] + + +def _rewrite_encoded_tensor( + tensor: model_editor.Tensor, + encoded_data: bytes, +) -> None: + """Rewrite a compressed tensor to hold encoded data. + + The original tensor contained uncompressed values with quantization. After + compression, it holds packed indices (or other encoded form) as raw bytes. + This function updates the tensor in place to reflect its new role. + + Args: + tensor: The tensor to rewrite. + encoded_data: The compressed/encoded data bytes. + """ + tensor.shape = (len(encoded_data), ) + tensor.dtype = tflite.TensorType.UINT8 + tensor.quantization = None + tensor.buffer.data = encoded_data + + +def insert_decode_operators( + model: model_editor.Model, + compression_results: dict[tuple[int, int], compressor.CompressionResult], +) -> None: + """Insert DECODE operators for all compressed tensors. + + This function modifies the model in-place, inserting DECODE operators + before any operator that uses a compressed tensor as input. + + A separate DECODE is inserted before each consumer, rather than sharing one + DECODE output among all consumers. This is required because the interpreter's + alternate decompression memory resets its allocation offset for each DECODE's + Prepare, causing all DECODE outputs to be allocated at the same address. If + two consumers share one DECODE and another DECODE runs between them, the + intervening DECODE overwrites the shared output, corrupting data for the + second consumer. + + For each consumer of a compressed tensor: + 1. Create an ancillary data tensor containing DCM + type-specific data + 2. Create an output tensor with the same shape/dtype as the decoded tensor + 3. Insert a DECODE operator immediately before the consumer + 4. Rewire the consumer to use the DECODE output + + Args: + model: The model to modify in-place. + compression_results: Map from (subgraph_idx, tensor_idx) to the + CompressionResult containing ancillary_data. + """ + # Group compressed tensors by subgraph + by_subgraph: dict[int, list[_CompressedTensorInfo]] = defaultdict(list) + + for (sg_idx, tensor_idx), result in compression_results.items(): + subgraph = model.subgraphs[sg_idx] + tensor = subgraph.tensors[tensor_idx] + consumers = _find_tensor_consumers(subgraph, tensor) + + if not consumers: + # Check if tensor is a subgraph output + is_output = tensor in subgraph.outputs + if is_output: + # TODO: Handle compressed tensors that are subgraph outputs. + # This occurs in multi-subgraph models using IF/WHILE where a + # compressed tensor flows out of a subgraph. + raise NotImplementedError( + f"Compressed tensor {tensor.name!r} (subgraph {sg_idx}, " + f"tensor {tensor_idx}) is a subgraph output with no consumers. " + "Compressed subgraph outputs are not yet supported.") + else: + warnings.warn( + f"Compressed tensor {tensor.name!r} (subgraph {sg_idx}, " + f"tensor {tensor_idx}) has no consumers and is not a subgraph " + "output. No DECODE operator will be inserted.", + stacklevel=2) + continue + + info = _CompressedTensorInfo( + subgraph_idx=sg_idx, + tensor_idx=tensor_idx, + tensor=tensor, + encoded_data=result.encoded_data, + ancillary_data=result.ancillary_data, + consumers=consumers, + ) + by_subgraph[sg_idx].append(info) + + # Process each subgraph + for sg_idx, tensor_infos in by_subgraph.items(): + subgraph = model.subgraphs[sg_idx] + + # Group tensor infos by consumer so multiple compressed inputs to the + # same operator get batched into a single DECODE. + consumer_to_infos: dict[model_editor.Operator, list[_CompressedTensorInfo]] + consumer_to_infos = defaultdict(list) + for info in tensor_infos: + for consumer in info.consumers: + if info not in consumer_to_infos[consumer]: + consumer_to_infos[consumer].append(info) + + # Sort consumers by position in reverse so insertions don't invalidate + # earlier positions. + sorted_consumers = sorted( + consumer_to_infos.keys(), + key=lambda op: subgraph.operators.index(op), + reverse=True, + ) + + # Cache ancillary tensors by original tensor to avoid duplicates. Each + # DECODE needs its own output tensor, but ancillary data is identical for + # all DECODEs of the same compressed tensor. + ancillary_cache: dict[model_editor.Tensor, model_editor.Tensor] = {} + + # Track tensors to rewrite after all output tensors are created, since + # _create_output_tensor reads the original tensor's shape/dtype/quantization. + tensors_to_rewrite: dict[model_editor.Tensor, bytes] = {} + + for consumer in sorted_consumers: + decode_inputs = [] + decode_outputs = [] + + for info in consumer_to_infos[consumer]: + # Reuse or create ancillary data tensor + if info.tensor not in ancillary_cache: + ancillary_tensor = _create_ancillary_tensor( + info.ancillary_data, + info.tensor, + ) + subgraph.tensors.append(ancillary_tensor) + ancillary_cache[info.tensor] = ancillary_tensor + tensors_to_rewrite[info.tensor] = info.encoded_data + else: + ancillary_tensor = ancillary_cache[info.tensor] + + # Create output tensor (one per compressed input) + output_tensor = _create_output_tensor(info.tensor) + subgraph.tensors.append(output_tensor) + + decode_inputs.extend([info.tensor, ancillary_tensor]) + decode_outputs.append(output_tensor) + + # Rewire this consumer to use the decoded output + _rewire_consumers([consumer], info.tensor, output_tensor) + + # Create single DECODE operator for all compressed inputs + decode_op = model_editor.Operator( + opcode=tflite.BuiltinOperator.CUSTOM, + custom_code=DECODE_CUSTOM_OP_NAME, + inputs=decode_inputs, + outputs=decode_outputs, + ) + + # Insert DECODE immediately before this consumer + insert_pos = subgraph.operators.index(consumer) + subgraph.operators.insert(insert_pos, decode_op) + + # Rewrite encoded tensors after all output tensors are created + for tensor, encoded_data in tensors_to_rewrite.items(): + _rewrite_encoded_tensor(tensor, encoded_data) diff --git a/tensorflow/lite/micro/compression/decode_insert_test.py b/tensorflow/lite/micro/compression/decode_insert_test.py new file mode 100644 index 00000000000..60965b46676 --- /dev/null +++ b/tensorflow/lite/micro/compression/decode_insert_test.py @@ -0,0 +1,559 @@ +# Copyright 2026 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. +"""Unit tests for DECODE operator insertion.""" + +import unittest +import warnings + +import numpy as np + +from tflite_micro.tensorflow.lite.micro.compression import compressor +from tflite_micro.tensorflow.lite.micro.compression import decode +from tflite_micro.tensorflow.lite.micro.compression import decode_insert +from tflite_micro.tensorflow.lite.micro.compression import lut +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + + +def _build_simple_fc_model(): + """Build a simple model with one FC operator and compressible weights.""" + # yapf: disable + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.array([[1, 2, 1, 2], + [3, 4, 3, 4], + [1, 2, 1, 2], + [3, 4, 3, 4]], dtype=np.int8), + name="weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + # yapf: enable + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, weights], + outputs=[output_t], + ) + ], + ) + ]) + return model + + +def _build_shared_weights_model(): + """Build model where one tensor is used by multiple operators.""" + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="shared_weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input1", + ) + input2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input2", + ) + output1 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output1", + ) + output2 = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output2", + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input1, weights], + outputs=[output1], + ), + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input2, weights], + outputs=[output2], + ), + ], + ) + ]) + return model + + +def _make_dummy_ancillary_data(bitwidth=4) -> bytes: + """Create dummy ancillary data for testing.""" + n_entries = 1 << bitwidth + value_tables = bytes(range(1, n_entries + 1)) + value_tables += b'\x00' * ((-len(value_tables)) % 16) + + lut_data = lut.LutAncillaryData( + bitwidth=bitwidth, + value_table_stride=n_entries, + value_tables=value_tables, + ) + dcm = decode.DecodeCommonMetadata( + decode_type=decode.DecodeType.LUT, + user_data=lut_data.to_user_data(), + ) + return dcm.to_bytes() + lut_data.to_bytes() + + +class TestDecodeInsertion(unittest.TestCase): + """Tests for insert_decode_operators function.""" + + def test_insert_single_decode_operator(self): + """DECODE operator inserted before FC that uses compressed weights.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + # Create compression result + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + # Insert DECODE operators + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # Should have 2 operators: DECODE then FC + self.assertEqual(len(sg.operators), 2) + self.assertEqual(sg.operators[0].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[0].custom_code, + decode_insert.DECODE_CUSTOM_OP_NAME) + self.assertEqual(sg.operators[1].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + + def test_decode_inputs_structure(self): + """DECODE operator has correct inputs: encoded tensor + ancillary.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + + # DECODE has 2 inputs + self.assertEqual(len(decode_op.inputs), 2) + # First input is the encoded tensor (original weights) + self.assertIs(decode_op.inputs[0], weights_tensor) + # Second input is ancillary tensor + self.assertEqual(decode_op.inputs[1].dtype, tflite.TensorType.UINT8) + + def test_decode_output_structure(self): + """DECODE operator output has correct shape and dtype.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + # Save original properties before rewrite + original_shape = weights_tensor.shape + original_dtype = weights_tensor.dtype + + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + output = decode_op.outputs[0] + + # Output matches original (pre-rewrite) tensor shape and dtype + self.assertEqual(output.shape, original_shape) + self.assertEqual(output.dtype, original_dtype) + + def test_consumer_rewired_to_decode_output(self): + """FC operator input rewired to use DECODE output.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + fc_op = model.subgraphs[0].operators[1] + + # FC's second input (weights) should now be DECODE's output + self.assertIs(fc_op.inputs[1], decode_op.outputs[0]) + # Original weights tensor should NOT be in FC inputs + self.assertNotIn(weights_tensor, fc_op.inputs) + + def test_shared_tensor_decode_per_consumer(self): + """Tensor used by multiple ops gets separate DECODE for each consumer.""" + model = _build_shared_weights_model() + weights_tensor = model.subgraphs[0].tensor_by_name("shared_weights") + + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # Should have 4 operators: 2 DECODEs + 2 FCs (DECODE before each FC) + self.assertEqual(len(sg.operators), 4) + self.assertEqual(sg.operators[0].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[1].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + self.assertEqual(sg.operators[2].opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(sg.operators[3].opcode, + tflite.BuiltinOperator.FULLY_CONNECTED) + + decode_op1 = sg.operators[0] + fc_op1 = sg.operators[1] + decode_op2 = sg.operators[2] + fc_op2 = sg.operators[3] + + # Each FC should use its own DECODE's output + self.assertIs(fc_op1.inputs[1], decode_op1.outputs[0]) + self.assertIs(fc_op2.inputs[1], decode_op2.outputs[0]) + # The two DECODEs should have different outputs + self.assertIsNot(decode_op1.outputs[0], decode_op2.outputs[0]) + # The two DECODEs should share the same ancillary tensor + self.assertIs(decode_op1.inputs[1], decode_op2.inputs[1]) + + def test_ancillary_tensor_contains_dcm(self): + """Ancillary tensor data contains valid DCM header.""" + model = _build_simple_fc_model() + + ancillary_data = _make_dummy_ancillary_data() + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=ancillary_data, + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + ancillary_tensor = decode_op.inputs[1] + + # Ancillary tensor data should match what we provided + self.assertEqual(bytes(ancillary_tensor.array), ancillary_data) + + # Verify DCM header + dcm_bytes = ancillary_tensor.array[:16] + self.assertEqual(dcm_bytes[0], 0) # decode_type = LUT + self.assertEqual(dcm_bytes[1], 1) # DCM version + + def test_no_consumers_no_decode(self): + """Tensor with no consumers gets no DECODE operator and emits warning.""" + # Create model where compressed tensor is not used as input + unused_tensor = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="unused", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + input_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="input", + ) + output_t = model_editor.Tensor( + shape=(1, 4), + dtype=tflite.TensorType.INT8, + name="output", + ) + other_weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="other_weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[unused_tensor, other_weights], + operators=[ + model_editor.Operator( + opcode=tflite.BuiltinOperator.FULLY_CONNECTED, + inputs=[input_t, other_weights], + outputs=[output_t], + ) + ], + ) + ]) + + # Compress the unused tensor + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + decode_insert.insert_decode_operators(model, compression_results) + + # Should emit a warning about no consumers + self.assertEqual(len(w), 1) + self.assertIn("no consumers", str(w[0].message)) + self.assertIn("unused", str(w[0].message)) + + # Should still have just 1 operator (no DECODE inserted) + self.assertEqual(len(model.subgraphs[0].operators), 1) + + def test_tensor_naming(self): + """Output and ancillary tensors get appropriate names.""" + model = _build_simple_fc_model() + + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x00', + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + decode_op = model.subgraphs[0].operators[0] + ancillary = decode_op.inputs[1] + output = decode_op.outputs[0] + + self.assertEqual(ancillary.name, "weights_ancillary") + self.assertEqual(output.name, "weights_decoded") + + def test_multiple_compressed_inputs_batched(self): + """CONCATENATION with two compressed inputs gets one batched DECODE.""" + weights_a = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="weights_a", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + weights_b = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="weights_b", + quantization=model_editor.Quantization(scales=0.25, zero_points=0), + ) + output_t = model_editor.Tensor( + shape=(4, 8), + dtype=tflite.TensorType.INT8, + name="output", + ) + + concat_op = model_editor.Operator( + opcode=tflite.BuiltinOperator.CONCATENATION, + inputs=[weights_a, weights_b], + outputs=[output_t], + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights_a, weights_b], + operators=[concat_op], + ) + ]) + + ancillary_a = _make_dummy_ancillary_data(bitwidth=2) + ancillary_b = _make_dummy_ancillary_data(bitwidth=4) + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x01', + ancillary_data=ancillary_a, + ), + (0, 1): + compressor.CompressionResult( + encoded_data=b'\x02\x03', + ancillary_data=ancillary_b, + ), + } + + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # One DECODE + one CONCATENATION + self.assertEqual(len(sg.operators), 2) + decode_op = sg.operators[0] + self.assertEqual(decode_op.opcode, tflite.BuiltinOperator.CUSTOM) + self.assertEqual(decode_op.custom_code, + decode_insert.DECODE_CUSTOM_OP_NAME) + + # DECODE has 4 inputs (enc_a, anc_a, enc_b, anc_b) and 2 outputs + self.assertEqual(len(decode_op.inputs), 4) + self.assertEqual(len(decode_op.outputs), 2) + + # Each ancillary tensor carries its own distinct data + self.assertNotEqual(ancillary_a, ancillary_b) + self.assertEqual(bytes(decode_op.inputs[1].array), ancillary_a) + self.assertEqual(bytes(decode_op.inputs[3].array), ancillary_b) + + # CONCATENATION rewired to DECODE outputs + self.assertIs(sg.operators[1].inputs[0], decode_op.outputs[0]) + self.assertIs(sg.operators[1].inputs[1], decode_op.outputs[1]) + + def test_mixed_compressed_and_uncompressed_inputs(self): + """CONCATENATION with one compressed and one plain input.""" + weights = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.ones((4, 4), dtype=np.int8), + name="weights", + quantization=model_editor.Quantization(scales=0.5, zero_points=0), + ) + plain = model_editor.Tensor( + shape=(4, 4), + dtype=tflite.TensorType.INT8, + data=np.zeros((4, 4), dtype=np.int8), + name="plain", + ) + output_t = model_editor.Tensor( + shape=(4, 8), + dtype=tflite.TensorType.INT8, + name="output", + ) + + concat_op = model_editor.Operator( + opcode=tflite.BuiltinOperator.CONCATENATION, + inputs=[weights, plain], + outputs=[output_t], + ) + + model = model_editor.Model(subgraphs=[ + model_editor.Subgraph( + tensors=[weights, plain], + operators=[concat_op], + ) + ]) + + # Only compress weights, not plain + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=b'\x00\x01', + ancillary_data=_make_dummy_ancillary_data(), + ), + } + + decode_insert.insert_decode_operators(model, compression_results) + + sg = model.subgraphs[0] + + # One DECODE + one CONCATENATION + self.assertEqual(len(sg.operators), 2) + decode_op = sg.operators[0] + + # DECODE has 2 inputs and 1 output (only the compressed tensor) + self.assertEqual(len(decode_op.inputs), 2) + self.assertEqual(len(decode_op.outputs), 1) + + # CONCATENATION: first input rewired to DECODE output, second unchanged + self.assertIs(sg.operators[1].inputs[0], decode_op.outputs[0]) + self.assertIs(sg.operators[1].inputs[1], plain) + + def test_encoded_tensor_rewritten(self): + """Compressed tensor is rewritten with encoded data, UINT8 type, no quant.""" + model = _build_simple_fc_model() + weights_tensor = model.subgraphs[0].tensor_by_name("weights") + + encoded_data = b'\xAB\xCD\xEF' + compression_results = { + (0, 0): + compressor.CompressionResult( + encoded_data=encoded_data, + ancillary_data=_make_dummy_ancillary_data(), + ) + } + + decode_insert.insert_decode_operators(model, compression_results) + + # Original tensor should be rewritten + self.assertEqual(weights_tensor.shape, (len(encoded_data), )) + self.assertEqual(weights_tensor.dtype, tflite.TensorType.UINT8) + self.assertIsNone(weights_tensor.quantization) + self.assertEqual(weights_tensor.buffer.data, encoded_data) + + +class TestHelperFunctions(unittest.TestCase): + """Tests for internal helper functions.""" + + def test_find_tensor_consumers(self): + """_find_tensor_consumers finds all ops using a tensor.""" + model = _build_shared_weights_model() + sg = model.subgraphs[0] + weights = sg.tensor_by_name("shared_weights") + + consumers = decode_insert._find_tensor_consumers(sg, weights) + + self.assertEqual(len(consumers), 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tensorflow/lite/micro/compression/huffman.py b/tensorflow/lite/micro/compression/huffman.py index 40d0be9284a..e539827eae4 100644 --- a/tensorflow/lite/micro/compression/huffman.py +++ b/tensorflow/lite/micro/compression/huffman.py @@ -25,7 +25,7 @@ from tflite_micro.tensorflow.lite.micro.compression import spec -class HuffmanCompressor: +class HuffmanCompressor(compressor.Compressor): """Huffman compression plugin (stub). This stub exists to validate the plugin architecture. The actual Huffman diff --git a/tensorflow/lite/micro/compression/lut.py b/tensorflow/lite/micro/compression/lut.py index def34059ac5..991288f54cc 100644 --- a/tensorflow/lite/micro/compression/lut.py +++ b/tensorflow/lite/micro/compression/lut.py @@ -241,7 +241,7 @@ def pack_lookup_tables(tables: list[np.ndarray], table_len: int) -> bytes: return bytes(buffer) -class LutCompressor: +class LutCompressor(compressor.Compressor): """LUT compression plugin implementing the Compressor protocol.""" @property diff --git a/tensorflow/lite/micro/compression/proprietary_integration_test.py b/tensorflow/lite/micro/compression/proprietary_integration_test.py new file mode 100644 index 00000000000..684805d0f56 --- /dev/null +++ b/tensorflow/lite/micro/compression/proprietary_integration_test.py @@ -0,0 +1,211 @@ +# Copyright 2026 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. +"""Integration tests for compression using proprietary models. + +These tests verify that compressed models produce correct inference results +when run through the TFLM Python interpreter. Tests compress models and +compare outputs against uncompressed originals using random inputs. + +This test is tagged `manual` and requires a path to a directory containing +.tflite model files. + +Usage: + bazel test //tensorflow/lite/micro/compression:proprietary_integration_test \ + --//:with_compression \ + --test_arg=/path/to/models + +Required files: + Each model requires a compression spec file: + model.spec.yaml (replacing .tflite extension) + + See spec.py for the YAML format. Example: + tensors: + - subgraph: 0 + tensor: 2 + compression: + - lut: + index_bitwidth: 4 + +Optional files: + model.config.json (replacing .tflite extension) + Tolerance overrides: {"rtol": 1e-5, "atol": 1e-6} + Default is exact match (rtol=0, atol=0). +""" + +import glob +import json +import os +import sys +import unittest + +import numpy as np + +from tflite_micro.python.tflite_micro import runtime +from tflite_micro.tensorflow.lite.micro.compression import compress +from tflite_micro.tensorflow.lite.micro.compression import model_editor +from tflite_micro.tensorflow.lite.micro.compression import spec +from tflite_micro.tensorflow.lite.python import schema_py_generated as tflite + + +def _dtype_to_numpy(dtype: tflite.TensorType) -> np.dtype: + """Convert TFLite dtype to numpy dtype.""" + type_map = { + tflite.TensorType.INT8: np.int8, + tflite.TensorType.INT16: np.int16, + tflite.TensorType.INT32: np.int32, + tflite.TensorType.INT64: np.int64, + tflite.TensorType.UINT8: np.uint8, + tflite.TensorType.UINT16: np.uint16, + tflite.TensorType.UINT32: np.uint32, + tflite.TensorType.FLOAT16: np.float16, + tflite.TensorType.FLOAT32: np.float32, + tflite.TensorType.FLOAT64: np.float64, + tflite.TensorType.BOOL: np.bool_, + } + return type_map.get(dtype, np.uint8) + + +class ProprietaryModelTest(unittest.TestCase): + """Integration tests using proprietary models.""" + + # Parsed from command line in main() + models_dir = None + + @classmethod + def setUpClass(cls): + if not cls.models_dir: + raise unittest.SkipTest( + "No models directory provided. " + "Usage: bazel test ... --test_arg=/path/to/models") + + cls.model_paths = sorted( + glob.glob(os.path.join(cls.models_dir, '*.tflite'))) + if not cls.model_paths: + raise unittest.SkipTest(f"No .tflite files found in {cls.models_dir}") + + def test_all_models(self): + """Run compression test on each discovered model.""" + for model_path in self.model_paths: + with self.subTest(model=os.path.basename(model_path)): + self._test_model_compression(model_path) + + def _test_model_compression(self, model_path): + """Test that a compressed model produces same outputs as original.""" + with open(model_path, 'rb') as f: + flatbuffer = f.read() + + # Load compression spec from sidecar file + specs = self._load_compression_spec(model_path) + + # Load tolerance config + rtol, atol = self._load_tolerance(model_path) + + # Compress the model + compressed_fb = compress.compress(flatbuffer, specs) + + # Create interpreters + original_interp = runtime.Interpreter.from_bytes(bytes(flatbuffer)) + compressed_interp = runtime.Interpreter.from_bytes(bytes(compressed_fb)) + + # Generate random inputs and compare outputs + np.random.seed(42) + model = model_editor.read(flatbuffer) + sg = model.subgraphs[0] + + for trial in range(5): + # Set inputs + for i, input_tensor in enumerate(sg.inputs): + test_input = self._generate_input(input_tensor) + original_interp.set_input(test_input, i) + compressed_interp.set_input(test_input, i) + + # Run inference + original_interp.invoke() + compressed_interp.invoke() + + # Compare outputs + for i in range(len(sg.outputs)): + expected = original_interp.get_output(i) + actual = compressed_interp.get_output(i) + self._compare_outputs(expected, actual, rtol, atol, + f"trial {trial}, output {i}") + + def _generate_input(self, tensor): + """Generate random input respecting tensor dtype.""" + shape = tensor.shape + dtype = _dtype_to_numpy(tensor.dtype) + + if np.issubdtype(dtype, np.floating): + return np.random.uniform(-1.0, 1.0, shape).astype(dtype) + elif np.issubdtype(dtype, np.integer): + info = np.iinfo(dtype) + return np.random.randint(info.min, info.max + 1, shape, dtype=dtype) + elif dtype == np.bool_: + return np.random.choice([False, True], shape) + return np.zeros(shape, dtype=dtype) + + def _load_compression_spec(self, model_path): + """Load compression spec from sidecar YAML file. + + Raises: + FileNotFoundError: If no spec file is found. + """ + spec_path = model_path.replace('.tflite', '.spec.yaml') + if os.path.exists(spec_path): + with open(spec_path) as f: + return spec.parse_yaml(f.read()) + + raise FileNotFoundError( + f"No compression spec file found for {model_path}. " + f"Expected: {spec_path}") + + def _load_tolerance(self, model_path): + """Load tolerance from sidecar config if present. + + Returns (0, 0) for exact match if no config file exists. + """ + config_path = model_path.replace('.tflite', '.config.json') + if os.path.exists(config_path): + with open(config_path) as f: + config = json.load(f) + return config.get('rtol', 0), config.get('atol', 0) + return 0, 0 + + def _compare_outputs(self, expected, actual, rtol, atol, context=""): + """Compare outputs with optional tolerance.""" + msg = f"Output mismatch ({context})" if context else "Output mismatch" + if rtol == 0 and atol == 0: + np.testing.assert_array_equal(expected, actual, err_msg=msg) + else: + np.testing.assert_allclose(expected, + actual, + rtol=rtol, + atol=atol, + err_msg=msg) + + +if __name__ == "__main__": + # Suppress TF C++ info/debug logs (0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR) + os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" + # Disable oneDNN to avoid non-deterministic floating point results + os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0" + + # Parse models directory from args, then strip it so tf.test doesn't see it + for arg in sys.argv[1:]: + if not arg.startswith('-') and os.path.isdir(arg): + ProprietaryModelTest.models_dir = arg + sys.argv.remove(arg) + break + + unittest.main() diff --git a/tensorflow/lite/micro/compression/pruning.py b/tensorflow/lite/micro/compression/pruning.py index 2181b73e34a..5c95e3e87e9 100644 --- a/tensorflow/lite/micro/compression/pruning.py +++ b/tensorflow/lite/micro/compression/pruning.py @@ -25,7 +25,7 @@ from tflite_micro.tensorflow.lite.micro.compression import spec -class PruningCompressor: +class PruningCompressor(compressor.Compressor): """Pruning compression plugin (stub). This stub exists to validate the plugin architecture. The actual pruning