qat

birdnet_stm32.training.qat

Keras 3 quantization-aware fine-tuning for full-INT8 TFLite deployment.

The training graph simulates both quantized kernels and per-tensor activation requantization. It shares variables with a clean deployment graph, so the saved checkpoint contains no FakeQuant operators or training-only wrappers.

FakeQuantActivation

Bases: Layer

Static per-tensor INT8 fake quantizer with an asymmetric zero point.

Source code in birdnet_stm32/training/qat.py

@tf.keras.utils.register_keras_serializable(package="birdnet_stm32")
class FakeQuantActivation(layers.Layer):
    """Static per-tensor INT8 fake quantizer with an asymmetric zero point."""

    def __init__(self, minimum: float, maximum: float, num_bits: int = 8, **kwargs):
        super().__init__(trainable=False, **kwargs)
        self.minimum, self.maximum = self._normalize(minimum, maximum)
        self.num_bits = int(num_bits)
        # The range lives in variables so it can be refreshed between epochs
        # without retracing the training step (see refresh_activation_ranges).
        self._range = None

    @staticmethod
    def _normalize(minimum: float, maximum: float) -> tuple[float, float]:
        minimum = min(float(minimum), 0.0)
        maximum = max(float(maximum), 0.0)
        if maximum - minimum < 1e-6:
            maximum = minimum + 1e-6
        return minimum, maximum

    def build(self, input_shape):
        self._range = self.add_weight(
            name="range",
            shape=(2,),
            initializer=tf.keras.initializers.Constant([self.minimum, self.maximum]),
            trainable=False,
        )
        super().build(input_shape)

    def set_range(self, minimum: float, maximum: float) -> None:
        """Move the grid, e.g. to the converter's range for the current weights."""
        self.minimum, self.maximum = self._normalize(minimum, maximum)
        if self._range is not None:
            self._range.assign([self.minimum, self.maximum])

    def call(self, inputs):
        """Apply the same scalar affine grid used for TFLite activations."""
        qmin = -(1 << (self.num_bits - 1))
        qmax = (1 << (self.num_bits - 1)) - 1
        minimum = tf.cast(self._range[0], inputs.dtype)
        maximum = tf.cast(self._range[1], inputs.dtype)
        scale = (maximum - minimum) / tf.cast(qmax - qmin, inputs.dtype)
        zero_point = tf.clip_by_value(
            tf.round(tf.cast(qmin, inputs.dtype) - minimum / scale),
            tf.cast(qmin, inputs.dtype),
            tf.cast(qmax, inputs.dtype),
        )
        quantized = tf.clip_by_value(tf.round(inputs / scale) + zero_point, qmin, qmax)
        dequantized = (quantized - zero_point) * scale
        # Explicit straight-through estimator. Unlike TensorFlow's FakeQuant
        # gradient kernel, this is supported by deterministic GPU execution.
        return inputs + tf.stop_gradient(dequantized - inputs)

    def get_config(self):
        """Return serializable quantizer settings."""
        config = super().get_config()
        config.update({"minimum": self.minimum, "maximum": self.maximum, "num_bits": self.num_bits})
        return config

set_range(minimum, maximum)

Move the grid, e.g. to the converter's range for the current weights.

Source code in birdnet_stm32/training/qat.py

def set_range(self, minimum: float, maximum: float) -> None:
    """Move the grid, e.g. to the converter's range for the current weights."""
    self.minimum, self.maximum = self._normalize(minimum, maximum)
    if self._range is not None:
        self._range.assign([self.minimum, self.maximum])

call(inputs)

Apply the same scalar affine grid used for TFLite activations.

Source code in birdnet_stm32/training/qat.py

