Skip to content

runners

birdnet_stm32.models.runners

Inference runners for Keras and TFLite models.

Provides a uniform predict(x_batch) interface for both model formats, enabling the evaluation pipeline to be agnostic to the model type.

KerasRunner

Thin wrapper for a Keras model to standardize batch prediction.

Source code in birdnet_stm32/models/runners.py
class KerasRunner:
    """Thin wrapper for a Keras model to standardize batch prediction."""

    def __init__(self, model: tf.keras.Model):
        """Initialize with a loaded Keras model.

        Args:
            model: Loaded Keras model (compiled=False is fine).
        """
        self.model = model
        # One traced graph per input shape instead of op-by-op eager dispatch,
        # which dominated evaluation time. The graph reads the model's variables,
        # so weights updated after construction are still used.
        self._forward = tf.function(lambda x: self.model(x, training=False), reduce_retracing=True)

    def predict(self, x_batch: np.ndarray) -> np.ndarray:
        """Run a forward pass on a batch.

        Args:
            x_batch: Input batch in the model's expected shape and dtype.

        Returns:
            Model outputs [B, C] as float32.
        """
        x_batch = x_batch.astype(np.float32, copy=False)
        return np.asarray(self._forward(tf.convert_to_tensor(x_batch)), dtype=np.float32)

__init__(model)

Initialize with a loaded Keras model.

Parameters:

Name Type Description Default
model Model

Loaded Keras model (compiled=False is fine).

required
Source code in birdnet_stm32/models/runners.py
def __init__(self, model: tf.keras.Model):
    """Initialize with a loaded Keras model.

    Args:
        model: Loaded Keras model (compiled=False is fine).
    """
    self.model = model
    # One traced graph per input shape instead of op-by-op eager dispatch,
    # which dominated evaluation time. The graph reads the model's variables,
    # so weights updated after construction are still used.
    self._forward = tf.function(lambda x: self.model(x, training=False), reduce_retracing=True)

predict(x_batch)

Run a forward pass on a batch.

Parameters:

Name Type Description Default
x_batch ndarray

Input batch in the model's expected shape and dtype.

required

Returns:

Type Description
ndarray

Model outputs [B, C] as float32.

Source code in birdnet_stm32/models/runners.py
def predict(self, x_batch: np.ndarray) -> np.ndarray:
    """Run a forward pass on a batch.

    Args:
        x_batch: Input batch in the model's expected shape and dtype.

    Returns:
        Model outputs [B, C] as float32.
    """
    x_batch = x_batch.astype(np.float32, copy=False)
    return np.asarray(self._forward(tf.convert_to_tensor(x_batch)), dtype=np.float32)

TFLiteRunner

TFLite model runner (XNNPACK, or the reference kernels where XNNPACK refuses the graph).

Source code in birdnet_stm32/models/runners.py
class TFLiteRunner:
    """TFLite model runner (XNNPACK, or the reference kernels where XNNPACK refuses the graph)."""

    def __init__(self, model_path: str, num_threads: int | None = None):
        """Initialize with a TFLite model file.

        Args:
            model_path: Path to a .tflite model file.
            num_threads: Interpreter threads (default: up to 8).
        """
        threads = num_threads if num_threads is not None else min(8, os.cpu_count() or 1)
        self.interpreter = allocated_interpreter(model_path=model_path, num_threads=threads)
        self.input_index = None
        self.output_index = None
        self._allocate()

    def _allocate(self):
        """Allocate tensors and cache input/output tensor indices."""
        self.interpreter.allocate_tensors()
        in_det = self.interpreter.get_input_details()[0]
        out_det = self.interpreter.get_output_details()[0]
        self.input_index = in_det["index"]
        self.output_index = out_det["index"]

    def _ensure_shape(self, shape: tuple):
        """Resize the interpreter input tensor to match the batch shape if needed.

        Args:
            shape: Desired input tensor shape.
        """
        in_det = self.interpreter.get_input_details()[0]
        cur = in_det["shape"]
        if list(cur) != list(shape):
            self.interpreter.resize_tensor_input(self.input_index, shape)
            self._allocate()

    def predict(self, x_batch: np.ndarray) -> np.ndarray:
        """Run a forward pass on a batch.

        Args:
            x_batch: Input batch in the model's expected shape and dtype.

        Returns:
            Model outputs [B, C] as float32.
        """
        x_batch = x_batch.astype(np.float32, copy=False)
        self._ensure_shape(x_batch.shape)
        self.interpreter.set_tensor(self.input_index, x_batch)
        self.interpreter.invoke()
        return np.asarray(self.interpreter.get_tensor(self.output_index), dtype=np.float32)

