litert-tunner 0.1.1__tar.gz

litert_tunner-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Krzysztof Kolasinski
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

litert_tunner-0.1.1/PKG-INFO

@@ -0,0 +1,149 @@
+Metadata-Version: 2.4
+Name: litert_tunner
+Version: 0.1.1
+Summary: LiteRT tuner Python library.
+Author: litert-tunner contributors
+License: MIT License
+        
+        Copyright (c) 2026 Krzysztof Kolasinski
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: keras>=3.0
+Requires-Dist: numpy
+Requires-Dist: tflite>=2.18.0
+Requires-Dist: flatbuffers
+Provides-Extra: dev
+Requires-Dist: tensorflow>=2.18.0; extra == "dev"
+Requires-Dist: ai-edge-litert; extra == "dev"
+Requires-Dist: bump-my-version>=0.26.0; extra == "dev"
+Requires-Dist: pre-commit>=3.7.0; extra == "dev"
+Requires-Dist: pytest>=8.2.0; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: pytest-sugar; extra == "dev"
+Requires-Dist: pytest-watch; extra == "dev"
+Requires-Dist: pytest-forked; extra == "dev"
+Requires-Dist: pytest-randomly; extra == "dev"
+Requires-Dist: pytest-xdist; extra == "dev"
+Requires-Dist: pytest-dotenv; extra == "dev"
+Requires-Dist: types-tensorflow; extra == "dev"
+Requires-Dist: ruff>=0.5.0; extra == "dev"
+Requires-Dist: ty>=0.0.1a8; extra == "dev"
+Dynamic: license-file
+
+# litert-tunner
+
+Fine-tune fully-quantized **INT8 LiteRT** (TFLite) models *after* export — without retraining from scratch.
+
+## What is this?
+
+**litert-tunner** parses an exported INT8 `.tflite` model, reconstructs it as a Keras 3 model with simulated quantization, lets you fine-tune parameters (biases, scales, weights), and writes them back into the flatbuffer. Graph topology stays intact.
+
+> **Note:** Even though this library simulates the quantization process during fine-tuning, in some cases the resulting INT8 model may still perform worse than a full-precision float32 model.
+
+## Installation
+
+### From source (recommended for development)
+
+```bash
+# Clone the repository
+git clone https://github.com/kmkolasinski/litert-tunner.git
+cd litert-tunner
+
+# Create a virtual environment and install uv
+make venv
+make init
+
+# Install the project in editable mode with dev dependencies
+make install
+```
+
+### As a dependency (pip)
+
+```bash
+pip install litert-tunner
+```
+
+*(Core dependencies: `keras>=3.0`, `numpy`, `tflite>=2.18.0`, `flatbuffers`. `tensorflow` and `ai-edge-litert` are only needed for dev/tests).*
+
+## Quickstart
+
+For a complete, runnable end-to-end example, see the **[Example Notebook](notebooks/quickstart_finetuning.ipynb)**.
+
+```python
+import litert_tunner
+
+# 1. Load an INT8 LiteRT model → trainable Keras replica
+model = litert_tunner.load_model("model_int8.tflite")
+
+# 2. Inference — should match LiteRT Interpreter output
+predictions = model.predict(inputs)
+
+# 3. Prepare for fine-tuning (freeze everything except biases & scales)
+litert_tunner.prepare_for_finetuning(model, trainable_pattern=".*bias")
+
+# 4. Fine-tune with any Keras optimizer / loss
+model.compile(optimizer="adam", loss="mse")
+model.fit(x_train, y_train, epochs=5)
+
+# 5. Export — writes updated parameters back into the flatbuffer
+litert_tunner.save_model(model, "model_int8_finetuned.tflite")
+```
+
+### Distillation Fine-Tuning (Recommended)
+
+```python
+import litert_tunner
+
+# 1. Load an INT8 LiteRT model → trainable Keras replica
+tunner_model = litert_tunner.load_model("model_int8.tflite")
+
+# 2. Freeze everything except biases
+litert_tunner.prepare_for_finetuning(tunner_model, trainable_pattern=".*bias")
+
+# 3. Fine-tune using Trainer (handles distillation & weight drift)
+trainer = litert_tunner.Trainer(
+    student_model=tunner_model,
+    teacher_model=teacher_model,  # Original float32 model
+)
+trainer.compile(optimizer="adam", loss="mse")
+trainer.fit(train_ds, validation_data=val_ds, epochs=5)
+
+# 4. Save updated parameters to flatbuffer
+litert_tunner.save_model(tunner_model, "model_int8_finetuned.tflite")
+```
+
+## Supported Operations
+
+- **Linear:** `FULLY_CONNECTED`, `CONV_2D`, `DEPTHWISE_CONV_2D`
+- **Arithmetic:** `ADD`, `SUB`, `MUL`, `DIV`, `SQUARED_DIFFERENCE`, `NEG`
+- **Activation:** `RELU`, `GELU`, `LOGISTIC`, `SOFTMAX`
+- **Pooling:** `AVERAGE_POOL_2D`, `MAX_POOL_2D`, `MEAN`
+- **Reshape/Resize:** `RESHAPE`, `TRANSPOSE`, `PACK`, `STRIDED_SLICE`, `RESIZE_NEAREST_NEIGHBOR`, `SHAPE`
+- **Other:** `CONCATENATION`, `RSQRT`, `QUANTIZE`, `DEQUANTIZE`
+
+Fused activations (`RELU`, `RELU6`, `RELU_N1_TO_1`) are supported.
+
+## License
+
+[MIT](LICENSE) — Copyright © 2026 Krzysztof Kolasinski

