From 38b5eed2498506ca0041df89cf29780ff05196b3 Mon Sep 17 00:00:00 2001 From: Xuban <59646791+EHxuban11@users.noreply.github.com> Date: Thu, 20 Aug 2026 23:51:49 +0200 Subject: [PATCH] Add TFLite INT8 export export(format="tflite", int8=True, data=...) now runs onnx2tf post-training quantization instead of rejecting the request, and returns the fully integer artifact. Calibration batches are the model's own preprocessed tensors, so onnx2tf is told not to normalize them again. The integer path runs on the tf_converter backend: flatbuffer_direct aborts on any op it cannot keep in int8 end to end, which YOLO9 hits on EQUAL. FP32 export keeps flatbuffer_direct unchanged. --- CHANGELOG.md | 10 ++ libreyolo/export/exporter.py | 15 +- libreyolo/export/tflite.py | 156 +++++++++++++++++- tests/unit/test_export_tflite.py | 271 ++++++++++++++++++++++++++++++- 4 files changed, 438 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index aa483f2f2..b957b568a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,16 @@ before 1.4.0 are documented in the ### Added +- **TFLite INT8 export.** `export(format="tflite", int8=True, data=...)` now + runs post-training quantization instead of rejecting the request, and + returns the fully integer artifact (int8 input and output tensors, which is + what EdgeTPU and int8-only MCU runtimes require). `TFLiteBackend` already + quantized inputs and dequantized outputs from the tensor scales, so + `predict()` on the artifact takes float images as before. Calibration data + is mandatory: there is no eight-image default for this format. Accuracy is + not parity-validated per family yet, so measure with `val()` before + deploying. + - **LibreGround** sibling factory: screenshot + instruction → `Results.points`. Shipped adapters are Florence-2-base (MIT), ShowUI-2B (MIT weights; Apache-2.0 code/base), and Qwen3-VL-2B (Apache-2.0). diff --git a/libreyolo/export/exporter.py b/libreyolo/export/exporter.py index deac189c6..36d35ce69 100644 --- a/libreyolo/export/exporter.py +++ b/libreyolo/export/exporter.py @@ -2306,9 +2306,13 @@ class TFLiteExporter(BaseExporter): format_name = "tflite" suffix = ".tflite" requires_onnx = True - supports_int8 = False + supports_int8 = True supports_fp16 = False apply_model_half = False + # Deliberately no default calibration set, matching TensorRT and OpenVINO. + # A full-integer TFLite graph calibrated on eight images is quietly wrong + # rather than loudly missing, so data= stays mandatory. + default_int8_calibration_data = False def __call__(self, *args, dynamic: bool = False, **kwargs) -> str: if dynamic: @@ -2326,11 +2330,6 @@ def _validate(self, half: bool, int8: bool, data: Optional[str]): raise ValueError( "TFLite FP16 export is not supported yet. Omit half=True for FP32." ) - if int8: - raise ValueError( - "TFLite INT8 quantization is not supported yet. " - "Omit int8=True for FP32." - ) return super()._validate(half, int8, data) def _preflight(self, **kwargs): @@ -2348,6 +2347,8 @@ def _export( metadata, onnx_path, half, + int8, + calibration_data, verbose, onnx2tf_args=None, **kwargs, @@ -2364,6 +2365,8 @@ def _export( onnx_path=onnx_path, output_path=output_path, half=half, + int8=int8, + calibration_data=calibration_data, verbose=verbose, onnx2tf_args=onnx2tf_args, metadata=metadata, diff --git a/libreyolo/export/tflite.py b/libreyolo/export/tflite.py index b6f3ef3da..29107b460 100644 --- a/libreyolo/export/tflite.py +++ b/libreyolo/export/tflite.py @@ -674,7 +674,31 @@ def _onnx2tf_command() -> list[str]: return [sys.executable, "-m", "onnx2tf"] -def _find_converted_tflite(output_dir: Path, onnx_path: Path) -> Path: +# onnx2tf emits every quantized variant it can build under -oiqt. Prefer the +# fully integer artifact: it is the only one EdgeTPU and int8-only MCU runtimes +# accept, and TFLiteBackend already quantizes the input and dequantizes the +# outputs from the tensor scales, so float32 callers see no difference. +_INT8_ARTIFACT_SUFFIXES = ("_full_integer_quant", "_integer_quant") + + +def _find_converted_tflite( + output_dir: Path, onnx_path: Path, *, int8: bool = False +) -> Path: + if int8: + for suffix in _INT8_ARTIFACT_SUFFIXES: + exact = output_dir / f"{onnx_path.stem}{suffix}.tflite" + if exact.exists(): + return exact + matches = sorted(output_dir.rglob(f"*{suffix}.tflite")) + if matches: + return matches[0] + # Never fall back to a float artifact here: returning one would hand + # back an FP32 model under int8 filename and precision metadata. + produced = sorted(str(f.relative_to(output_dir)) for f in output_dir.rglob("*")) + raise RuntimeError( + f"onnx2tf did not produce an INT8 TFLite file. Files found: {produced[:20]}" + ) + exact = output_dir / f"{onnx_path.stem}_float32.tflite" if exact.exists(): return exact @@ -693,6 +717,73 @@ def _find_converted_tflite(output_dir: Path, onnx_path: Path) -> Path: ) +def _onnx_input_name(onnx_path: Path) -> str: + """Return the name of the first ONNX graph input.""" + import onnx + + model = onnx.load(str(onnx_path), load_external_data=False) + if not model.graph.input: + raise ValueError(f"ONNX graph has no inputs: {onnx_path}") + return model.graph.input[0].name + + +def _write_int8_calibration_npy(calibration_data: Any, output_dir: Path) -> Path: + """Write the calibration batches as one NHWC float32 ``.npy`` for onnx2tf. + + ``CalibrationDataLoader`` yields already-preprocessed NCHW batches and pads + the final batch by repeating its last image so TensorRT sees a full batch. + onnx2tf instead reads the whole file as a representative dataset, so the + padding is trimmed here and every image is weighted once. The array is + built through a memmap because a few hundred 640x640 float32 images are + larger than the process should hold at once. + """ + npy_path = output_dir / "_int8_calibration.npy" + total = int(calibration_data.num_samples) + array = None + written = 0 + + for batch in calibration_data: + if written >= total: + break + if batch.ndim != 4: + raise ValueError( + "TFLite INT8 calibration requires rank-4 NCHW batches; " + f"got shape {batch.shape}." + ) + chunk = np.ascontiguousarray(np.transpose(batch, (0, 2, 3, 1))) + chunk = chunk[: total - written] + if array is None: + array = np.lib.format.open_memmap( + npy_path, + mode="w+", + dtype=np.float32, + shape=(total, *chunk.shape[1:]), + ) + array[written : written + len(chunk)] = chunk + written += len(chunk) + + if array is None or written == 0: + raise RuntimeError( + "The calibration dataset produced no usable images for TFLite " + "INT8 export. Check that the images in data= can be read." + ) + array.flush() + del array + + if written < total: + # Unreadable images are skipped by the loader, so the memmap can be + # longer than what was filled. Copy the used prefix into a second file + # rather than truncating one that may still be mapped on Windows. + trimmed_path = output_dir / "_int8_calibration_trimmed.npy" + source = np.load(npy_path, mmap_mode="r") + np.save(trimmed_path, np.asarray(source[:written])) + del source + npy_path = trimmed_path + + logger.info("INT8 calibration tensor: %d images -> %s", written, npy_path) + return npy_path + + def _write_metadata_sidecar(output_path: Path, metadata: dict) -> None: sidecar_path = Path(str(output_path) + ".json") with open(sidecar_path, "w") as f: @@ -705,12 +796,19 @@ def export_tflite( output_path: str, *, half: bool = False, + int8: bool = False, + calibration_data: Any = None, verbose: bool = False, onnx2tf_args: Iterable[str] | None = None, metadata: dict | None = None, ) -> str: """Convert a static ONNX model to TensorFlow Lite using onnx2tf. + With ``int8=True`` the converter runs post-training quantization against + ``calibration_data`` and the fully integer artifact is returned. The + calibration batches are the model's own preprocessed tensors, so onnx2tf is + told not to normalize them a second time (mean 0, std 1). + Note: ``onnx2tf_args`` is forwarded only on the YOLO9 CLI path. It is not applicable to the RF-DETR Python-API path and will be ignored there. """ @@ -718,11 +816,23 @@ def export_tflite( raise ValueError( "TFLite FP16 export is not supported yet. Omit half=True for FP32." ) + if int8 and calibration_data is None: + raise ValueError( + "TFLite INT8 export requires calibration data. " + "Pass data= to export()." + ) check_tflite_export_available() model_family = ((metadata or {}).get("model_family") or "").lower() if model_family == "rfdetr": + if int8: + raise NotImplementedError( + "TFLite INT8 export is not implemented for RF-DETR. That family " + "converts through the onnx2tf Python API with a bespoke " + "GridSample and position-embedding fixup that the quantization " + "path has not been run against." + ) if onnx2tf_args is not None: logger.warning( "onnx2tf_args is not supported on the RF-DETR TFLite path " @@ -747,10 +857,40 @@ def export_tflite( "-o", str(tmp_output), "-tb", - "flatbuffer_direct", + # flatbuffer_direct is the validated FP32 lowering, but its integer + # path is strict: it aborts on any op it cannot represent in int8 + # end to end, which YOLO9 hits on EQUAL. tf_converter runs the + # standard TFLite representative-dataset quantization instead and + # leaves such ops in float. + "tf_converter" if int8 else "flatbuffer_direct", "-v", "info" if verbose else "warn", ] + if int8: + calib_npy = _write_int8_calibration_npy(calibration_data, tmp_output) + # mean 0 / std 1: onnx2tf applies (value - mean) / std to the + # calibration tensor, and these batches already went through the + # model's own preprocessing. + cmd += [ + "-oiqt", + "-cind", + _onnx_input_name(onnx_file), + str(calib_npy), + "0.0", + "1.0", + ] + logger.warning( + "TFLite INT8 is post-training quantization; accuracy is not " + "parity-validated per family. Measure the exported artifact " + "with val() before deploying it." + ) + if shutil.which("onnxsim") is None: + logger.warning( + "The onnxsim executable is not on PATH. onnx2tf shells out " + "to it by name, and the tf_converter backend needs the " + "shapes it propagates. Install onnx-simplifier into the " + "active environment if the conversion below fails." + ) if onnx2tf_args is not None: cmd.extend(str(arg) for arg in onnx2tf_args) @@ -762,14 +902,22 @@ def export_tflite( if result.returncode != 0: stdout = result.stdout or "" stderr = result.stderr or "" + hint = "" + if int8 and shutil.which("onnxsim") is None: + hint = ( + "\nHint: onnxsim was not found on PATH. onnx2tf invokes it " + "as a bare command, and without the shapes it propagates " + "the INT8 backend fails while building the Keras graph." + ) raise RuntimeError( - f"onnx2tf failed with exit code {result.returncode}.\n" + f"onnx2tf failed with exit code {result.returncode}.{hint}\n" f"Command: {' '.join(cmd)}\n" f"stdout: {stdout}\n" f"stderr: {stderr}" ) - converted = _find_converted_tflite(tmp_output, onnx_file) + converted = _find_converted_tflite(tmp_output, onnx_file, int8=int8) + logger.info("Selected converted artifact: %s", converted.name) shutil.copy2(converted, dst) if metadata is not None: diff --git a/tests/unit/test_export_tflite.py b/tests/unit/test_export_tflite.py index 8ade158b1..c06129e26 100644 --- a/tests/unit/test_export_tflite.py +++ b/tests/unit/test_export_tflite.py @@ -9,6 +9,7 @@ from pathlib import Path from unittest.mock import MagicMock +import cv2 import numpy as np import pytest import torch @@ -67,9 +68,10 @@ def test_tflite_format_registered(): assert "tflite" in BaseExporter._registry assert TFLiteExporter.suffix == ".tflite" assert TFLiteExporter.requires_onnx is True - assert TFLiteExporter.supports_int8 is False + assert TFLiteExporter.supports_int8 is True assert TFLiteExporter.supports_fp16 is False assert TFLiteExporter.apply_model_half is False + assert TFLiteExporter.default_int8_calibration_data is False def test_tflite_family_support_scaffold(): @@ -103,11 +105,11 @@ def test_tflite_rejects_dynamic_export(): exporter(dynamic=True) -def test_tflite_rejects_int8_export(): +def test_tflite_int8_requires_calibration_data(): exporter = TFLiteExporter(_make_wrapper()) - with pytest.raises(ValueError, match="INT8"): - exporter(output_path="unused.tflite", int8=True, data="coco8") + with pytest.raises(ValueError, match="requires calibration data"): + exporter(output_path="unused.tflite", int8=True) def test_tflite_rejects_fp16_export(): @@ -185,7 +187,10 @@ def test_tflite_export_copies_float32_output(monkeypatch, tmp_path): monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) monkeypatch.setattr(tflite_module, "_onnx2tf_command", lambda: ["onnx2tf"]) + captured = {} + def fake_run(cmd, capture_output, text): + captured["cmd"] = list(cmd) output_dir = Path(cmd[cmd.index("-o") + 1]) output_dir.mkdir(parents=True, exist_ok=True) (output_dir / "model_float32.tflite").write_bytes(b"fp32") @@ -203,6 +208,8 @@ def fake_run(cmd, capture_output, text): assert result == str(fp32_dst) assert fp32_dst.read_bytes() == b"fp32" + assert captured["cmd"][captured["cmd"].index("-tb") + 1] == "flatbuffer_direct" + assert "-oiqt" not in captured["cmd"] sidecar = json.loads(Path(str(fp32_dst) + ".json").read_text()) assert sidecar["model_family"] == "yolo9" @@ -302,6 +309,262 @@ def fake_export_tflite(**kwargs): assert not Path(captured["tflite"]["onnx_path"]).exists() +class _FakeCalibration: + """Stand-in for CalibrationDataLoader: NCHW batches, padded final batch.""" + + def __init__(self, batches, num_samples): + self._batches = batches + self.num_samples = num_samples + + def __iter__(self): + return iter(self._batches) + + +def _constant_images(*values): + return np.concatenate( + [np.full((1, 3, 2, 2), value, dtype=np.float32) for value in values] + ) + + +def _int8_converter(monkeypatch, produced): + """Patch onnx2tf away and record the command it would have been given.""" + from libreyolo.export import tflite as tflite_module + + monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) + monkeypatch.setattr(tflite_module, "_onnx2tf_command", lambda: ["onnx2tf"]) + monkeypatch.setattr(tflite_module, "_onnx_input_name", lambda _path: "images") + captured = {} + + def fake_run(cmd, capture_output, text): + captured["cmd"] = list(cmd) + output_dir = Path(cmd[cmd.index("-o") + 1]) + output_dir.mkdir(parents=True, exist_ok=True) + for name, payload in produced.items(): + (output_dir / name).write_bytes(payload) + return subprocess.CompletedProcess(cmd, 0, stdout="", stderr="") + + monkeypatch.setattr(subprocess, "run", fake_run) + return captured + + +def test_int8_calibration_npy_is_nhwc_and_drops_batch_padding(tmp_path): + from libreyolo.export import tflite as tflite_module + + # Second batch repeats its last image, the way CalibrationDataLoader pads. + loader = _FakeCalibration( + [_constant_images(1.0, 2.0), _constant_images(3.0, 3.0)], + num_samples=3, + ) + + npy_path = tflite_module._write_int8_calibration_npy(loader, tmp_path) + data = np.load(npy_path) + + assert data.shape == (3, 2, 2, 3) + assert data.dtype == np.float32 + np.testing.assert_array_equal(data[:, 0, 0, 0], [1.0, 2.0, 3.0]) + + +def test_int8_calibration_npy_trims_when_images_are_skipped(tmp_path): + from libreyolo.export import tflite as tflite_module + + # num_samples counts files on disk; unreadable ones never reach a batch. + loader = _FakeCalibration([_constant_images(1.0, 2.0)], num_samples=4) + + data = np.load(tflite_module._write_int8_calibration_npy(loader, tmp_path)) + + assert data.shape == (2, 2, 2, 3) + np.testing.assert_array_equal(data[:, 0, 0, 0], [1.0, 2.0]) + + +def test_int8_calibration_npy_rejects_empty_dataset(tmp_path): + from libreyolo.export import tflite as tflite_module + + with pytest.raises(RuntimeError, match="no usable images"): + tflite_module._write_int8_calibration_npy( + _FakeCalibration([], num_samples=4), tmp_path + ) + + +def test_tflite_int8_helper_requires_calibration_data(tmp_path): + from libreyolo.export import tflite as tflite_module + + with pytest.raises(ValueError, match="requires calibration data"): + tflite_module.export_tflite( + str(tmp_path / "model.onnx"), + str(tmp_path / "model.tflite"), + int8=True, + ) + + +def test_tflite_int8_rejected_for_rfdetr(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + monkeypatch.setattr(tflite_module, "check_tflite_export_available", lambda: None) + + with pytest.raises(NotImplementedError, match="RF-DETR"): + tflite_module.export_tflite( + str(tmp_path / "model.onnx"), + str(tmp_path / "model.tflite"), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], num_samples=1), + metadata={"model_family": "rfdetr"}, + ) + + +def test_tflite_int8_export_selects_full_integer_artifact(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + captured = _int8_converter( + monkeypatch, + { + "model_float32.tflite": b"fp32", + "model_integer_quant.tflite": b"int8", + "model_full_integer_quant.tflite": b"full-int8", + }, + ) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + result = tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0, 2.0)], 2), + metadata={"model_family": "yolo9", "precision": "int8"}, + ) + + assert result == str(dst) + assert dst.read_bytes() == b"full-int8" + + cmd = captured["cmd"] + # The integer path runs on tf_converter; flatbuffer_direct aborts on any op + # it cannot keep in int8 end to end. + assert cmd[cmd.index("-tb") + 1] == "tf_converter" + assert "-oiqt" in cmd + cind = cmd.index("-cind") + assert cmd[cind + 1] == "images" + assert cmd[cind + 2].endswith(".npy") + # mean 0 / std 1: the calibration batches are already preprocessed, so a + # second normalization inside onnx2tf would calibrate the wrong ranges. + assert cmd[cind + 3 : cind + 5] == ["0.0", "1.0"] + + sidecar = json.loads(Path(str(dst) + ".json").read_text()) + assert sidecar["precision"] == "int8" + + +def test_tflite_int8_export_falls_back_to_integer_quant(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + _int8_converter( + monkeypatch, + { + "model_float32.tflite": b"fp32", + "model_integer_quant.tflite": b"int8", + }, + ) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], 1), + ) + + assert dst.read_bytes() == b"int8" + + +def test_tflite_int8_export_never_returns_a_float_artifact(monkeypatch, tmp_path): + from libreyolo.export import tflite as tflite_module + + _int8_converter(monkeypatch, {"model_float32.tflite": b"fp32"}) + onnx_path = tmp_path / "model.onnx" + onnx_path.write_bytes(b"fake onnx") + dst = tmp_path / "model.tflite" + + with pytest.raises(RuntimeError, match="did not produce an INT8"): + tflite_module.export_tflite( + str(onnx_path), + str(dst), + int8=True, + calibration_data=_FakeCalibration([_constant_images(1.0)], 1), + ) + + assert not dst.exists() + + +def test_tflite_exporter_forwards_int8_and_calibration(monkeypatch, tmp_path): + import libreyolo.export.exporter as exporter_module + + _mock_onnx_available(monkeypatch) + + image_dir = tmp_path / "images" + image_dir.mkdir() + rng = np.random.default_rng(0) + for idx in range(3): + cv2.imwrite( + str(image_dir / f"{idx}.jpg"), + rng.integers(0, 256, size=(24, 24, 3), dtype=np.uint8), + ) + data_yaml = tmp_path / "data.yaml" + data_yaml.write_text( + "\n".join( + [ + f"path: {tmp_path.as_posix()}", + "train: images", + "val: images", + "nc: 1", + "names:", + " 0: object", + ] + ), + encoding="utf-8", + ) + + def preprocess(img_rgb, imgsz): + h, w = imgsz if isinstance(imgsz, tuple) else (imgsz, imgsz) + resized = cv2.resize(img_rgb, (w, h)).astype(np.float32) / 255.0 + return np.transpose(resized, (2, 0, 1)), 1.0 + + wrapper = _make_wrapper() + wrapper._get_preprocess_numpy.return_value = preprocess + exporter = TFLiteExporter(wrapper) + output_path = tmp_path / "model.tflite" + captured = {} + + def fake_export_onnx(_nn_model, _dummy, **kwargs): + Path(kwargs["output_path"]).write_bytes(b"onnx") + return kwargs["output_path"] + + def fake_export_tflite(**kwargs): + captured.update(kwargs) + Path(kwargs["output_path"]).write_bytes(b"tflite") + return kwargs["output_path"] + + monkeypatch.setattr(exporter_module, "export_onnx", fake_export_onnx) + monkeypatch.setattr( + "libreyolo.export.tflite.check_tflite_export_available", lambda: None + ) + monkeypatch.setattr("libreyolo.export.tflite.export_tflite", fake_export_tflite) + + exporter( + output_path=str(output_path), + imgsz=16, + simplify=False, + int8=True, + data=str(data_yaml), + ) + + assert captured["int8"] is True + assert captured["half"] is False + assert captured["metadata"]["precision"] == "int8" + assert captured["calibration_data"].num_samples == 3 + assert next(iter(captured["calibration_data"])).shape[1:] == (3, 16, 16) + + def test_tflite_backend_restores_yolonas_channel_first_outputs(): from libreyolo.backends.tflite import TFLiteBackend