__init__(model_path, num_threads=None)

Initialize with a TFLite model file.

Parameters:

Name Type Description Default
model_path str

Path to a .tflite model file.

required
num_threads int | None

Interpreter threads (default: up to 8).

None
Source code in birdnet_stm32/models/runners.py
def __init__(self, model_path: str, num_threads: int | None = None):
    """Initialize with a TFLite model file.

    Args:
        model_path: Path to a .tflite model file.
        num_threads: Interpreter threads (default: up to 8).
    """
    threads = num_threads if num_threads is not None else min(8, os.cpu_count() or 1)
    self.interpreter = allocated_interpreter(model_path=model_path, num_threads=threads)
    self.input_index = None
    self.output_index = None
    self._allocate()

predict(x_batch)

Run a forward pass on a batch.

Parameters:

Name Type Description Default
x_batch ndarray

Input batch in the model's expected shape and dtype.

required

Returns:

Type Description
ndarray

Model outputs [B, C] as float32.

Source code in birdnet_stm32/models/runners.py
def predict(self, x_batch: np.ndarray) -> np.ndarray:
    """Run a forward pass on a batch.

    Args:
        x_batch: Input batch in the model's expected shape and dtype.

    Returns:
        Model outputs [B, C] as float32.
    """
    x_batch = x_batch.astype(np.float32, copy=False)
    self._ensure_shape(x_batch.shape)
    self.interpreter.set_tensor(self.input_index, x_batch)
    self.interpreter.invoke()
    return np.asarray(self.interpreter.get_tensor(self.output_index), dtype=np.float32)

ChainedTFLiteRunner

Run a split backbone and classifier head as one model.

Conversion can emit the classifier head as its own artifact so it can be updated over a narrowband link without reflashing the backbone. Evaluation and deployment then need the two halves to behave like the model they came from, which is what this runner provides.

Source code in birdnet_stm32/models/runners.py
class ChainedTFLiteRunner:
    """Run a split backbone and classifier head as one model.

    Conversion can emit the classifier head as its own artifact so it can be
    updated over a narrowband link without reflashing the backbone. Evaluation
    and deployment then need the two halves to behave like the model they came
    from, which is what this runner provides.
    """

    def __init__(self, backbone_path: str, classifier_path: str):
        """Initialize with the two halves of a split model.

        Args:
            backbone_path: Path to the .tflite backbone (audio -> embeddings).
            classifier_path: Path to the .tflite head (embeddings -> scores).
        """
        self.backbone = TFLiteRunner(backbone_path)
        self.classifier = TFLiteRunner(classifier_path)

    def predict(self, x_batch: np.ndarray) -> np.ndarray:
        """Run the backbone and feed its embeddings to the classifier head.

        Args:
            x_batch: Input batch in the backbone's expected shape and dtype.

        Returns:
            Model outputs [B, C] as float32.
        """
        embeddings = self.backbone.predict(x_batch)
        return self.classifier.predict(np.asarray(embeddings, dtype=np.float32))

__init__(backbone_path, classifier_path)

Initialize with the two halves of a split model.

Parameters:

Name Type Description Default
backbone_path str

Path to the .tflite backbone (audio -> embeddings).

required
classifier_path str

Path to the .tflite head (embeddings -> scores).

required
Source code in birdnet_stm32/models/runners.py
def __init__(self, backbone_path: str, classifier_path: str):
    """Initialize with the two halves of a split model.

    Args:
        backbone_path: Path to the .tflite backbone (audio -> embeddings).
        classifier_path: Path to the .tflite head (embeddings -> scores).
    """
    self.backbone = TFLiteRunner(backbone_path)
    self.classifier = TFLiteRunner(classifier_path)

predict(x_batch)

Run the backbone and feed its embeddings to the classifier head.

Parameters:

Name Type Description Default
x_batch ndarray

Input batch in the backbone's expected shape and dtype.

required

Returns:

Type Description
ndarray

Model outputs [B, C] as float32.

Source code in birdnet_stm32/models/runners.py
def predict(self, x_batch: np.ndarray) -> np.ndarray:
    """Run the backbone and feed its embeddings to the classifier head.

    Args:
        x_batch: Input batch in the backbone's expected shape and dtype.

    Returns:
        Model outputs [B, C] as float32.
    """
    embeddings = self.backbone.predict(x_batch)
    return self.classifier.predict(np.asarray(embeddings, dtype=np.float32))

load_keras_model(model_path)

