trainer

birdnet_stm32.training.trainer

Training loop with cosine LR schedule, early stopping, and checkpointing.

WarmupCosineDecay

Bases: LearningRateSchedule

Linear warmup followed by cosine decay to zero.

Also carries a step offset so a resumed run continues along the same schedule instead of jumping back to the peak learning rate.

Parameters:

Name	Type	Description	Default
initial_learning_rate	float	Peak learning rate reached at the end of warmup.	required
decay_steps	int	Total number of steps in the full run (warmup included).	required
warmup_steps	int	Number of steps spent ramping up linearly from ~0.	0
offset_steps	int	Steps already completed by a previous run.	0

Source code in birdnet_stm32/training/trainer.py

class WarmupCosineDecay(tf.keras.optimizers.schedules.LearningRateSchedule):
    """Linear warmup followed by cosine decay to zero.

    Also carries a step offset so a resumed run continues along the same
    schedule instead of jumping back to the peak learning rate.

    Args:
        initial_learning_rate: Peak learning rate reached at the end of warmup.
        decay_steps: Total number of steps in the full run (warmup included).
        warmup_steps: Number of steps spent ramping up linearly from ~0.
        offset_steps: Steps already completed by a previous run.
    """

    def __init__(
        self,
        initial_learning_rate: float,
        decay_steps: int,
        warmup_steps: int = 0,
        offset_steps: int = 0,
    ):
        super().__init__()
        self.initial_learning_rate = float(initial_learning_rate)
        self.decay_steps = max(1, int(decay_steps))
        self.warmup_steps = max(0, int(warmup_steps))
        self.offset_steps = max(0, int(offset_steps))

    def __call__(self, step):
        step = tf.cast(step, tf.float32) + float(self.offset_steps)
        peak = tf.constant(self.initial_learning_rate, tf.float32)
        warmup = tf.constant(float(self.warmup_steps), tf.float32)
        total = tf.constant(float(self.decay_steps), tf.float32)

        warmup_lr = peak * (step + 1.0) / tf.maximum(warmup, 1.0)

        progress = (step - warmup) / tf.maximum(total - warmup, 1.0)
        progress = tf.clip_by_value(progress, 0.0, 1.0)
        cosine_lr = peak * 0.5 * (1.0 + tf.cos(tf.constant(3.14159265, tf.float32) * progress))

        if self.warmup_steps == 0:
            return cosine_lr
        return tf.where(step < warmup, warmup_lr, cosine_lr)

    def get_config(self) -> dict:
        """Return a serializable configuration dict."""
        return {
            "initial_learning_rate": self.initial_learning_rate,
            "decay_steps": self.decay_steps,
            "warmup_steps": self.warmup_steps,
            "offset_steps": self.offset_steps,
        }

get_config()

Return a serializable configuration dict.

Source code in birdnet_stm32/training/trainer.py

def get_config(self) -> dict:
    """Return a serializable configuration dict."""
    return {
        "initial_learning_rate": self.initial_learning_rate,
        "decay_steps": self.decay_steps,
        "warmup_steps": self.warmup_steps,
        "offset_steps": self.offset_steps,
    }

monitor_mode(monitor)

Return "max" or "min" for a validation metric name.

Source code in birdnet_stm32/training/trainer.py

def monitor_mode(monitor: str) -> str:
    """Return ``"max"`` or ``"min"`` for a validation metric name."""
    stem = monitor[4:] if monitor.startswith("val_") else monitor
    return "max" if stem in _MAXIMISED_METRICS else "min"

load_for_resume(checkpoint_path)

Load a checkpoint to continue training it, keeping each layer's saved dtype policy.

load_keras_model rebuilds a mixed-precision checkpoint in float32, which evaluation and conversion need; a resumed run must keep the policy it was trained with, or it trains on in float32 (2026-10-14: 87 instead of 55 ms per step on the probe that resumed after a stop).

Source code in birdnet_stm32/training/trainer.py