litert_tunner-0.1.1/README.md

@@ -0,0 +1,96 @@
+# litert-tunner
+
+Fine-tune fully-quantized **INT8 LiteRT** (TFLite) models *after* export — without retraining from scratch.
+
+## What is this?
+
+**litert-tunner** parses an exported INT8 `.tflite` model, reconstructs it as a Keras 3 model with simulated quantization, lets you fine-tune parameters (biases, scales, weights), and writes them back into the flatbuffer. Graph topology stays intact.
+
+> **Note:** Even though this library simulates the quantization process during fine-tuning, in some cases the resulting INT8 model may still perform worse than a full-precision float32 model.
+
+## Installation
+
+### From source (recommended for development)
+
+```bash
+# Clone the repository
+git clone https://github.com/kmkolasinski/litert-tunner.git
+cd litert-tunner
+
+# Create a virtual environment and install uv
+make venv
+make init
+
+# Install the project in editable mode with dev dependencies
+make install
+```
+
+### As a dependency (pip)
+
+```bash
+pip install litert-tunner
+```
+
+*(Core dependencies: `keras>=3.0`, `numpy`, `tflite>=2.18.0`, `flatbuffers`. `tensorflow` and `ai-edge-litert` are only needed for dev/tests).*
+
+## Quickstart
+
+For a complete, runnable end-to-end example, see the **[Example Notebook](notebooks/quickstart_finetuning.ipynb)**.
+
+```python
+import litert_tunner
+
+# 1. Load an INT8 LiteRT model → trainable Keras replica
+model = litert_tunner.load_model("model_int8.tflite")
+
+# 2. Inference — should match LiteRT Interpreter output
+predictions = model.predict(inputs)
+
+# 3. Prepare for fine-tuning (freeze everything except biases & scales)
+litert_tunner.prepare_for_finetuning(model, trainable_pattern=".*bias")
+
+# 4. Fine-tune with any Keras optimizer / loss
+model.compile(optimizer="adam", loss="mse")
+model.fit(x_train, y_train, epochs=5)
+
+# 5. Export — writes updated parameters back into the flatbuffer
+litert_tunner.save_model(model, "model_int8_finetuned.tflite")
+```
+
+### Distillation Fine-Tuning (Recommended)
+
+```python
+import litert_tunner
+
+# 1. Load an INT8 LiteRT model → trainable Keras replica
+tunner_model = litert_tunner.load_model("model_int8.tflite")
+
+# 2. Freeze everything except biases
+litert_tunner.prepare_for_finetuning(tunner_model, trainable_pattern=".*bias")
+
+# 3. Fine-tune using Trainer (handles distillation & weight drift)
+trainer = litert_tunner.Trainer(
+    student_model=tunner_model,
+    teacher_model=teacher_model,  # Original float32 model
+)
+trainer.compile(optimizer="adam", loss="mse")
+trainer.fit(train_ds, validation_data=val_ds, epochs=5)
+
+# 4. Save updated parameters to flatbuffer
+litert_tunner.save_model(tunner_model, "model_int8_finetuned.tflite")
+```
+
+## Supported Operations
+
+- **Linear:** `FULLY_CONNECTED`, `CONV_2D`, `DEPTHWISE_CONV_2D`
+- **Arithmetic:** `ADD`, `SUB`, `MUL`, `DIV`, `SQUARED_DIFFERENCE`, `NEG`
+- **Activation:** `RELU`, `GELU`, `LOGISTIC`, `SOFTMAX`
+- **Pooling:** `AVERAGE_POOL_2D`, `MAX_POOL_2D`, `MEAN`
+- **Reshape/Resize:** `RESHAPE`, `TRANSPOSE`, `PACK`, `STRIDED_SLICE`, `RESIZE_NEAREST_NEIGHBOR`, `SHAPE`
+- **Other:** `CONCATENATION`, `RSQRT`, `QUANTIZE`, `DEQUANTIZE`
+
+Fused activations (`RELU`, `RELU6`, `RELU_N1_TO_1`) are supported.
+
+## License
+
+[MIT](LICENSE) — Copyright © 2026 Krzysztof Kolasinski