def call(self, inputs):
    """Apply the same scalar affine grid used for TFLite activations."""
    qmin = -(1 << (self.num_bits - 1))
    qmax = (1 << (self.num_bits - 1)) - 1
    minimum = tf.cast(self._range[0], inputs.dtype)
    maximum = tf.cast(self._range[1], inputs.dtype)
    scale = (maximum - minimum) / tf.cast(qmax - qmin, inputs.dtype)
    zero_point = tf.clip_by_value(
        tf.round(tf.cast(qmin, inputs.dtype) - minimum / scale),
        tf.cast(qmin, inputs.dtype),
        tf.cast(qmax, inputs.dtype),
    )
    quantized = tf.clip_by_value(tf.round(inputs / scale) + zero_point, qmin, qmax)
    dequantized = (quantized - zero_point) * scale
    # Explicit straight-through estimator. Unlike TensorFlow's FakeQuant
    # gradient kernel, this is supported by deterministic GPU execution.
    return inputs + tf.stop_gradient(dequantized - inputs)

get_config()

Return serializable quantizer settings.

Source code in birdnet_stm32/training/qat.py

def get_config(self):
    """Return serializable quantizer settings."""
    config = super().get_config()
    config.update({"minimum": self.minimum, "maximum": self.maximum, "num_bits": self.num_bits})
    return config

RangeRefresh

Bases: Callback

Recalibrate QAT activation ranges on the current weights after each epoch.

QAT otherwise trains against ranges measured once on the starting weights, while conversion recalibrates the weights it is given. The longer and the faster fine-tuning moves the weights, the further the two grids drift.

Source code in birdnet_stm32/training/qat.py

class RangeRefresh(tf.keras.callbacks.Callback):
    """Recalibrate QAT activation ranges on the current weights after each epoch.

    QAT otherwise trains against ranges measured once on the starting weights,
    while conversion recalibrates the weights it is given. The longer and the
    faster fine-tuning moves the weights, the further the two grids drift.
    """

    def __init__(self, qat_student, deployment, calibration, sync, samples: int):
        super().__init__()
        self.qat_student, self.deployment = qat_student, deployment
        self.calibration, self.sync, self.samples = calibration, sync, samples

    def on_epoch_end(self, epoch, logs=None):
        self.sync()
        ranges = calibrate_activation_ranges(
            self.deployment, self.calibration, max_samples=self.samples, percentile=range_percentile
        )
        moved = refresh_activation_ranges(self.qat_student, ranges)
        print(f"[QAT] Refreshed {moved} activation ranges on epoch {epoch + 1} weights")

fake_quantize_weights(w, num_bits=8, per_channel=True, channel_axis=-1)

Quantize and dequantize a kernel on TFLite's symmetric INT8 grid.

Source code in birdnet_stm32/training/qat.py

def fake_quantize_weights(
    w: np.ndarray,
    num_bits: int = 8,
    per_channel: bool = True,
    channel_axis: int = -1,
) -> np.ndarray:
    """Quantize and dequantize a kernel on TFLite's symmetric INT8 grid."""
    qmax = (1 << (num_bits - 1)) - 1
    if per_channel and w.ndim > 1:
        axis = channel_axis % w.ndim
        reduce_axes = tuple(i for i in range(w.ndim) if i != axis)
        amax = np.max(np.abs(w), axis=reduce_axes, keepdims=True)
    else:
        amax = np.max(np.abs(w))
    scale = np.maximum(amax / qmax, 1e-12)
    return np.asarray(np.clip(np.round(w / scale), -qmax, qmax) * scale, dtype=np.float32)

add_to_histogram(hist, name, values, ranges, bins)

Accumulate values into hist[name] over the tensor's calibrated range (skipped when it cannot be binned).

Source code in birdnet_stm32/training/qat.py

def add_to_histogram(hist: dict, name: str, values: np.ndarray, ranges: dict, bins: int) -> None:
    """Accumulate ``values`` into ``hist[name]`` over the tensor's calibrated range (skipped when it cannot be binned)."""
    from birdnet_stm32.conversion.quantize import _binnable

    lo, hi = ranges.get(name, (0.0, 0.0))
    if not _binnable(lo, hi):
        return
    counts = np.histogram(np.asarray(values, np.float32), bins=bins, range=(np.float32(lo), np.float32(hi)))[0]
    hist[name] = hist.get(name, 0) + counts

percentile_range(lo, hi, counts, percentile)

Cut a histogrammed range to percentile at each end as the converter does (the low end only if negative).

Source code in birdnet_stm32/training/qat.py