def load_for_resume(checkpoint_path: str) -> tf.keras.Model:
    """Load a checkpoint to continue training it, keeping each layer's saved dtype policy.

    ``load_keras_model`` rebuilds a mixed-precision checkpoint in float32, which evaluation and conversion need; a
    resumed run must keep the policy it was trained with, or it trains on in float32 (2026-10-14: 87 instead of
    55 ms per step on the probe that resumed after a stop).
    """
    from birdnet_stm32.models.runners import _KERAS_CUSTOM_OBJECTS

    return tf.keras.models.load_model(checkpoint_path, compile=False, custom_objects=_KERAS_CUSTOM_OBJECTS)

train_model(model, train_dataset, val_dataset, epochs=50, learning_rate=0.001, batch_size=32, patience=10, checkpoint_path='checkpoints/best_model.keras', steps_per_epoch=None, val_steps=None, optimizer='adam', weight_decay=0.0, loss_fn=None, gradient_clip_norm=1.0, resume=False, extra_callbacks=None, jit_compile=False, checkpoint_model=None, checkpoint_sync=None, checkpoint_monitor=_MONITOR, checkpoint_mode=_MONITOR_MODE, checkpoint_start_epoch=0, checkpoint_managed=False, training_wrapper=None, warmup_epochs=_WARMUP_EPOCHS)

Train a model with cosine LR schedule, early stopping, and checkpointing.

The classifier head is always sigmoid + binary crossentropy (multi-label). Soundscape recordings are inherently multi-label even when the source label is a single species, so we always optimise per-class probabilities.

Checkpointing and early stopping track validation cmAP; the best model is saved as a full .keras file.

Parameters:

Name	Type	Description	Default
model	Model	Model to train.	required
train_dataset	Dataset	Training dataset (infinite).	required
val_dataset	Dataset	Validation dataset (infinite).	required
epochs	int	Number of epochs.	50
learning_rate	float	Initial learning rate for cosine schedule.	0.001
batch_size	int	Unused; kept for API symmetry with data loader.	32
patience	int	Early stopping patience (epochs). Early stopping only starts after EARLY_STOP_FROM of epochs: under a cosine schedule the learning rate is still high in the middle of a long run and the validation metric plateaus noisily there (a 200-epoch run stopped at epoch 35 of 200 on a 10-epoch plateau).	10
checkpoint_path	str	Path to save the best .keras model.	'checkpoints/best_model.keras'
steps_per_epoch	int | None	Training steps per epoch (> 0 required).	None
val_steps	int | None	Validation steps per epoch (defaults to 1 if <= 0).	None
optimizer	str	Optimizer name ('adam', 'sgd', or 'adamw').	'adam'
weight_decay	float	Weight decay factor (only used by 'adamw').	0.0
loss_fn	str | Loss | None	Optional custom loss function. Defaults to binary_crossentropy.	None
gradient_clip_norm	float	Max gradient norm for clipping (0 = disabled).	1.0
resume	bool	If True, reload the model from the checkpoint and continue from the recorded epoch (the learning-rate schedule is advanced to match).	False
extra_callbacks	list[Callback] | None	Additional Keras callbacks (e.g. QAT callback).	None
jit_compile	bool	XLA-compile the training step (float training; not with fake-quant QAT models).	False
checkpoint_model	Model | None	Optional deployment model that shares weights with model. When supplied, checkpoints contain this clean model instead of training-only wrappers such as fake-quant layers.	None
checkpoint_sync	Callable[[], None] | None	Optional hook that copies separately cloned weights into checkpoint_model immediately before each save.	None
checkpoint_monitor	str	Validation metric used for checkpoint selection and early stopping.	_MONITOR
checkpoint_mode	str	Whether a larger (max) or smaller (min) monitored value is better.	_MONITOR_MODE
checkpoint_managed	bool	An external callback owns checkpoint artifacts (INT8 selection).	False
checkpoint_start_epoch	int	First epoch (0-based) eligible for checkpoint selection and early stopping. Compression schedules use this so a model that has not yet reached its target sparsity or quantization noise cannot win on the strength of the perturbation it is missing.	0
training_wrapper	Callable[[Model], Model] | None	Optional function that wraps the model for training only (e.g. with a training-only loss head). It is applied after a resume reload, and the unwrapped model becomes the checkpoint.	None
warmup_epochs	float	Length of the linear learning-rate warm-up, in epochs (may be fractional: with one epoch per pass over a large dataset, half an epoch).	_WARMUP_EPOCHS