litert_tunner-0.1.1/pyproject.toml

@@ -0,0 +1,181 @@
+[build-system]
+requires = ["setuptools>=69", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "litert_tunner"
+version = "0.1.1"
+description = "LiteRT tuner Python library."
+readme = "README.md"
+requires-python = ">=3.11"
+license = { file = "LICENSE" }
+authors = [{ name = "litert-tunner contributors" }]
+dependencies = [
+  "keras>=3.0",
+  "numpy",
+  "tflite>=2.18.0",
+  "flatbuffers",
+]
+
+[project.optional-dependencies]
+dev = [
+  "tensorflow>=2.18.0",
+  "ai-edge-litert",
+  "bump-my-version>=0.26.0",
+  "pre-commit>=3.7.0",
+  "pytest>=8.2.0",
+  "pytest-cov",
+  "pytest-mock",
+  "pytest-sugar",
+  "pytest-watch",
+  "pytest-forked",
+  "pytest-randomly",
+  "pytest-xdist",
+  "pytest-dotenv",
+  "types-tensorflow",
+  "ruff>=0.5.0",
+  "ty>=0.0.1a8",
+]
+
+[tool.setuptools]
+package-dir = { "" = "src" }
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.pytest.ini_options]
+testpaths = [
+    "tests",
+]
+
+# Load environment variables from pytest.env file
+env_files = [
+    "pytest.env",
+]
+
+addopts = [
+    "-vv",
+    "--junit-xml=test-results/junit.xml",
+    "--randomly-seed=32",
+]
+
+[tool.coverage.run]
+parallel = true
+patch = ["_exit"]
+data_file = "test-results/.coverage"
+
+[tool.ruff]
+target-version = "py311"
+line-length = 100
+indent-width = 4
+exclude = [
+    "metadata_pb2.py",
+]
+
+# Apply fixes automatically
+fix = true
+
+[tool.ruff.format]
+preview = true
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# By default, Ruff enables Flake8's F rules, along with a subset of the E rules,
+# omitting any stylistic rules that overlap with the use of a formatter, like ruff format or Black.
+# For more details see https://docs.astral.sh/ruff/rules/
+[tool.ruff.lint]
+exclude = ["*.ipynb"]
+select = ["ALL"]
+
+ignore = [
+    # Mostly Missing docstring errors - most of them are False Positives
+    "D100", "D101", "D102", "D103", "D104", "D105", "D107", "D205",
+    "ANN204", # Missing type annotation for self in method
+    "ANN201", # Missing return type annotation for public function
+    "ANN002", "ANN003", # Missing type annotation for `*args` or `**kwargs`
+    "COM812", # Conflicts with another rule
+    "TD002", "TD003", "FIX002", # Todos related formatting rules
+    "TRY003", #  Avoid specifying long messages outside the exception class
+    "EM101", # Exception must not use a string literal, assign to variable first
+    "G004", # Logging uses f-string formatting
+    "PLR0913", # Too many arguments in function - some functions are complex by nature
+]
+
+[tool.ruff.lint.per-file-ignores]
+"src/litert_tunner/testing_utils.py" = [
+    "S101", # Use of assert detected
+]
+"tests/*" = [
+    "SLF001", #  Private member accessed
+    "PERF", # Performance checks are disabled for tests
+    "PLR2004", #  Magic value used in comparison
+    "ARG001", # Unused argument
+    "PT019",  # Fixture `_` without value is injected as parameter
+    "S101", # Use of assert detected
+    "ANN", # Missing type annotation in tests is allowed
+    "ISC001", # single-line-implicit-string-concatenation - conflicts with another rule
+    "PT011", # ValueError is too broad
+    "FBT", # Boolean positional arguments are allowed in tests
+    "C901", # Complexity checks are disabled for test code
+    "PLR0912", # Branch complexity checks are disabled for test code
+    "PLR0915", # Statement complexity checks are disabled for test code
+    "S108", # Insecure temp files are allowed in test code
+]
+
+# Only allow autofix for these import-related and docstring rules
+fixable = ["F401", "F403", "I001", "D411", "D415", "D208", "D209"]
+
+[tool.ruff.lint.pydocstyle]
+convention = "google"
+
+[tool.bumpversion]
+current_version = "0.1.1"
+commit = true
+tag = true
+tag_name = "v{new_version}"
+allow_dirty = false
+
+[[tool.bumpversion.files]]
+filename = "pyproject.toml"
+search = 'version = "{current_version}"'
+replace = 'version = "{new_version}"'
+
+[[tool.bumpversion.files]]
+filename = "src/litert_tunner/__init__.py"
+search = '__version__ = "{current_version}"'
+replace = '__version__ = "{new_version}"'
+
+[tool.interrogate]
+ignore-init-method = true
+ignore-init-module = true
+ignore-magic = false
+ignore-semiprivate = false
+ignore-private = false
+ignore-property-decorators = false
+ignore-module = true
+ignore-nested-functions = false
+ignore-nested-classes = true
+ignore-overloaded-functions = true
+ignore-setters = false
+fail-under = 80
+exclude = ["setup.py", "docs", "build"]
+verbose = 1
+quiet = false
+whitelist-regex = []
+color = true
+omit-covered-files = false
+
+[tool.pyright]
+include = ["src", "tests"]
+extraPaths = ["."]
+venvPath = "."
+venv = ".venv"