def percentile_range(lo: float, hi: float, counts: np.ndarray, percentile: float) -> tuple[float, float]:
    """Cut a histogrammed range to ``percentile`` at each end as the converter does (the low end only if negative)."""
    total = counts.sum()
    if total == 0:
        return lo, hi
    bins = len(counts)
    tail = (100.0 - percentile) / 100.0
    edges = np.linspace(lo, hi, bins + 1)
    cdf = np.cumsum(counts) / total
    new_hi = min(hi, float(edges[min(bins, np.searchsorted(cdf, 1.0 - tail) + 1)]))
    new_lo = max(lo, float(edges[np.searchsorted(cdf, tail)])) if lo < 0 else lo
    return (lo, hi) if new_hi <= new_lo else (min(new_lo, 0.0), max(new_hi, 0.0))

calibrate_activation_ranges(model, dataset, max_samples=64, percentile=None)

Measure activation ranges on real inputs, as the converter does.

Min/max by default. With percentile (run_qat passes the converter's range_percentile, a function of the tensor name), a second pass histograms every tensor over its min/max and cuts the range to that percentile, as quantize.clip_activation_ranges cuts the deployed model's ranges since 2.0; the model input and the output logits keep min/max there and here. Without it QAT would train against coarser grids than the converted model has: on the B0 probe, min/max ranges outside the frontend turned a 0.997-cosine conversion into a 0.71 one.

Parameters:

Name	Type	Description	Default
model	Model	Deployment model to probe.	required
dataset	Iterable	Calibration data (iterated twice with percentile).	required
max_samples	int	Number of samples to observe.	64
percentile	float | Callable[[str], float] | None	None for min/max, or the per-tensor percentile kept at each end.	None

Source code in birdnet_stm32/training/qat.py

def calibrate_activation_ranges(
    model: tf.keras.Model,
    dataset: Iterable,
    max_samples: int = 64,
    percentile: float | Callable[[str], float] | None = None,
) -> dict[str, tuple[float, float]]:
    """Measure activation ranges on real inputs, as the converter does.

    Min/max by default. With ``percentile`` (``run_qat`` passes the converter's
    ``range_percentile``, a function of the tensor name), a second pass histograms every tensor over
    its min/max and cuts the range to that percentile, as
    ``quantize.clip_activation_ranges`` cuts the deployed model's ranges since 2.0;
    the model input and the output logits keep min/max there and here. Without it
    QAT would train against coarser grids than the converted model has: on the B0
    probe, min/max ranges outside the frontend turned a 0.997-cosine conversion into
    a 0.71 one.

    Args:
        model: Deployment model to probe.
        dataset: Calibration data (iterated twice with ``percentile``).
        max_samples: Number of samples to observe.
        percentile: None for min/max, or the per-tensor percentile kept at each end.
    """
    from birdnet_stm32.conversion.quantize import _RANGE_BINS

    if max_samples <= 0:
        raise ValueError("Calibration sample count must be positive")
    boundaries = [layer for layer in model.layers if _is_activation_boundary(layer)]
    if not boundaries:
        raise ValueError("Model has no supported activation quantization boundaries")
    sigmoid_heads = [layer for layer in boundaries if _is_sigmoid_dense(layer)]
    probe = tf.keras.Model(
        model.inputs, [layer.output for layer in boundaries] + [layer.input for layer in sigmoid_heads]
    )
    frontends = [layer for layer in all_layers(model) if layer.__class__.__name__ == "AudioFrontendLayer"]

    def observe(hook, record) -> int:
        """Run up to ``max_samples`` through the probe; ``record(name, array)`` sees every observed tensor."""
        for frontend in frontends:
            frontend.set_quantization_hook(hook)
        seen = 0
        try:
            for batch in dataset:
                inputs = np.asarray(batch[0] if isinstance(batch, (tuple, list)) else batch)
                for sample in inputs:
                    record("__input__", sample)
                    outputs = probe([sample[None]], training=False)
                    if not isinstance(outputs, (tuple, list)):
                        outputs = [outputs]
                    for layer, output in zip(boundaries, outputs[: len(boundaries)], strict=True):
                        record(layer.name, np.asarray(output))
                    for layer, features in zip(sigmoid_heads, outputs[len(boundaries) :], strict=True):
                        logits = np.asarray(features) @ np.asarray(layer.kernel)
                        if layer.bias is not None:
                            logits += np.asarray(layer.bias)
                        hook.activation(f"{layer.name}__logits", logits)
                    seen += 1
                    if seen >= max_samples:
                        return seen
        finally:
            for frontend in frontends:
                frontend.set_quantization_hook(None)
        return seen

    ranges = {layer.name: [float("inf"), -float("inf")] for layer in boundaries}
    ranges["__input__"] = [float("inf"), -float("inf")]
    collector = _ActivationRangeCollector()

    def record_minmax(name, array):
        ranges[name][0] = min(ranges[name][0], float(np.min(array)), 0.0)
        ranges[name][1] = max(ranges[name][1], float(np.max(array)), 0.0)

    seen = observe(collector, record_minmax)
    if seen == 0:
        raise ValueError("Activation calibration dataset yielded no samples")
    ranges.update(collector.ranges)
    result = {name: (values[0], values[1]) for name, values in ranges.items()}

    if percentile is not None:
        hist_hook = _ActivationHistogramCollector(result, _RANGE_BINS)
        observe(hist_hook, lambda name, array: add_to_histogram(hist_hook.hist, name, array, result, _RANGE_BINS))
        outputs = {layer.name for layer in model.layers if any(layer.output is o for o in model.outputs)}
        keep = {"__input__"} | outputs | {f"{name}__logits" for name in outputs}
        clipped = 0
        for name, counts in hist_hook.hist.items():
            if name in keep:
                continue
            new = percentile_range(*result[name], counts, percentile(name) if callable(percentile) else percentile)
            clipped += new != result[name]
            result[name] = new
        label = "per-tensor" if callable(percentile) else percentile
        print(f"[QAT] Percentile {label} ranges: {clipped} of {len(result)} tensors cut below min/max")

    print(f"[QAT] Calibrated {len(result)} activation tensors on {seen} samples")
    for layer in sigmoid_heads:
        # TFLite LOGISTIC has fixed scale 1/256 and zero point -128.
        result[layer.name] = (0.0, 255.0 / 256.0)
    return result

build_qat_model(deployment_model, activation_ranges, logit_output=True)

Build an activation-fake-quant graph sharing deployment model weights.

A sigmoid head is simulated as the release converts it: INT8 logits, the sigmoid applied after (convert --output_activation logit, every release from 1.5). logit_output=False also puts the probabilities on the INT8 output grid of a sigmoid-output conversion (1/256 steps), which floors every score below ~0.002: QAT trained against that grid while the releases shipped logits (2.2's N6-B0: simulated logits -15 where the float model has -9).

Source code in birdnet_stm32/training/qat.py

def build_qat_model(
    deployment_model: tf.keras.Model,
    activation_ranges: dict[str, tuple[float, float]],
    logit_output: bool = True,
) -> tf.keras.Model:
    """Build an activation-fake-quant graph sharing deployment model weights.

    A sigmoid head is simulated as the release converts it: INT8 logits, the
    sigmoid applied after (``convert --output_activation logit``, every release
    from 1.5). ``logit_output=False`` also puts the probabilities on the INT8
    output grid of a sigmoid-output conversion (1/256 steps), which floors every
    score below ~0.002: QAT trained against that grid while the releases shipped
    logits (2.2's N6-B0: simulated logits -15 where the float model has -9).
    """

    def clone_function(layer):
        if layer.__class__.__name__ != "AudioFrontendLayer":
            return layer
        clone = layer.__class__.from_config(layer.get_config())
        clone.build(tuple(layer.input.shape))
        clone.set_weights(layer.get_weights())
        freeze_batch_norm(clone)
        clone.set_quantization_hook(_FrontendQuantizationHook(activation_ranges))
        return clone

    def call_function(layer, *args, **kwargs):
        if _is_quantizable(layer):
            kernel_call = _QuantizedKernelCall(layer, name=f"{layer.name}_quantized_kernel")
            if _is_sigmoid_dense(layer):
                output = kernel_call(*args, linear=True, **kwargs)
                lo, hi = activation_ranges[f"{layer.name}__logits"]
                output = FakeQuantActivation(lo, hi, name=f"{layer.name}_logits_fake_quant")(output)
                output = layers.Activation("sigmoid")(output)
            else:
                output = kernel_call(*args, **kwargs)
        else:
            output = layer(*args, **kwargs)
        if layer.name in activation_ranges and not (logit_output and _is_sigmoid_dense(layer)):
            minimum, maximum = activation_ranges[layer.name]
            output = FakeQuantActivation(
                minimum,
                maximum,
                name=f"{layer.name}_fake_quant",
            )(output)
        return output

    inner_model = tf.keras.models.clone_model(
        deployment_model,
        clone_function=clone_function,
        call_function=call_function,
    )
    raw_inputs = tf.keras.Input(
        shape=deployment_model.input_shape[1:],
        dtype=deployment_model.input_dtype,
        name="qat_input",
    )
    minimum, maximum = activation_ranges["__input__"]
    quantized_inputs = FakeQuantActivation(minimum, maximum, name="input_fake_quant")(raw_inputs)
    return tf.keras.Model(raw_inputs, inner_model(quantized_inputs), name=f"{deployment_model.name}_qat")

refresh_activation_ranges(qat_student, ranges)

Move every fake quantizer to a new range; returns how many moved.

Source code in birdnet_stm32/training/qat.py

def refresh_activation_ranges(qat_student: tf.keras.Model, ranges: dict[str, tuple[float, float]]) -> int:
    """Move every fake quantizer to a new range; returns how many moved."""
    moved = 0
    for key, quantizer in _fake_quantizers(qat_student).items():
        if key in ranges:
            quantizer.set_range(*ranges[key])
            moved += 1
    return moved

sync_frontend_weights(qat_model, deployment_model)

Copy separately cloned custom-frontend weights into the clean model.

Source code in birdnet_stm32/training/qat.py

def sync_frontend_weights(qat_model: tf.keras.Model, deployment_model: tf.keras.Model) -> None:
    """Copy separately cloned custom-frontend weights into the clean model."""
    qat_frontends = {
        layer.name: layer
        for layer in all_layers(qat_model)
        if layer.__class__.__name__ == "AudioFrontendLayer" and getattr(layer, "_quantization_hook", None) is not None
    }
    deployment_frontends = {
        layer.name: layer for layer in all_layers(deployment_model) if layer.__class__.__name__ == "AudioFrontendLayer"
    }
    if qat_frontends.keys() != deployment_frontends.keys():
        raise ValueError("QAT and deployment frontend layers do not match")
    for name, frontend in qat_frontends.items():
        deployment_frontends[name].set_weights(frontend.get_weights())

freeze_batch_norm(model)

Freeze all BatchNormalization layers, including nested frontend BN.

Source code in birdnet_stm32/training/qat.py

def freeze_batch_norm(model: tf.keras.Model) -> int:
    """Freeze all BatchNormalization layers, including nested frontend BN."""
    batch_norms = [layer for layer in all_layers(model) if isinstance(layer, layers.BatchNormalization)]
    for layer in batch_norms:
        layer.trainable = False
    return len(batch_norms)

run_qat(args)

Fine-tune a pretrained model against weight and activation INT8 noise.

Source code in birdnet_stm32/training/qat.py

def run_qat(args: argparse.Namespace) -> None:
    """Fine-tune a pretrained model against weight and activation INT8 noise."""
    from birdnet_stm32.conversion.quantize import representative_data_gen, stratified_sample_paths
    from birdnet_stm32.data.dataset import (
        load_classes_file,
        load_file_paths_from_directory,
    )
    from birdnet_stm32.data.generator import load_dataset
    from birdnet_stm32.models.runners import load_keras_model
    from birdnet_stm32.training.config import ModelConfig
    from birdnet_stm32.training.trainer import train_model
    from birdnet_stm32.training.validation import VALIDATION_SUBSET_SEED, Int8Selection, stratified_validation_subset

    if not args.checkpoint_path.endswith(".keras"):
        raise ValueError("QAT checkpoint must end in .keras")
    qat_path = args.checkpoint_path.replace(".keras", "_qat.keras")
    if os.path.exists(qat_path) or os.path.exists(qat_path.replace(".keras", "_selection.json")):
        raise FileExistsError("QAT outputs already exist; use a new run directory")
    if args.resume:
        raise ValueError("QAT resume is unsupported: start from an explicit checkpoint in a new output directory")
    if not args.data_path_val:
        raise ValueError("QAT requires an explicit, disjoint --data_path_val")
    # Teacher targets and the teacher crop reach QAT's training loader as they do
    # float training's (opt-in); without them QAT trains on hard labels and energy
    # crops, and the knowledge float training added still reaches it through its
    # frozen teacher, the pre-QAT checkpoint itself.
    teacher_cache = getattr(args, "teacher_cache", None)
    teacher_weight = float(getattr(args, "teacher_weight", 0.0) or 0.0)
    crop_policy = getattr(args, "crop_policy", "energy")
    if (teacher_weight > 0 or crop_policy == "teacher") and not teacher_cache:
        raise ValueError("QAT --teacher_weight > 0 or --crop_policy teacher needs a --teacher_cache")
    if not os.path.isfile(args.checkpoint_path):
        raise FileNotFoundError(f"QAT requires a pretrained model: {args.checkpoint_path}")
    print(f"[QAT] Loading pretrained model from {args.checkpoint_path}")
    deployment_model = load_keras_model(args.checkpoint_path)
    teacher_model = load_keras_model(args.checkpoint_path)
    teacher_model.trainable = False

    cfg_path = getattr(args, "model_config", "") or os.path.splitext(args.checkpoint_path)[0] + "_model_config.json"
    if not os.path.isfile(cfg_path):
        raise FileNotFoundError(f"Model config not found: {cfg_path}")
    cfg = ModelConfig.load(cfg_path)

    classes = load_classes_file(args.classes_file) if args.classes_file else list(cfg.class_names)
    if not classes:
        raise ValueError("QAT requires class_names in the model config or --classes_file")
    if classes != cfg.class_names:
        raise ValueError("QAT class order must exactly match the pretrained model config")
    if len(classes) != deployment_model.output_shape[-1]:
        raise ValueError("QAT dataset class count does not match the pretrained model output")

    train_paths, train_classes = load_file_paths_from_directory(args.data_path_train, classes=classes)
    if args.data_path_val:
        val_paths, val_classes = load_file_paths_from_directory(args.data_path_val, classes=classes)
    else:
        rng = np.random.default_rng(args.seed)
        rng.shuffle(train_paths)
        split_idx = int(len(train_paths) * (1 - args.val_split))
        train_paths, val_paths = train_paths[:split_idx], train_paths[split_idx:]
    if not train_paths or not val_paths:
        raise ValueError("QAT requires non-empty training and validation datasets")
    if train_classes != classes:
        raise ValueError(
            f"QAT training data is missing configured classes: {[name for name in classes if name not in train_classes]}"
        )
    if args.data_path_val and val_classes != classes:
        raise ValueError(
            f"QAT validation data is missing configured classes: {[name for name in classes if name not in val_classes]}"
        )

    if set(map(os.path.realpath, train_paths)) & set(map(os.path.realpath, val_paths)):
        raise ValueError("QAT training and validation manifests overlap")

    # Conversion samples the physical training manifest, never the optionally
    # duplicated epoch-balancing list used by the trainer.
    calibration_source_paths = list(train_paths)

    common_kwargs = dict(
        sample_rate=cfg.sample_rate,
        max_duration=args.max_duration,
        chunk_duration=cfg.chunk_duration,
        spec_width=cfg.spec_width,
        mel_bins=cfg.num_mels,
        fft_length=cfg.fft_length,
        mag_scale=cfg.mag_scale,
        input_compression=cfg.input_compression,
        num_workers=args.num_workers,
        max_chunks_per_file=args.max_chunks_per_file,
        prefetch_batches=args.prefetch_batches,
    )
    teacher_kwargs: dict = {}
    if teacher_cache:
        from birdnet_stm32.data.teacher import TeacherTargets

        teacher = TeacherTargets(teacher_cache, classes)  # fail here rather than in every loader worker
        covered = sum(os.path.splitext(os.path.basename(p))[0] in teacher for p in train_paths)
        print(
            f"[QAT] Teacher targets: {covered}/{len(train_paths)} training files covered, "
            f"weight {teacher_weight}, crop policy {crop_policy}"
        )
        teacher_kwargs = dict(teacher_cache=teacher_cache, teacher_weight=teacher_weight, crop_policy=crop_policy)
    if getattr(args, "label_sidecar", None):
        from birdnet_stm32.data.sidecar import load_label_sidecar

        teacher_kwargs["label_sidecar"] = load_label_sidecar(args.label_sidecar, classes)
        print(f"[QAT] Label sidecar: {len(teacher_kwargs['label_sidecar'])} files with label additions")
    # The float run's draw: without the class cap QAT trains on the raw, unbalanced stream (and without the field
    # table on focal audio only), which moved 2.2's N6-B0 away from its float model each epoch (INT8 selection cMAP
    # 0.739 -> 0.729 -> 0.712).
    draw_kwargs: dict = {}
    if getattr(args, "class_cap_per_epoch", 0):
        draw_kwargs["class_cap_per_epoch"] = args.class_cap_per_epoch
        if getattr(args, "field_table", None):
            from birdnet_stm32.data.sidecar import load_field_table

            draw_kwargs.update(
                field_table=load_field_table(args.field_table),
                field_share=args.field_share,
                field_per_recording=args.field_per_recording,
            )
        print(
            f"[QAT] Class cap {args.class_cap_per_epoch} per epoch"
            + (f", field share {args.field_share}" if "field_table" in draw_kwargs else "")
        )
    train_dataset = load_dataset(
        train_paths,
        classes,
        audio_frontend=cfg.audio_frontend,
        batch_size=args.batch_size,
        mixup_alpha=0.0,
        mixup_probability=0.0,
        random_offset=True,
        snr_threshold=0.1,
        spec_augment=False,
        **teacher_kwargs,
        **draw_kwargs,
        **common_kwargs,
    )
    # Selection converts and scores an INT8 model every epoch, twice over the
    # manifest. A fixed stratified subset keeps that affordable; it must be the
    # same draw for every arm and epoch or the comparison means nothing, so it
    # is seeded and its hash goes into the selection report.
    subset = int(getattr(args, "validation_subset", 0) or 0)
    selection_paths = stratified_validation_subset(val_paths, subset)
    if len(selection_paths) < len(val_paths):
        print(
            f"[QAT] Selecting on a fixed stratified subset of {len(selection_paths)} "
            f"of {len(val_paths)} validation files (seed {VALIDATION_SUBSET_SEED})"
        )
    elif subset:
        print(f"[QAT] Requested subset {subset} >= {len(val_paths)} validation files; using all")
    # Keras' per-epoch validation pass reads the selection files, not the whole
    # manifest: its val_loss selects nothing and the full pass re-decoded every
    # validation file each epoch.
    val_dataset = load_dataset(
        selection_paths,
        classes,
        audio_frontend=cfg.audio_frontend,
        batch_size=args.batch_size,
        mixup_alpha=0.0,
        mixup_probability=0.0,
        random_offset=False,
        snr_threshold=0.0,
        spec_augment=False,
        **common_kwargs,
    )

    n_frozen = freeze_batch_norm(deployment_model)
    print(f"[QAT] Frozen {n_frozen} BatchNorm layers")

    calibration_count = int(args.qat_calibration_samples)
    calibration_paths = stratified_sample_paths(calibration_source_paths, calibration_count, seed=42)
    if len(calibration_paths) != calibration_count:
        raise ValueError(
            f"QAT requested {calibration_count} calibration paths but only {len(calibration_paths)} are available"
        )
    calibration_data = list(
        representative_data_gen(
            calibration_paths,
            cfg.to_dict(),
            num_samples=calibration_count,
        )
    )
    if len(calibration_data) != calibration_count:
        raise ValueError("Calibration skipped files; refusing to train on an incomplete manifest")
    activation_ranges = calibrate_activation_ranges(
        deployment_model,
        calibration_data,
        max_samples=calibration_count,
        percentile=range_percentile,
    )
    print(f"[QAT] Activation ranges use the converter's exact stratified {calibration_count}-sample manifest (seed=42)")
    extra_callbacks: list[tf.keras.callbacks.Callback] = []

    loss_weights = {
        "distillation_weight": float(args.qat_distillation_weight),
        "cosine_weight": float(args.qat_cosine_weight),
        "cosine_tail_weight": float(args.qat_cosine_tail_weight),
        "cosine_tail_fraction": float(args.qat_cosine_tail_fraction),
    }
    validate_loss_weights(loss_weights)
    qat_student = build_qat_model(deployment_model, activation_ranges)
    qat_model = _DistilledQATModel(qat_student, teacher_model, **loss_weights)
    print(
        "[QAT] Enabled frozen-teacher consistency losses "
        f"(Bernoulli KL weight={loss_weights['distillation_weight']:.3g}, "
        f"cosine mean weight={loss_weights['cosine_weight']:.3g}, "
        f"worst {loss_weights['cosine_tail_fraction']:.1%} cosine "
        f"weight={loss_weights['cosine_tail_weight']:.3g})"
    )

    qat_path = args.checkpoint_path.replace(".keras", "_qat.keras")
    ranges_path = qat_path.replace(".keras", "_activation_ranges.json")
    with open(ranges_path, "w", encoding="utf-8") as handle:
        json.dump({name: {"min": lo, "max": hi} for name, (lo, hi) in activation_ranges.items()}, handle, indent=2)

    cfg.save(qat_path.replace(".keras", "_model_config.json"))
    with open(qat_path.replace(".keras", "_labels.txt"), "w", encoding="utf-8") as handle:
        handle.write("\n".join(classes) + "\n")
    if getattr(args, "qat_range_refresh", False):
        # Before the selector, so the simulated score it logs uses the new grid.
        extra_callbacks.append(
            RangeRefresh(
                qat_student,
                deployment_model,
                calibration_data,
                sync=lambda: sync_frontend_weights(qat_model, deployment_model),
                samples=calibration_count,
            )
        )
    selector = Int8Selection(
        deployment_model,
        teacher_model,
        calibration_data,
        qat_path,
        sync=lambda: sync_frontend_weights(qat_model, deployment_model),
        simulated=qat_student,
        files=selection_paths,
        classes=classes,
        cfg=cfg.to_dict(),
        overlap=args.validation_overlap,
        pooling=args.validation_pooling,
        batch_size=args.batch_size,
    )
    extra_callbacks.append(selector)
    steps_per_epoch = getattr(args, "steps_per_epoch", 0) or max(
        1, math.ceil(len(train_paths) / float(args.batch_size))
    )
    val_steps = max(1, math.ceil(len(selection_paths) / float(args.batch_size)))
    print(f"[QAT] Training on {len(train_paths)} files, validating on {len(val_paths)} files")
    print(f"[QAT] Fine-tuning for {args.epochs} epochs at LR={args.learning_rate}")
    train_model(
        qat_model,
        train_dataset,
        val_dataset,
        epochs=args.epochs,
        learning_rate=args.learning_rate,
        batch_size=args.batch_size,
        checkpoint_path=qat_path,
        steps_per_epoch=steps_per_epoch,
        val_steps=val_steps,
        optimizer=args.optimizer,
        weight_decay=args.weight_decay,
        loss_fn=_detect_loss(deployment_model),
        gradient_clip_norm=args.grad_clip,
        checkpoint_model=deployment_model,
        checkpoint_sync=lambda: sync_frontend_weights(qat_model, deployment_model),
        checkpoint_monitor="val_int8_cmap",
        checkpoint_mode="max",
        checkpoint_managed=True,
        extra_callbacks=extra_callbacks,
    )
    print(f"[QAT] Best converted INT8 cMAP: {selector.best:.6f}; checkpoint: {qat_path}")
    print(f"[QAT] Selected development TFLite: {selector.int8_path}")
    print(f"[QAT] Activation calibration ranges saved to {ranges_path}")