[3/3] audioprotopnet: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED models by nkundiushuti · Pull Request #184 · earthspecies/avex

nkundiushuti · 2026-05-14T14:16:15Z

Add ConvNeXt, AudioProtoPNet, AudioProtoPNet SED, prototype-based
classification, and two-stage training

New models

ConvNeXt (avex/models/convnext.py)

Plain ConvNextForImageClassification audio classifier (1-channel mel
spectrogram input)
Architecture and mel parameters loaded from packaged convnext_base.yml via
ConvNextCfg — no hardcoded constants
backbone property exposes model.convnext (encoder only) for correct
freeze/unfreeze during two-stage training
Hook-based embedding extraction from the 4 encoder stages, compatible with
all existing probes
Transparent PyTorch-Lightning checkpoint loading (strips model.model. /
model._orig_mod.model. prefixes)

AudioProtoPNet (avex/models/audioprotopnet.py)

ConvNeXt backbone + prototype classification head, loaded from HuggingFace
(DBD-research-group/AudioProtoPNet-{1,5,10,20}-BirdSet-XCL)
backbone property exposes model.model.backbone for correct freeze targeting
Exposes ebird_codes (raw eBird codes) and label_mapping (common names via
bundled eBird taxonomy)
Calls ModelBase.init directly to skip the ConvNeXt mel pipeline

AudioProtoPNet SED (avex/models/audioprotopnet_sed.py)

ConvNeXt-Base backbone + 2× bilinear add-on layers + spatial cosine
prototype head for sound event detection
forward() — clip-level logits [B, C] via global H×W max-pool
forward_frames() — per-frame sigmoid probabilities [B, T, C] via frequency
max-pool only; temporal resolution ≈7.6 Hz at 32 kHz
Checkpoint loaded from
gs://representation-learning/models/sed/audioprotopnet-20/ via
universal_torch_load (cached locally; same pattern as BEATs)
LinearLayerWithoutNegativeConnections: block-diagonal per-class
non-negative linear layer; weights initialised with kaiming_uniform
Mel preprocessing via ConvNextCfg; backbone architecture from packaged
convnext_base.yml — no HuggingFace network calls beyond the checkpoint itself

Config

ConvNextCfg (avex/configs.py)

New Pydantic model for mel preprocessing parameters (sample_rate, n_fft,
hop_length, n_mels, norm_mean, norm_std); defaults match BirdSet XCL training
setup
Consistent with BEATsConfig pattern; packaged YAMLs carry a top-level
convnext_cfg block
convnext_cfg and checkpoint_dir added to ModelSpec and
_add_model_spec_params so user YAMLs flow through the factory correctly

eBird taxonomy (avex/data/)

Bundled ebird_taxonomy.json — 18,130 taxa (eBird v2021 + v2025), superset of
all 9,736 BirdSet XCL species
load() with @lru_cache for zero-cost repeated access

New losses (avex/training/losses.py)

AsymmetricLoss (Ridnik et al., ICCV 2021) — separate focusing parameters for
positives/negatives, probability margin clipping; registered as "asymmetric"
/ "asl" in build_criterion
OrthogonalLoss — encourages prototype vectors within the same class to be
orthogonal (‖P_c Pᵀ_c − I‖_F / numel)

Prototypical probe (avex/models/probes/prototypical_probe.py)

New "prototypical" probe type alongside linear, mlp, attention, lstm,
transformer
Architecture: cosine-similarity prototype layer → [0, 1] shift →
non-negative linear classifier
num_prototypes_per_class added to ProbeConfig
New "prototypical" probe type alongside linear, mlp, attention, lstm, transformer
Architecture: cosine-similarity prototype layer → [0, 1] shift → non-negative linear classifier
num_prototypes_per_class added to ProbeConfig

Example

examples/audioprotopnet_species_detection.py — clip-level and frame-level species detection on a directory of audio
files; --demo mode generates synthetic audio with no real recordings needed

Tests

24 unit tests for ConvNeXt and AudioProtoPNet (architecture, label mapping, embedding extraction, hook
registration, layer discovery)
29 unit tests for AudioProtoPNet SED (forward shapes, frame-level properties, batch independence, cosine
activation, _LinearLayerWithoutNegativeConnections, registry)

…al losses, phase 2 in training with prototypical heads

Custom feature extractor returns Tensor, not BatchFeature dict. **-unpacking a Tensor raises TypeError. Inherited ConvNext forward passes process_audio result positionally — correct.

