Skip to content

Export

TFLite export and C header generation for edge deployment.

compressionkit.export.tflite.export_encoder_tflite(encoder, *, rep_dataset, output_dir, tflite_name='encoder.tflite', header_name='encoder.h', c_array_name='encoder', quantization='INT8', io_type='int8')

Export encoder to INT8 TFLite and C header using helia_edge.

Parameters:

Name Type Description Default
encoder Model

Keras encoder model to export.

required
rep_dataset ndarray

Representative dataset array for calibration.

required
output_dir Path

Directory to write .tflite and .h files.

required
tflite_name str

Filename for the TFLite model.

'encoder.tflite'
header_name str

Filename for the C header.

'encoder.h'
c_array_name str

Name for the C array in the header.

'encoder'
quantization str

Quantization mode (e.g. "INT8").

'INT8'
io_type str

I/O type string (e.g. "int8").

'int8'

Returns:

Type Description
tuple[Path, Path]

Tuple of (tflite_path, header_path).

Source code in compressionkit/export/tflite.py
def export_encoder_tflite(
    encoder: keras.Model,
    *,
    rep_dataset: np.ndarray,
    output_dir: Path,
    tflite_name: str = "encoder.tflite",
    header_name: str = "encoder.h",
    c_array_name: str = "encoder",
    quantization: str = "INT8",
    io_type: str = "int8",
) -> tuple[Path, Path]:
    """Export encoder to INT8 TFLite and C header using helia_edge.

    Args:
        encoder: Keras encoder model to export.
        rep_dataset: Representative dataset array for calibration.
        output_dir: Directory to write ``.tflite`` and ``.h`` files.
        tflite_name: Filename for the TFLite model.
        header_name: Filename for the C header.
        c_array_name: Name for the C array in the header.
        quantization: Quantization mode (e.g. ``"INT8"``).
        io_type: I/O type string (e.g. ``"int8"``).

    Returns:
        Tuple of ``(tflite_path, header_path)``.
    """
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    converter = helia.converters.tflite.TfLiteKerasConverter(model=encoder)
    with open(os.devnull, "w") as devnull, contextlib.redirect_stdout(devnull), contextlib.redirect_stderr(devnull):
        converter.convert(
            test_x=rep_dataset,
            quantization=quantization,
            io_type=io_type,
            mode="KERAS",
            strict=False,
            verbose=0,
        )
    tflite_path = output_dir / tflite_name
    header_path = output_dir / header_name
    converter.export(tflite_path=tflite_path)
    converter.export_header(header_path=header_path, name=c_array_name)
    logger.info("Exported encoder TFLite to %s and header to %s", tflite_path, header_path)
    return tflite_path, header_path