litert_tunner-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

litert_tunner-0.1.1/src/litert_tunner/__init__.py

@@ -0,0 +1,46 @@
+"""litert_tunner package."""
+
+import keras
+
+from litert_tunner import flatbuffer, graph
+from litert_tunner.testing_utils import (
+    assert_allclose_with_mismatch_tolerance,
+    assert_cosine_similarity,
+)
+from litert_tunner.trainer import Trainer, cosine_similarity, prepare_for_finetuning
+
+__version__ = "0.1.1"
+
+
+def load_model(path: str) -> keras.Model:
+    """Load a .tflite INT8 model and return a trainable Keras replica.
+
+    Args:
+        path: Path to the .tflite file.
+
+    Returns:
+        A trainable Keras Model replica.
+    """
+    graph_def = flatbuffer.parse_tflite(path)
+    return graph.build_keras_model(graph_def)
+
+
+def save_model(model: keras.Model, path: str) -> None:
+    """Save updated parameters back to a .tflite file.
+
+    Args:
+        model: The trained Keras Model replica.
+        path: Path to write the updated .tflite file.
+    """
+    flatbuffer.save_tflite(model, path)
+
+
+__all__ = [
+    "Trainer",
+    "assert_allclose_with_mismatch_tolerance",
+    "assert_cosine_similarity",
+    "cosine_similarity",
+    "load_model",
+    "prepare_for_finetuning",
+    "save_model",
+]

litert_tunner-0.1.1/src/litert_tunner/export.py