Load a project checkpoint with the canonical custom-layer registry.

Source code in birdnet_stm32/models/runners.py
def load_keras_model(model_path: str) -> tf.keras.Model:
    """Load a project checkpoint with the canonical custom-layer registry."""
    return tf.keras.models.load_model(
        model_path,
        compile=False,
        custom_objects=_KERAS_CUSTOM_OBJECTS,
    )

allocated_interpreter(model_path=None, model_content=None, num_threads=None)

Return a TFLite interpreter with its tensors allocated.

TFLite applies the XNNPACK delegate by default. XNNPACK refuses some valid INT8 graphs outright, for example a requantization scale of 256 or more, which PTQ produces when a branch calibrates to an almost empty range (a PWL hinge that never fires). The builtin reference kernels run those graphs, so the interpreter falls back to them with a warning instead of failing. Graphs XNNPACK accepts keep running on it, so their outputs are unchanged.

Parameters:

Name Type Description Default
model_path str | None

Path to a .tflite file.

None
model_content bytes | None

Serialized model, instead of a path.

None
num_threads int | None

Interpreter threads.

None

Returns:

Type Description
Interpreter

Interpreter with tensors allocated.

Source code in birdnet_stm32/models/runners.py
def allocated_interpreter(
    model_path: str | None = None, model_content: bytes | None = None, num_threads: int | None = None
) -> tf.lite.Interpreter:
    """Return a TFLite interpreter with its tensors allocated.

    TFLite applies the XNNPACK delegate by default. XNNPACK refuses some valid
    INT8 graphs outright, for example a requantization scale of 256 or more,
    which PTQ produces when a branch calibrates to an almost empty range (a PWL
    hinge that never fires). The builtin reference kernels run those graphs, so
    the interpreter falls back to them with a warning instead of failing.
    Graphs XNNPACK accepts keep running on it, so their outputs are unchanged.

    Args:
        model_path: Path to a .tflite file.
        model_content: Serialized model, instead of a path.
        num_threads: Interpreter threads.

    Returns:
        Interpreter with tensors allocated.
    """
    kwargs = {"model_path": model_path, "model_content": model_content, "num_threads": num_threads}
    interpreter = tf.lite.Interpreter(**kwargs, experimental_delegates=[])
    try:
        interpreter.allocate_tensors()
        return interpreter
    except RuntimeError as exc:
        if "XNNPACK" not in str(exc):
            raise
        warnings.warn(
            f"XNNPACK cannot run {model_path or 'this model'} ({exc}); using the builtin reference kernels.",
            RuntimeWarning,
            stacklevel=2,
        )
    interpreter = tf.lite.Interpreter(
        **kwargs,
        experimental_op_resolver_type=tf.lite.experimental.OpResolverType.BUILTIN_WITHOUT_DEFAULT_DELEGATES,
    )
    interpreter.allocate_tensors()
    return interpreter

load_model_runner(model_path, classifier_path='')

Load a .keras or .tflite model and return a runner with predict().

Parameters:

Name Type Description Default
model_path str

Path to a saved model (.keras or .tflite). When classifier_path is given, this is the .tflite backbone.

required
classifier_path str

Optional .tflite classifier head. Supplying it runs the split pair as one chained model.

''

Returns:

Type Description
KerasRunner | TFLiteRunner | ChainedTFLiteRunner

KerasRunner, TFLiteRunner, or ChainedTFLiteRunner instance.

Raises:

Type Description
ValueError

If a classifier head is paired with a non-TFLite backbone.

Source code in birdnet_stm32/models/runners.py
def load_model_runner(
    model_path: str,
    classifier_path: str = "",
) -> KerasRunner | TFLiteRunner | ChainedTFLiteRunner:
    """Load a .keras or .tflite model and return a runner with predict().

    Args:
        model_path: Path to a saved model (.keras or .tflite). When
            ``classifier_path`` is given, this is the .tflite backbone.
        classifier_path: Optional .tflite classifier head. Supplying it runs
            the split pair as one chained model.

    Returns:
        KerasRunner, TFLiteRunner, or ChainedTFLiteRunner instance.

    Raises:
        ValueError: If a classifier head is paired with a non-TFLite backbone.
    """
    if classifier_path:
        if not model_path.lower().endswith(".tflite"):
            raise ValueError("A classifier head can only be chained onto a .tflite backbone")
        return ChainedTFLiteRunner(model_path, classifier_path)
    if model_path.lower().endswith(".tflite"):
        return TFLiteRunner(model_path)
    return KerasRunner(load_keras_model(model_path))