Training¶
Basic usage¶
python -m birdnet_stm32 train \
--data_path_train data/train \
--checkpoint_path checkpoints/my_model.keras
The defaults are the release recipe: raw frontend, pwl magnitude scaling,
24 kHz, 2.5 s chunks, 512-d embedding, 50 epochs at learning rate 5e-4.
For a leakage-safe precomputed validation split and stable output order, pass a separate validation root and a one-label-per-line classes file:
python -m birdnet_stm32 train \
--data_path_train data/train \
--data_path_val data/validation \
--classes_file data/labels.txt
The class file controls model-output order. Folders named noise, silence,
background, or other are still loaded as all-zero examples and must not be
listed as outputs. When --data_path_val is present, --val_split is ignored.
The script saves these files alongside the checkpoint:
my_model.keras— trained Keras modelmy_model_model_config.json— conversion metadata (frontend, shapes, etc.)my_model_labels.txt— ordered class namesmy_model_history.csv— per-epoch training metrics (loss, ROC-AUC)my_model_curves.png— loss and exact validation cMAP curves plotmy_model_train_state.json— epoch counter for--resume
Audio frontends¶
| Frontend | Input to model | Description |
|---|---|---|
raw (default) |
Peak-normalized waveform | Model applies a mel-seeded, trainable Gabor quadrature filterbank. The release frontend: the whole pipeline runs on the NPU. |
hybrid |
Linear magnitude STFT | Model applies a learned mel mixer and magnitude scaling. The STFT runs outside the model (host, or the Cortex-M55). With --input_compression sqrt the best INT8 accuracy measured, at 148 ms per file against 75 ms for raw (1.4 models). |
librosa |
Mel spectrogram | Mel spectrogram computed outside the model. Smallest quantization loss, weakest float model. |
hybrid and librosa inputs are defined by birdnet_stm32.audio.stft and
specified in Spectrogram Input.
These three are the only frontends. mfcc and log_mel were removed in 1.2.0:
both were host-precomputed variants of the same librosa path, and no release
ever used them.
Raw frontend memory limit
The raw input must contain fewer than 65,536 samples. At 24 kHz, use a chunk no longer than about 2.7 seconds; 2.5 seconds is the tested setting.
Magnitude scaling¶
| Mode | Description | Quantization friendliness |
|---|---|---|
pwl (default) |
Piecewise-linear learned compression | Excellent — recommended for deployment |
none |
No compression | Baseline only, for ablations |
pcen and db were removed in 1.2.0. dB's log op produces exactly the wide
dynamic range INT8 cannot hold — the failure this frontend exists to avoid —
and PCEN was never used by a release. cpwl, a compressive variant of pwl,
was removed in 1.3.0: best float of any raw model, worst INT8.
Model architecture¶
The DS-CNN is scaled with two knobs:
--alpha(width multiplier): scales channel counts across all stages. Default 1.0. Values like 0.5 or 0.75 produce smaller models.--depth_multiplier: repeats each depthwise-separable block. Default 1. Increase to 2 for deeper models.
Two experimental options change the backbone shape. Both default to the release architecture, and neither has been validated on INT8 or on the board yet:
--head_pooling:gap(default) averages the final feature map over frequency and time.freq_mean_time_maxmeanaverages over frequency, then adds the max and the mean over time, so a short call is not averaged away. It adds no weights.--dw_kernel_size: depthwise kernel size in stages 2–4,3(default) or5. Stage 1 stays 3 × 3.--stage_widths: base output channels of the four stages before--alpha, default32 64 128 256. Use it to widen only the late stages, where the feature map is small and extra channels are cheap in activation memory. When the last stage is as wide as--embeddings_size, the separate 1 × 1 embedding convolution is left out.
Channel alignment
Keep channel counts as multiples of 8 for optimal NPU vectorization. The
model builder enforces this automatically via _make_divisible.
Training options¶
Data augmentation¶
- Mixup: controlled by
--mixup_alpha(default 0.2, 0 disables) and--mixup_probability(default 0.25). Uses Dirichlet multi-source mixing (2–3 sources per sample) to realistically emulate overlapping bird vocalizations. Labels are combined via element-wise max. - SpecAugment: enabled by default. Applies random frequency and time
masking to spectrograms during training. Disable with
--no_spec_augment. Control mask widths with--freq_mask_max(default 8 bins) and--time_mask_max(default 25 frames). - Smart crop: long recordings (> 2 chunks) are automatically cropped to salient regions using short-time energy (STE) analysis, reducing label noise from silent or irrelevant segments.
- Multi-chunk I/O reuse: long files (e.g. 60 s recordings) yield up to
--max_chunks_per_file(default 3) salient chunks per file open, stored in a memory-bounded shuffled reservoir. This avoids redundant FLAC decode + resample for the same file across epochs.
Loss function¶
The classifier head is always sigmoid + binary crossentropy. Soundscape recordings are inherently multi-label, so we always optimise per-class probabilities even when the source label is single-class.
Optimizer¶
Select with --optimizer (default adam):
| Optimizer | Description |
|---|---|
adam |
Adaptive moment estimation (default) |
sgd |
SGD with momentum 0.9 |
adamw |
AdamW with decoupled weight decay |
Set weight decay with --weight_decay (default 0, only used by adamw).
Deterministic mode¶
Training is always deterministic — all random seeds (Python, NumPy,
TensorFlow) are set and TF_DETERMINISTIC_OPS is enabled automatically.
Use --seed (default 42) to change the RNG seed.
Gradient clipping¶
Gradient clipping by global norm is enabled by default (--grad_clip 1.0).
Set to 0 to disable. Prevents exploding gradients, especially useful with
large models or unstable training.
Mixed precision¶
Use --mixed_precision to enable FP16 compute with FP32 accumulation.
Reduces memory usage and speeds up training on GPUs with Tensor Cores.
Resumable training¶
Use --resume to continue training from a previously saved checkpoint.
The optimizer state is recompiled and training resumes from the last saved
epoch. Example:
# Initial training (interrupted or completed at epoch 30)
python -m birdnet_stm32 train --epochs 30 --checkpoint_path ckpt/model.keras ...
# Resume and extend to 50 epochs
python -m birdnet_stm32 train --epochs 50 --resume --checkpoint_path ckpt/model.keras ...
Quantization-Aware Training (QAT)¶
Use --qat to fine-tune a pretrained model with simulated INT8 quantization
noise. This closes the accuracy gap between the float Keras model and the
quantized TFLite model by teaching the weights to survive quantization.
QAT requires a pretrained model
Always train normally first, then fine-tune with --qat. Do not use
--qat from scratch — the quantization noise destabilizes randomly
initialized weights and the model will not converge. The dataset must
have the same classes as the pretrained model; use --linear_probe to
adapt to a different class set first.
QAT calibrates activation ranges on the converter's exact deterministic, class-stratified training manifest (1,024 samples by default), then trains with per-channel INT8 kernel grids and per-tensor INT8 activation grids. The simulation includes the quantized waveform input and the kernels and elementwise boundaries inside the custom raw frontend. BatchNorm layers are frozen. Standard variables are shared with a clean deployment graph; cloned frontend variables are synchronized before each checkpoint. Only the clean graph is saved, so no FakeQuant ops remain in the model.
A frozen copy of the untouched checkpoint acts as the teacher. QAT minimizes the normal label loss, per-output Bernoulli KL divergence, and both mean and worst-sample per-sample cosine distance from that teacher. The defaults apply the tail loss to the worst 10% of each batch with 0.75 weight; both values are configurable. This constrains background and low-confidence probabilities while directly optimizing the lower tail that the release parity gate measures.
With --qat, the defaults switch to the fine-tuning schedule: 8 epochs at
learning rate 2e-5, with activation ranges recalibrated on the current weights
after every epoch (--qat_range_refresh, on by default). The simulation
quantizes every tensor the converter quantizes, including each partial
filterbank convolution, so its validation cMAP (val_sim_int8_cmap) tracks
the converted model's (val_int8_cmap) to within a few thousandths.
The full raw pipeline:
# Step 1: Normal training
python -m birdnet_stm32 train --data_path_train data/train \
--data_path_val data/validation --classes_file data/labels.txt \
--checkpoint_path checkpoints/model.keras
# Step 2: Equalize the frontend's per-band ranges (raw only; exact in float)
python -m birdnet_stm32 equalize --checkpoint_path checkpoints/model.keras \
--data_path_train data/train --output_path checkpoints/model_eq.keras
# Step 3: QAT fine-tuning
python -m birdnet_stm32 train --data_path_train data/train \
--data_path_val data/validation --classes_file data/labels.txt --qat \
--checkpoint_path checkpoints/model_eq.keras
# Step 4: Convert the QAT model
python -m birdnet_stm32 convert \
--checkpoint_path checkpoints/model_eq_qat.keras \
--model_config checkpoints/model_eq_model_config.json \
--data_path_train data/train
The QAT model is saved as {name}_qat.keras alongside the original.
equalize rescales each band of the raw filterbank and PWL so that every band
gets the same share of the tensors' INT8 grids, and refuses to save if the
float output changes. See INT8 quality for the
measurements.
For hybrid and librosa with --input_compression, skip QAT: it scored
below plain post-training quantization at every epoch. Convert the trained
checkpoint directly.
Linear probing¶
Use --linear_probe to freeze a pretrained backbone and train only a new
classification head on your custom species dataset. This is useful when you
have a pretrained model (e.g. a large BirdNET checkpoint) and want to adapt
it to a different set of species with limited data.
python -m birdnet_stm32 train --data_path_train data/my_species \
--linear_probe --checkpoint_path checkpoints/pretrained.keras \
--data_path_val data/my_species_val \
--classes_file data/my_species/labels.txt \
--epochs 20 --learning_rate 0.001
The probe model is saved as {name}_probe.keras with a new labels file and
{name}_probe_model_config.json.
Pass --classes_file whenever the head will be shipped. The head's output order
is its labels file, and without an explicit schema that order comes from a
directory listing, which is not a contract. --data_path_val gives the probe a
fixed validation root instead of a random slice of the training set, so repeated
runs are comparable.
Probing a checkpoint from an earlier compression step needs --model_config:
a QAT checkpoint writes no config of its own, since its architecture is the base
model's.
python -m birdnet_stm32 train --data_path_train data/my_species \
--linear_probe --checkpoint_path checkpoints/pretrained_qat.keras \
--model_config checkpoints/pretrained_model_config.json \
--classes_file data/my_species/labels.txt --epochs 20
To ship the resulting head on its own, convert it against the backbone already on the device — see Updating the head against a flashed backbone.
Learning rate¶
A two-epoch linear warmup reaches --learning_rate (default 0.001), followed
by cosine decay to near-zero over --epochs (default 50). Best-checkpoint
selection and early stopping maximize exact validation cMAP. Standard CLI
training evaluates files with the configured overlap and pooling. QAT selects
on converted INT8 file cMAP, including epoch zero, and saves the exact
TFLite artifact that was scored alongside its matching Keras checkpoint.
The chunk PR-AUC metric is logged as pr_auc and does not select checkpoints.
Full argument reference¶
| Argument | Default | Description |
|---|---|---|
--data_path_train |
(required) | Path to training data |
--data_path_val |
None | Separate validation root; disables random validation splitting |
--classes_file |
None | Ordered one-label-per-line output schema |
--max_classes |
None | Use only the N most populated classes |
--max_samples |
None | Max files per class |
--upsample_ratio |
0.5 | Minority class upsample ratio |
--sample_rate |
24000 | Audio sample rate (Hz) |
--num_mels |
64 | Number of mel frequency bins |
--spec_width |
256 | Spectrogram width (frames) |
--fft_length |
512 | FFT window length |
--chunk_duration |
2.5 | Chunk duration (seconds) |
--max_duration |
60 | Max seconds to load per file |
--audio_frontend |
raw | raw, hybrid, or librosa — raw models trained before the NPU fixes must be retrained |
--mag_scale |
pwl | pwl or none |
--input_compression |
none | none, sqrt or log: compress a librosa/hybrid spectrogram before its first INT8 quantization; the firmware applies the same compression |
--embeddings_size |
512 | Embedding channels before head |
--alpha |
1.0 | Model width scaling |
--depth_multiplier |
1 | Block repeats per stage |
--head_pooling |
gap | gap or freq_mean_time_maxmean (experimental) |
--dw_kernel_size |
3 | Depthwise kernel in stages 2–4: 3 or 5 (experimental) |
--stage_widths |
32 64 128 256 | Base channels of the four stages, before --alpha (experimental) |
--frontend_trainable |
False | Make frontend weights trainable |
--mixup_alpha |
0.2 | Mixup alpha (0 disables) |
--mixup_probability |
0.25 | Fraction of batch to mix |
--no_spec_augment |
False | Disable SpecAugment masking (on by default) |
--freq_mask_max |
8 | Max frequency mask width (bins) |
--time_mask_max |
25 | Max time mask width (frames) |
--dropout |
0.5 | Dropout rate before classifier head |
--optimizer |
adam | adam, sgd, or adamw |
--weight_decay |
0.0 | Weight decay (adamw only) |
--grad_clip |
1.0 | Max gradient norm for clipping (0 = disabled) |
--mixed_precision |
False | Enable FP16 mixed precision training |
--resume |
False | Resume training from checkpoint |
--seed |
42 | Random seed |
--batch_size |
32 | Batch size |
--num_workers |
8 | Parallel data loading workers (0 = sequential) |
--max_chunks_per_file |
1 | Max salient chunks per file open (reduces redundant I/O) |
--prefetch_batches |
2 | Loader prefetch depth in batches |
--epochs |
50 (8 with --qat) |
Number of epochs |
--learning_rate |
5e-4 (2e-5 with --qat, 1e-3 with --linear_probe) |
Initial learning rate |
--val_split |
0.2 | Validation split fraction when --data_path_val is not supplied |
--checkpoint_path |
checkpoints/best_model.keras | Output path (.keras) |
--qat |
False | Quantization-aware fine-tuning |
--qat_calibration_samples |
1024 | Exact stratified samples used for QAT and conversion calibration |
--qat_distillation_weight |
1.0 | Frozen-teacher Bernoulli-KL weight |
--qat_cosine_weight |
0.10 | Mean teacher/student cosine-loss weight |
--qat_cosine_tail_weight |
0.75 | Worst-sample cosine-loss weight |
--qat_cosine_tail_fraction |
0.10 | Fraction of each batch included in the worst-sample loss |
--qat_range_refresh |
on | Recalibrate activation ranges on the current weights after every epoch, so QAT trains against the grid conversion will use; --no-qat_range_refresh disables it |
--linear_probe |
False | Freeze backbone and train only classifier head |
--model_config |
(inferred) | Architecture config for --qat, --linear_probe; required when the checkpoint has no sibling config |
--validation_overlap |
half the chunk duration | File validation overlap in seconds |
--validation_pooling |
max | File validation pooling |
--validation_subset |
0 | Score checkpoint selection (training and QAT) on a fixed class-balanced draw of N validation files, seed 1234 (0 = all). Biased upward; not comparable to full-manifest cMAP |
Data pipeline¶
The training pipeline uses a multiprocessing pool for parallel data loading, bypassing the GIL so FLAC decode, resampling, smart-crop, and spectrogram computation run across separate CPU cores.
When --max_chunks_per_file is greater than 1 (default 3), each file open
extracts multiple salient chunks which are buffered in a shuffled in-memory
reservoir sized from the sample representation and the loader's memory
budget. This dramatically reduces I/O for
long recordings: a 60 s file decoded once yields 3 usable chunks instead of
re-opening the same file 3 times across epochs.
The reservoir maintains batch diversity by shuffling samples from many different files before yielding them.
Training also checks host-available memory every 25 batches. It aborts before available RAM falls below an adaptive reserve (20% of host RAM, bounded between 2 and 12 GiB), leaving the last completed-epoch checkpoint available for a safer restart. This protects the host from multiprocessing and TensorFlow memory spikes; it is not a substitute for choosing a conservative worker count and batch size.
Tune with:
--num_workers N— number of worker processes (default 8, 0 = sequential)--max_chunks_per_file N— chunks per file open (default 3, 1 = original behavior)--prefetch_batches N— queued batches (default 2; higher uses more RAM)
Noise classes¶
Place audio in folders named noise, silence, background, or other
under data/train/. These receive all-zero label vectors and help the model
learn to reject non-bird sounds.