Returns:

Type	Description
History	Keras training history.

Raises:

Type	Description
ValueError	If steps_per_epoch is not positive.

Source code in birdnet_stm32/training/trainer.py

def train_model(
    model: tf.keras.Model,
    train_dataset: tf.data.Dataset,
    val_dataset: tf.data.Dataset,
    epochs: int = 50,
    learning_rate: float = 0.001,
    batch_size: int = 32,
    patience: int = 10,
    checkpoint_path: str = "checkpoints/best_model.keras",
    steps_per_epoch: int | None = None,
    val_steps: int | None = None,
    optimizer: str = "adam",
    weight_decay: float = 0.0,
    loss_fn: str | tf.keras.losses.Loss | None = None,
    gradient_clip_norm: float = 1.0,
    resume: bool = False,
    extra_callbacks: list[tf.keras.callbacks.Callback] | None = None,
    jit_compile: bool = False,
    checkpoint_model: tf.keras.Model | None = None,
    checkpoint_sync: Callable[[], None] | None = None,
    checkpoint_monitor: str = _MONITOR,
    checkpoint_mode: str = _MONITOR_MODE,
    checkpoint_start_epoch: int = 0,
    checkpoint_managed: bool = False,
    training_wrapper: Callable[[tf.keras.Model], tf.keras.Model] | None = None,
    warmup_epochs: float = _WARMUP_EPOCHS,
) -> tf.keras.callbacks.History:
    """Train a model with cosine LR schedule, early stopping, and checkpointing.

    The classifier head is always sigmoid + binary crossentropy (multi-label).
    Soundscape recordings are inherently multi-label even when the source
    label is a single species, so we always optimise per-class probabilities.

    Checkpointing and early stopping track validation cmAP; the best model
    is saved as a full .keras file.

    Args:
        model: Model to train.
        train_dataset: Training dataset (infinite).
        val_dataset: Validation dataset (infinite).
        epochs: Number of epochs.
        learning_rate: Initial learning rate for cosine schedule.
        batch_size: Unused; kept for API symmetry with data loader.
        patience: Early stopping patience (epochs). Early stopping only starts
            after ``EARLY_STOP_FROM`` of ``epochs``: under a cosine schedule the
            learning rate is still high in the middle of a long run and the
            validation metric plateaus noisily there (a 200-epoch run stopped at
            epoch 35 of 200 on a 10-epoch plateau).
        checkpoint_path: Path to save the best .keras model.
        steps_per_epoch: Training steps per epoch (> 0 required).
        val_steps: Validation steps per epoch (defaults to 1 if <= 0).
        optimizer: Optimizer name ('adam', 'sgd', or 'adamw').
        weight_decay: Weight decay factor (only used by 'adamw').
        loss_fn: Optional custom loss function. Defaults to ``binary_crossentropy``.
        gradient_clip_norm: Max gradient norm for clipping (0 = disabled).
        resume: If True, reload the model from the checkpoint and continue from
            the recorded epoch (the learning-rate schedule is advanced to match).
        extra_callbacks: Additional Keras callbacks (e.g. QAT callback).
        jit_compile: XLA-compile the training step (float training; not with fake-quant QAT models).
        checkpoint_model: Optional deployment model that shares weights with
            ``model``. When supplied, checkpoints contain this clean model
            instead of training-only wrappers such as fake-quant layers.
        checkpoint_sync: Optional hook that copies separately cloned weights
            into ``checkpoint_model`` immediately before each save.
        checkpoint_monitor: Validation metric used for checkpoint selection
            and early stopping.
        checkpoint_mode: Whether a larger (``max``) or smaller (``min``)
            monitored value is better.
        checkpoint_managed: An external callback owns checkpoint artifacts (INT8 selection).
        checkpoint_start_epoch: First epoch (0-based) eligible for checkpoint
            selection and early stopping. Compression schedules use this so a
            model that has not yet reached its target sparsity or quantization
            noise cannot win on the strength of the perturbation it is missing.
        training_wrapper: Optional function that wraps the model for training
            only (e.g. with a training-only loss head). It is applied after a
            ``resume`` reload, and the unwrapped model becomes the checkpoint.
        warmup_epochs: Length of the linear learning-rate warm-up, in epochs (may be
            fractional: with one epoch per pass over a large dataset, half an epoch).

    Returns:
        Keras training history.

    Raises:
        ValueError: If ``steps_per_epoch`` is not positive.
    """
    if steps_per_epoch is None or steps_per_epoch <= 0:
        raise ValueError("steps_per_epoch must be > 0")
    if val_steps is None or val_steps <= 0:
        val_steps = 1
    if checkpoint_mode not in {"min", "max"}:
        raise ValueError("checkpoint_mode must be 'min' or 'max'")
    checkpoint_start_epoch = max(0, int(checkpoint_start_epoch))

    os.makedirs(os.path.dirname(checkpoint_path) or ".", exist_ok=True)

    # Resume: reload model from checkpoint if it exists
    initial_epoch = 0
    state_path = checkpoint_path.replace(".keras", "_train_state.json")
    if resume and os.path.isfile(checkpoint_path):
        print(f"[resume] Loading model from {checkpoint_path}")
        model = load_for_resume(checkpoint_path)
        if os.path.isfile(state_path):
            with open(state_path) as f:
                state = json.load(f)
            initial_epoch = state.get("epoch", 0)
            print(f"[resume] Resuming from epoch {initial_epoch}")

    if training_wrapper is not None:
        if checkpoint_model is not None:
            raise ValueError("training_wrapper and checkpoint_model cannot be combined")
        checkpoint_model = model
        model = training_wrapper(model)

    warmup_steps = int(round(max(0.0, warmup_epochs) * steps_per_epoch))
    lr_schedule = WarmupCosineDecay(
        initial_learning_rate=learning_rate,
        decay_steps=epochs * steps_per_epoch,
        warmup_steps=warmup_steps,
        offset_steps=initial_epoch * steps_per_epoch,
    )
    print(f"LR schedule: {warmup_epochs:g} warmup epoch(s) -> cosine decay over {epochs} epochs.")

    opt = _build_optimizer(optimizer, lr_schedule, weight_decay, gradient_clip_norm)

    if loss_fn is None:
        loss_fn = "binary_crossentropy"

    num_labels = None
    try:
        num_labels = int(model.output_shape[-1])
    except (TypeError, IndexError):
        num_labels = None

    auc_metric = tf.keras.metrics.AUC(
        curve="ROC",
        multi_label=True,
        num_labels=num_labels,
        name="roc_auc",
    )

    # Class-macro average precision, tracked alongside ROC-AUC because the two
    # do not move together on a long-tail multi-label problem. Measured on the
    # 100-output v0.2 schema, quantization cost 0.005 val ROC-AUC and 0.035
    # catalog cmAP: a run that looks healthy on ROC-AUC can be losing an order
    # of magnitude more of what the model is actually selected on. Threshold-free
    # ROC-AUC is dominated by the easy negative mass and saturates; AP is not.
    #
    # PR-AUC is only a training diagnostic. An exact AP callback supplies the
    # selection metric before checkpointing and early stopping run.
    pr_metric = tf.keras.metrics.AUC(
        curve="PR",
        multi_label=True,
        num_labels=num_labels,
        name="pr_auc",
    )

    model.compile(
        optimizer=opt,
        loss=loss_fn,
        metrics=[auc_metric, pr_metric],
        # Off by default: fake-quant operators (QAT) compile fragile and memory
        # hungry under XLA. Float training of the raw 448 model is another matter:
        # the depthwise backbone runs ~3x faster with XLA and mixed precision
        # (``train --jit_compile --mixed_precision``).
        jit_compile=jit_compile,
    )

    class _SaveTrainState(tf.keras.callbacks.Callback):
        """Save epoch counter alongside checkpoint for resume support."""

        def on_epoch_end(self, epoch, logs=None):
            with open(state_path, "w") as f:
                json.dump({"epoch": epoch + 1}, f)

    class _CSVHistoryLogger(tf.keras.callbacks.Callback):
        """Append per-epoch metrics to a CSV file alongside the checkpoint."""

        def __init__(self, csv_path):
            super().__init__()
            self.csv_path = csv_path
            self._header_written = os.path.isfile(csv_path) and resume

        def on_epoch_end(self, epoch, logs=None):
            logs = logs or {}
            import csv

            write_header = not self._header_written
            with open(self.csv_path, "w" if write_header else "a", newline="") as f:
                writer = csv.DictWriter(f, fieldnames=["epoch"] + sorted(logs.keys()))
                if write_header:
                    writer.writeheader()
                    self._header_written = True
                row = {"epoch": epoch + 1}
                row.update({k: f"{v:.6f}" for k, v in logs.items()})
                writer.writerow(row)

    class _SharedWeightsCheckpoint(tf.keras.callbacks.Callback):
        """Checkpoint a clean model that shares variables with the train graph."""

        def __init__(self, target: tf.keras.Model):
            super().__init__()
            self.target = target
            self.best = -float("inf") if checkpoint_mode == "max" else float("inf")

        def on_epoch_end(self, epoch, logs=None):
            value = (logs or {}).get(checkpoint_monitor)
            if value is None or epoch < checkpoint_start_epoch:
                return
            improved = value > self.best if checkpoint_mode == "max" else value < self.best
            if improved:
                self.best = float(value)
                if checkpoint_sync is not None:
                    checkpoint_sync()
                self.target.save(checkpoint_path)

    csv_path = checkpoint_path.replace(".keras", "_history.csv")

    class _DelayedModelCheckpoint(tf.keras.callbacks.ModelCheckpoint):
        """Ignore epochs before ``checkpoint_start_epoch`` when selecting."""

        def on_epoch_end(self, epoch, logs=None):
            if epoch < checkpoint_start_epoch:
                return
            super().on_epoch_end(epoch, logs)

    checkpoint_callback: tf.keras.callbacks.Callback
    if checkpoint_model is None or checkpoint_model is model:
        checkpoint_callback = _DelayedModelCheckpoint(
            checkpoint_path,
            monitor=checkpoint_monitor,
            save_best_only=True,
            mode=checkpoint_mode,
            save_weights_only=False,
        )
    else:
        checkpoint_callback = _SharedWeightsCheckpoint(checkpoint_model)

    from birdnet_stm32.training.validation import ExactCmap, FileCmap

    callbacks = list(extra_callbacks or [])
    if not any(isinstance(callback, FileCmap) for callback in callbacks):
        callbacks.insert(0, ExactCmap(val_dataset, val_steps))
    callbacks += [
        tf.keras.callbacks.EarlyStopping(
            monitor=checkpoint_monitor,
            patience=patience,
            restore_best_weights=not checkpoint_managed,
            mode=checkpoint_mode,
            start_from_epoch=max(checkpoint_start_epoch, int(EARLY_STOP_FROM * epochs)),
        ),
        *([] if checkpoint_managed else [checkpoint_callback]),
        _SaveTrainState(),
        _CSVHistoryLogger(csv_path),
    ]
    history = model.fit(
        train_dataset,
        validation_data=val_dataset,
        epochs=epochs,
        initial_epoch=initial_epoch,
        steps_per_epoch=steps_per_epoch,
        validation_steps=val_steps,
        callbacks=callbacks,
    )

    # Save training curves as PNG
    _save_training_curves(history, checkpoint_path.replace(".keras", "_curves.png"))

    return history

compute_hop_length(sample_rate, chunk_duration, spec_width)

Compute hop length to produce spec_width frames from an input chunk.

Parameters:

Name	Type	Description	Default
sample_rate	int	Sampling rate (Hz).	required
chunk_duration	float	Chunk duration (seconds).	required
spec_width	int	Desired number of frames.	required

Returns:

Type	Description
int	Hop length in samples (floor(T / spec_width), at least 1).

Source code in birdnet_stm32/training/trainer.py

def compute_hop_length(sample_rate: int, chunk_duration: float, spec_width: int) -> int:
    """Compute hop length to produce spec_width frames from an input chunk.

    Args:
        sample_rate: Sampling rate (Hz).
        chunk_duration: Chunk duration (seconds).
        spec_width: Desired number of frames.

    Returns:
        Hop length in samples (floor(T / spec_width), at least 1).
    """
    T = int(sample_rate * chunk_duration)
    return max(1, T // int(spec_width))