- Replace sklearn train_test_split with esp_data PandasBackend in TrainValSplitTransform - Add worker_timeout param to build_dataloaders - Add birdset_train_splits module (registers birdset_train dataset) - Split ebird_taxonomy.json into v2021/v2025; load() now requires version arg - Add fill_labels_from_answer transform to beans benchmark configs - Rename birdset -> birdset_train in birdset benchmark configs - Add slurm benchmark config variants for beans

- Add _write_embedding_metadata helper; store aggregation in HDF5 attrs - HDF5EmbeddingDataset exposes embedding_aggregation attribute - _embedding_cache_matches validates aggregation before reuse - embedding_manager: use uncompressed size (shape*dtype) for memory routing - __getstate__ strips _window_cache to avoid N_workers x cache duplication - Window loads lazily on first __getitem__ (not prefilled in __init__) - _load_window scales capacity by num_workers to stay within budget

- Use _embedding_cache_matches for cache validation (aggregation-aware) - Fix operator precedence bugs in need_recompute_embeddings_*_clustering - Configurable offline probe embedding aggregation via probe_storage_aggregation - Hoist memory_limit_bytes before probing + retrieval/clustering branches - Graceful per-file fallback when embeddings missing (avoids base_model=None crash) - Release probe HDF5 caches before clustering to free ~46 GB - Reuse probe embeddings for retrieval/clustering when aggregation matches - Pass worker_timeout=120 to build_dataloaders

- Wire probe_num_workers (default 0) into all embedding DataLoaders - Gate pin_memory on probe_workers > 0 (pointless with no workers) - Fix multi-label target shape: one_hot 1D/squeezed-2D before BCEWithLogitsLoss

- Fix forward pass and feature extraction bugs - Remove example script (superseded by checkpoint configs) - Add beans evaluation config and slurm job

- Fix SED forward pass, feature extraction, and checkpoint loading bugs - Update audioprotopnet_sed checkpoint config - Add beans/birdset evaluation configs and slurm jobs

- BaseProbe2D._to_2d centralises 2D/3D/4D → 2D reduction logic - _pool4d pools W/time dim of 4-D tensors (mean or max) - input_processing param ('pooled'/'flatten'/'sequence') threads through all probe types - Remove duplicated reshape branches across probes

- Export ESP_DATA_HOME=gs://esp-ml-datasets in all existing job scripts - Add setup_slurm_cuda124_venv.sh and sync_to_slurm.sh

nkundiushuti · 2026-05-29T14:56:16Z

This PR has been split into three parts to make review easier. #184 must be merged last.

#	PR	Issue	Scope	Status
1	#186	#188	Data layer: taxonomy versioning, birdset splits, esp_data transforms	must merge first
2	#187	#189	Eval pipeline: embedding cache, run_evaluate fixes, finetune fixes	depends on #186
3	#184	—	Models: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED	depends on #186 and #187

Once #186 and #187 are merged into main, this branch will be rebased onto main and the diff will narrow to only the model changes.

- transforms: add DataFrame path using iloc/unique for pandas compatibility; backend API path retained for esp_data Polars/other backends - pre-commit: narrow large-file and codespell excludes to ebird_taxonomy_v*.json; remove ruff/ruff-format avex/data/ exclude (ruff only targets Python anyway) - ebird_taxonomy: remove unused logger; fix FileNotFoundError docstring indent - birdset_train_splits: add missing Returns/Raises/Yields docstring sections - dataset: document worker_timeout param in build_dataloaders docstring - test_ebird_taxonomy: annotate version param as EbirdTaxonomyVersion