@@ -0,0 +1,100 @@
+from collections.abc import Callable, Iterable
+from pathlib import Path
+from typing import Literal, TypeAlias
+
+import keras
+import numpy as np
+import tensorflow as tf
+
+RepresentativeDataset: TypeAlias = Callable[
+    [], tf.data.Dataset | Iterable[list[np.ndarray] | dict[str, np.ndarray]]
+]
+
+
+def export_litert_model(
+    model: keras.Model,
+    export_dir: str | Path,
+    *,
+    quantization: Literal["int8", "float32"] = "int8",
+    float_io: bool = True,
+    representative_dataset: RepresentativeDataset | None = None,
+    run_debugger: bool = False,
+    denylisted_ops: list[str] | None = None,
+    denylisted_nodes: list[str] | None = None,
+) -> tuple[Path, "tf.lite.experimental.QuantizationDebugger | None"]:
+    """Export a Keras model to a TFLite model.
+
+    Args:
+        model: The Keras model to export.
+        export_dir: Directory where the exported model and stats will be saved.
+        quantization: The quantization mode. Either "int8" or "float32". Defaults to "int8".
+        float_io: If True, use float32 inputs/outputs even for int8 quantization. Defaults to True.
+        representative_dataset: A generator yielding sample inputs for quantization.
+            Required if quantization is "int8".
+        run_debugger: If True, run the TFLite QuantizationDebugger and save stats.
+            Only supported for int8 quantization. Defaults to False.
+        denylisted_ops: List of ops to exclude from quantization.
+        denylisted_nodes: List of nodes to exclude from quantization.
+
+    Returns:
+        A tuple containing the path to the exported .tflite model and the
+        QuantizationDebugger instance (if run_debugger is True, otherwise None).
+    """
+    if (denylisted_ops or denylisted_nodes) and not run_debugger:
+        raise ValueError(
+            "denylisted_ops and denylisted_nodes can only be provided when run_debugger is True"
+        )
+    if quantization == "int8" and representative_dataset is None:
+        raise ValueError("representative_dataset must be provided for int8 quantization")
+    if quantization == "float32" and run_debugger:
+        raise ValueError("Debugger only supported for int8 quantization")
+    if quantization not in ("int8", "float32"):
+        msg = f"Unknown quantization: {quantization}"
+        raise ValueError(msg)
+
+    converter = tf.lite.TFLiteConverter.from_keras_model(model)
+
+    match quantization:
+        case "int8":
+            converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            converter.representative_dataset = representative_dataset
+            if denylisted_ops or denylisted_nodes:
+                converter.target_spec.supported_ops = [
+                    tf.lite.OpsSet.TFLITE_BUILTINS_INT8,
+                    tf.lite.OpsSet.TFLITE_BUILTINS,
+                ]
+            else:
+                converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
+        case "float32":
+            pass
+
+    io_type = tf.float32 if float_io or quantization == "float32" else tf.int8
+    converter.inference_input_type = io_type
+    converter.inference_output_type = io_type
+
+    export_dir = Path(export_dir)
+    export_dir.mkdir(parents=True, exist_ok=True)
+
+    debugger = None
+    if run_debugger:
+        debug_options = tf.lite.experimental.QuantizationDebugOptions(
+            denylisted_ops=denylisted_ops,
+            denylisted_nodes=denylisted_nodes,
+        )
+        debugger = tf.lite.experimental.QuantizationDebugger(
+            converter=converter,
+            debug_dataset=representative_dataset,
+            debug_options=debug_options,
+        )
+        debugger.run()
+        tflite_model = debugger.get_nondebug_quantized_model()
+
+        with (export_dir / "quantization_stats.csv").open("w") as f:
+            debugger.layer_statistics_dump(f)
+    else:
+        tflite_model = converter.convert()
+
+    tflite_model_filepath = export_dir / "model.tflite"
+    tflite_model_filepath.write_bytes(tflite_model)
+
+    return tflite_model_filepath, debugger

litert_tunner-0.1.1/src/litert_tunner/flatbuffer/__init__.py

@@ -0,0 +1,9 @@
+"""Flatbuffer module for litert_tunner.
+
+Exposes parse_tflite and save_tflite.
+"""
+
+from __future__ import annotations
+
+from litert_tunner.flatbuffer.parser import parse_tflite as parse_tflite
+from litert_tunner.flatbuffer.writer import save_tflite as save_tflite