Skip to content

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 model
  • my_model_model_config.json — conversion metadata (frontend, shapes, etc.)
  • my_model_labels.txt — ordered class names
  • my_model_history.csv — per-epoch training metrics (loss, ROC-AUC)
  • my_model_curves.png — loss and exact validation cMAP curves plot
  • my_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_maxmean averages 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) or 5. Stage 1 stays 3 × 3.
  • --stage_widths: base output channels of the four stages before --alpha, default 32 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.