- embedding_utils: mark cache incomplete when batches were skipped; skipped rows leave unwritten gaps in pre-allocated HDF5 datasets - run_evaluate: _reuse_probe_embeddings_for_eval rejects aggregation=none (3D embeddings can't be fed into retrieval/clustering) - run_evaluate: replace path.exists() with _embedding_cache_matches for clustering cache; use pooled aggregation (mean/max) for the check

- tests: add heavy regression coverage for the BaseProbe3D 2-D-embedding semantic change (LSTM/attention/transformer end-to-end + shape primitives); locks in 'single timestep, full feature dim' over the old 'F scalar timesteps' - prototypical_utils: derive the dim-probing dummy length from the backbone audio config instead of hardcoding 32 kHz x 5 s - prototypical_utils: document that virtual 'last_layer' only exists on AudioProtoPNet(SED); ConvNeXt/EfficientNet fall back to base discovery - audioprotopnet: add trust_remote_code security warning and a revision param threaded through all HF load calls for pinning - audioprotopnet_sed: allocate the empty prototype_vectors placeholder only on the birdcode path (split path constructs it directly) - configs: include 'asl' in the multilabel loss validation message

…chable CI runs 'pytest tests/unittests' without -m 'not integration', so the @pytest.mark.integration SED test that pulls a checkpoint from a private GCS bucket ran anonymously and failed with a 401. Skip cleanly on any remote/auth/IO failure so the test only runs where the asset is reachable.

…ration test

nkundiushuti · 2026-06-01T12:49:41Z

@benjaminsshoffman I removed the stage 2 part in run_train (as you suggested) and I left only the models and the evaluation part.

nkundiushuti · 2026-06-04T06:06:41Z

@benjaminsshoffman let me know, @GaganNarula do you want to take a look as well? I just added these models to be able to benchmark them also the probe head, no training happens

convnext, audioprotopnet, prototypical head, asymetrical and orthogon…

8a8be0a

…al losses, phase 2 in training with prototypical heads

nkundiushuti requested a review from david-rx as a code owner May 14, 2026 14:16

nkundiushuti added 2 commits May 14, 2026 16:38

fix CI issues

735c5b2

fix CI

55d13a1

nkundiushuti self-assigned this May 14, 2026

nkundiushuti added 3 commits May 14, 2026 18:59

bug fixes

2b53434

fix: remove forward override in AudioProtoPNet

3b47cbd

Custom feature extractor returns Tensor, not BatchFeature dict. **-unpacking a Tensor raises TypeError. Inherited ConvNext forward passes process_audio result positionally — correct.

audioprotopnet_sed

3b8c9ae

nkundiushuti requested a review from benjaminsshoffman May 15, 2026 11:19

benjaminsshoffman requested changes May 15, 2026

View reviewed changes

Comment thread avex/data/ebird_taxonomy.py Outdated

Comment thread examples/audioprotopnet_species_detection.py Outdated

Comment thread avex/training/train.py Outdated

nkundiushuti added 9 commits May 29, 2026 16:46

eval: finetune probe_num_workers + multi-label target fix

2ce0de2

- Wire probe_num_workers (default 0) into all embedding DataLoaders - Gate pin_memory on probe_workers > 0 (pointless with no workers) - Fix multi-label target shape: one_hot 1D/squeezed-2D before BCEWithLogitsLoss

eval: add probe_storage_aggregation and probe_num_workers to config

b3c4f0a

models: audioprotopnet bug fixes and eval config

2daed32

- Fix forward pass and feature extraction bugs - Remove example script (superseded by checkpoint configs) - Add beans evaluation config and slurm job

models: audioprotopnet_sed bug fixes and eval configs

e77dda8

- Fix SED forward pass, feature extraction, and checkpoint loading bugs - Update audioprotopnet_sed checkpoint config - Add beans/birdset evaluation configs and slurm jobs

jobs: add ESP_DATA_HOME and slurm utility scripts

21c3cdc

- Export ESP_DATA_HOME=gs://esp-ml-datasets in all existing job scripts - Add setup_slurm_cuda124_venv.sh and sync_to_slurm.sh

This was referenced May 29, 2026

[1/3] audioprotopnet: data layer updates #186

Merged

[2/3] audioprotopnet: eval pipeline fixes #187

Merged

nkundiushuti changed the title ~~audioprotopnet models, probe heads, losses, training phase~~ [3/3] audioprotopnet: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED models May 29, 2026

This was referenced May 29, 2026

[1/3] audioprotopnet: data layer updates #188

Closed

[2/3] audioprotopnet: eval pipeline fixes #189

Closed

nkundiushuti added 5 commits May 29, 2026 17:55

fix: skip test_embedding_metadata when esp_data not installed

9c57f53

nkundiushuti added 2 commits June 1, 2026 14:11

Merge remote-tracking branch 'origin/main' into marius/audioprotopnet

bfb4014

feat: wire orthogonal loss + clamp into prototypical probe, add integ…

96972ee

…ration test

nkundiushuti mentioned this pull request Jun 30, 2026

add audioprotopnet models #195

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[3/3] audioprotopnet: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED models - #184

[3/3] audioprotopnet: ConvNeXt + AudioProtoPNet + AudioProtoPNetSED models#184
nkundiushuti wants to merge 22 commits into
mainfrom
marius/audioprotopnet

nkundiushuti commented May 14, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nkundiushuti commented May 29, 2026

Uh oh!

nkundiushuti commented Jun 1, 2026

Uh oh!

nkundiushuti commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

nkundiushuti commented May 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

nkundiushuti commented May 29, 2026

Uh oh!

nkundiushuti commented Jun 1, 2026

Uh oh!

nkundiushuti commented Jun 4, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

nkundiushuti commented May 14, 2026 •

edited

Loading