From 13e6f97fa71492c6500dccda7db910c8b1daa63d Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:45:29 +0100 Subject: [PATCH 01/10] Add setting for compiling the model before train --- src/batdetect2/models/__init__.py | 4 ++++ src/batdetect2/train/config.py | 3 +++ src/batdetect2/train/train.py | 12 ++++++++++++ 3 files changed, 19 insertions(+) diff --git a/src/batdetect2/models/__init__.py b/src/batdetect2/models/__init__.py index ee96d93..8564d05 100644 --- a/src/batdetect2/models/__init__.py +++ b/src/batdetect2/models/__init__.py @@ -112,6 +112,9 @@ class ModelConfig(BaseConfig): Attributes ---------- + compile : bool + If ``True``, compile the model before training. Defaults to + ``False``. samplerate : int Expected input audio sample rate in Hz. Audio must be resampled to this rate before being passed to the model. Defaults to @@ -129,6 +132,7 @@ class ModelConfig(BaseConfig): ``PostprocessConfig()``. """ + compile: bool = False samplerate: int = Field(default=TARGET_SAMPLERATE_HZ, gt=0) architecture: BackboneConfig = Field(default_factory=UNetBackboneConfig) preprocess: PreprocessingConfig = Field( diff --git a/src/batdetect2/train/config.py b/src/batdetect2/train/config.py index 94584cb..d12f332 100644 --- a/src/batdetect2/train/config.py +++ b/src/batdetect2/train/config.py @@ -1,3 +1,5 @@ +from typing import Literal + from pydantic import Field from batdetect2.core.configs import BaseConfig @@ -39,6 +41,7 @@ class PLTrainerConfig(BaseConfig): class TrainingConfig(BaseConfig): + precision: Literal["medium", "high"] | None = None train_loader: TrainLoaderConfig = Field(default_factory=TrainLoaderConfig) val_loader: ValLoaderConfig = Field(default_factory=ValLoaderConfig) optimizer: OptimizerConfig = Field(default_factory=AdamOptimizerConfig) diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index f59138b..ed3218b 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -2,6 +2,7 @@ from collections.abc import Sequence from pathlib import Path from typing import Optional +import torch from lightning import Trainer, seed_everything from lightning.pytorch.loggers import Logger from loguru import logger @@ -206,6 +207,17 @@ def run_train( run_name=run_name, ) + if model_config.compile: + logger.info("Compiling model...") + module.compile() + + if train_config.precision is not None: + logger.info( + "Setting precision float precision to {}", + train_config.precision, + ) + torch.set_float32_matmul_precision(train_config.precision) + logger.info("Starting main training loop...") trainer.fit( module, From 573a0ca39295771e62b5e55f21e04e65dcb10303 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:46:32 +0100 Subject: [PATCH 02/10] Allow passing the logger object directly to the train workflow --- src/batdetect2/api_v2.py | 3 +++ src/batdetect2/train/train.py | 3 ++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/batdetect2/api_v2.py b/src/batdetect2/api_v2.py index d2c0a8f..e84b876 100644 --- a/src/batdetect2/api_v2.py +++ b/src/batdetect2/api_v2.py @@ -1,4 +1,5 @@ from __future__ import annotations +from lightning.pytorch.loggers import Logger from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -192,6 +193,7 @@ class BatDetect2API: train_config: TrainingConfig | None = None, logger_config: LoggerConfig | None = None, logging_callbacks: Sequence[LoggingCallback[TrainLoggingContext]] = (), + train_logger: Logger | None = None, ): """Train the current model on a set of annotations. @@ -255,6 +257,7 @@ class BatDetect2API: audio_config=audio_config or self.audio_config, logger_config=logger_config or self.logging_config.train, logging_callbacks=logging_callbacks, + train_logger=train_logger, ) self.model.eval() return self diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index ed3218b..0f4d8aa 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -72,6 +72,7 @@ def run_train( run_name: str | None = None, seed: int | None = None, logging_callbacks: Sequence[LoggingCallback[TrainLoggingContext]] = (), + train_logger: Logger | None = None, ): if seed is not None: seed_everything(seed) @@ -165,7 +166,7 @@ def run_train( roi_mapper=roi_mapper, ) - train_logger = build_logger( + train_logger = train_logger or build_logger( logger_config or CSVLoggerConfig(), log_dir=log_dir, experiment_name=experiment_name, From 1bab80f8c6cd880adfc306d4632e582ddfe45653 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:46:51 +0100 Subject: [PATCH 03/10] Expand the cosine annealing config --- src/batdetect2/train/schedulers.py | 39 ++++++++++++++++++------------ 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/src/batdetect2/train/schedulers.py b/src/batdetect2/train/schedulers.py index 73ebd09..a21ef5b 100644 --- a/src/batdetect2/train/schedulers.py +++ b/src/batdetect2/train/schedulers.py @@ -23,21 +23,6 @@ __all__ = [ ] -class CosineAnnealingSchedulerConfig(BaseConfig): - """Configuration for ``CosineAnnealingLR``. - - Attributes - ---------- - name : Literal["cosine_annealing"] - Discriminator field used by the scheduler registry. - t_max : int - Number of epochs to complete one cosine cycle. - """ - - name: Literal["cosine_annealing"] = "cosine_annealing" - t_max: int = 200 - - scheduler_registry: Registry[LRScheduler, [Optimizer]] = Registry("scheduler") @@ -53,6 +38,24 @@ class SchedulerImportConfig(ImportConfig): name: Literal["import"] = "import" +class CosineAnnealingSchedulerConfig(BaseConfig): + """Configuration for ``CosineAnnealingLR``. + + Attributes + ---------- + name : Literal["cosine_annealing"] + Discriminator field used by the scheduler registry. + t_max : int + Number of epochs to complete one cosine cycle. + eta_min : float, optional + Minimum learning rate. Defaults to 0. + """ + + name: Literal["cosine_annealing"] = "cosine_annealing" + t_max: int = 200 + eta_min: float = 0 + + @scheduler_registry.register(CosineAnnealingSchedulerConfig) def build_cosine_scheduler( config: CosineAnnealingSchedulerConfig, @@ -63,7 +66,11 @@ def build_cosine_scheduler( ``t_max`` is interpreted in epochs because Lightning steps the scheduler once per epoch when ``interval="epoch"`` is used. """ - return CosineAnnealingLR(optimizer, T_max=config.t_max) + return CosineAnnealingLR( + optimizer, + T_max=config.t_max, + eta_min=config.eta_min, + ) SchedulerConfig = Annotated[ From 74922802185dbf19b290c202ac797ea8747f1bd1 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:45:29 +0100 Subject: [PATCH 04/10] Add setting for compiling the model before train --- src/batdetect2/models/__init__.py | 4 ++++ src/batdetect2/train/config.py | 3 +++ src/batdetect2/train/train.py | 12 ++++++++++++ 3 files changed, 19 insertions(+) diff --git a/src/batdetect2/models/__init__.py b/src/batdetect2/models/__init__.py index ee96d93..8564d05 100644 --- a/src/batdetect2/models/__init__.py +++ b/src/batdetect2/models/__init__.py @@ -112,6 +112,9 @@ class ModelConfig(BaseConfig): Attributes ---------- + compile : bool + If ``True``, compile the model before training. Defaults to + ``False``. samplerate : int Expected input audio sample rate in Hz. Audio must be resampled to this rate before being passed to the model. Defaults to @@ -129,6 +132,7 @@ class ModelConfig(BaseConfig): ``PostprocessConfig()``. """ + compile: bool = False samplerate: int = Field(default=TARGET_SAMPLERATE_HZ, gt=0) architecture: BackboneConfig = Field(default_factory=UNetBackboneConfig) preprocess: PreprocessingConfig = Field( diff --git a/src/batdetect2/train/config.py b/src/batdetect2/train/config.py index 94584cb..d12f332 100644 --- a/src/batdetect2/train/config.py +++ b/src/batdetect2/train/config.py @@ -1,3 +1,5 @@ +from typing import Literal + from pydantic import Field from batdetect2.core.configs import BaseConfig @@ -39,6 +41,7 @@ class PLTrainerConfig(BaseConfig): class TrainingConfig(BaseConfig): + precision: Literal["medium", "high"] | None = None train_loader: TrainLoaderConfig = Field(default_factory=TrainLoaderConfig) val_loader: ValLoaderConfig = Field(default_factory=ValLoaderConfig) optimizer: OptimizerConfig = Field(default_factory=AdamOptimizerConfig) diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index f59138b..ed3218b 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -2,6 +2,7 @@ from collections.abc import Sequence from pathlib import Path from typing import Optional +import torch from lightning import Trainer, seed_everything from lightning.pytorch.loggers import Logger from loguru import logger @@ -206,6 +207,17 @@ def run_train( run_name=run_name, ) + if model_config.compile: + logger.info("Compiling model...") + module.compile() + + if train_config.precision is not None: + logger.info( + "Setting precision float precision to {}", + train_config.precision, + ) + torch.set_float32_matmul_precision(train_config.precision) + logger.info("Starting main training loop...") trainer.fit( module, From ead0adf284a210d75958c56260587a2fc42dda1d Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:46:32 +0100 Subject: [PATCH 05/10] Allow passing the logger object directly to the train workflow --- src/batdetect2/api_v2.py | 3 +++ src/batdetect2/train/train.py | 3 ++- 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/batdetect2/api_v2.py b/src/batdetect2/api_v2.py index d2c0a8f..e84b876 100644 --- a/src/batdetect2/api_v2.py +++ b/src/batdetect2/api_v2.py @@ -1,4 +1,5 @@ from __future__ import annotations +from lightning.pytorch.loggers import Logger from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -192,6 +193,7 @@ class BatDetect2API: train_config: TrainingConfig | None = None, logger_config: LoggerConfig | None = None, logging_callbacks: Sequence[LoggingCallback[TrainLoggingContext]] = (), + train_logger: Logger | None = None, ): """Train the current model on a set of annotations. @@ -255,6 +257,7 @@ class BatDetect2API: audio_config=audio_config or self.audio_config, logger_config=logger_config or self.logging_config.train, logging_callbacks=logging_callbacks, + train_logger=train_logger, ) self.model.eval() return self diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index ed3218b..0f4d8aa 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -72,6 +72,7 @@ def run_train( run_name: str | None = None, seed: int | None = None, logging_callbacks: Sequence[LoggingCallback[TrainLoggingContext]] = (), + train_logger: Logger | None = None, ): if seed is not None: seed_everything(seed) @@ -165,7 +166,7 @@ def run_train( roi_mapper=roi_mapper, ) - train_logger = build_logger( + train_logger = train_logger or build_logger( logger_config or CSVLoggerConfig(), log_dir=log_dir, experiment_name=experiment_name, From bcc9cb6d005186e2a353bbb871f7a06654626a49 Mon Sep 17 00:00:00 2001 From: Santiago Martinez Balvanera Date: Tue, 4 Aug 2026 08:46:51 +0100 Subject: [PATCH 06/10] Expand the cosine annealing config --- src/batdetect2/train/schedulers.py | 39 ++++++++++++++++++------------ 1 file changed, 23 insertions(+), 16 deletions(-) diff --git a/src/batdetect2/train/schedulers.py b/src/batdetect2/train/schedulers.py index 73ebd09..a21ef5b 100644 --- a/src/batdetect2/train/schedulers.py +++ b/src/batdetect2/train/schedulers.py @@ -23,21 +23,6 @@ __all__ = [ ] -class CosineAnnealingSchedulerConfig(BaseConfig): - """Configuration for ``CosineAnnealingLR``. - - Attributes - ---------- - name : Literal["cosine_annealing"] - Discriminator field used by the scheduler registry. - t_max : int - Number of epochs to complete one cosine cycle. - """ - - name: Literal["cosine_annealing"] = "cosine_annealing" - t_max: int = 200 - - scheduler_registry: Registry[LRScheduler, [Optimizer]] = Registry("scheduler") @@ -53,6 +38,24 @@ class SchedulerImportConfig(ImportConfig): name: Literal["import"] = "import" +class CosineAnnealingSchedulerConfig(BaseConfig): + """Configuration for ``CosineAnnealingLR``. + + Attributes + ---------- + name : Literal["cosine_annealing"] + Discriminator field used by the scheduler registry. + t_max : int + Number of epochs to complete one cosine cycle. + eta_min : float, optional + Minimum learning rate. Defaults to 0. + """ + + name: Literal["cosine_annealing"] = "cosine_annealing" + t_max: int = 200 + eta_min: float = 0 + + @scheduler_registry.register(CosineAnnealingSchedulerConfig) def build_cosine_scheduler( config: CosineAnnealingSchedulerConfig, @@ -63,7 +66,11 @@ def build_cosine_scheduler( ``t_max`` is interpreted in epochs because Lightning steps the scheduler once per epoch when ``interval="epoch"`` is used. """ - return CosineAnnealingLR(optimizer, T_max=config.t_max) + return CosineAnnealingLR( + optimizer, + T_max=config.t_max, + eta_min=config.eta_min, + ) SchedulerConfig = Annotated[ From 8166db8f9c59c566f940602928cae0261ff0ca42 Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Sat, 8 Aug 2026 12:02:00 +0100 Subject: [PATCH 07/10] feat: add runtime model compilation options --- justfile | 2 +- src/batdetect2/api_v2.py | 35 +++++++++++- src/batdetect2/inference/batch.py | 4 ++ src/batdetect2/inference/config.py | 1 + src/batdetect2/models/__init__.py | 17 ++++-- src/batdetect2/train/config.py | 1 + src/batdetect2/train/train.py | 12 ++-- tests/conftest.py | 38 +++++++++++-- tests/test_api_v2/test_api_v2.py | 34 +++++++++++ tests/test_inference/test_batch.py | 27 +++++++++ tests/test_train/test_lightning.py | 91 +++++++++++++++++++++++++----- 11 files changed, 229 insertions(+), 33 deletions(-) diff --git a/justfile b/justfile index 8e889f1..142239d 100644 --- a/justfile +++ b/justfile @@ -136,7 +136,7 @@ clean: clean-build clean-pyc clean-test clean-docs # Train on example data. example-train OPTIONS="": - uv run batdetect2 train \ + uv run batdetect2 -v train \ --val-dataset example_data/dataset.yaml \ --base-dir . \ --targets example_data/targets.yaml \ diff --git a/src/batdetect2/api_v2.py b/src/batdetect2/api_v2.py index e84b876..3928dc7 100644 --- a/src/batdetect2/api_v2.py +++ b/src/batdetect2/api_v2.py @@ -1,5 +1,4 @@ from __future__ import annotations -from lightning.pytorch.loggers import Logger from pathlib import Path from typing import TYPE_CHECKING, Literal @@ -9,6 +8,7 @@ if TYPE_CHECKING: import numpy as np import torch + from lightning.pytorch.loggers import Logger from soundevent import data from batdetect2.audio import AudioConfig, AudioLoader @@ -153,6 +153,19 @@ class BatDetect2API: self.model.eval() + def compile(self) -> "BatDetect2API": + """Compile the detector path used by inference. + + Returns + ------- + BatDetect2API + This API instance with the detector compiled. + """ + from batdetect2.models import compile_model + + compile_model(self.model) + return self + def load_annotations( self, path: data.PathLike, @@ -982,6 +995,7 @@ class BatDetect2API: inference_config: InferenceConfig | None = None, outputs_config: OutputsConfig | None = None, logging_config: AppLoggingConfig | None = None, + compile_model: bool = False, ) -> "BatDetect2API": """Build an API instance from config objects. @@ -1007,6 +1021,8 @@ class BatDetect2API: Output config. If omitted, the default outputs config is used. logging_config : AppLoggingConfig | None, optional Logging config. If omitted, the default logging config is used. + compile_model : bool, optional + If ``True``, compile the detector path after building the API. Returns ------- @@ -1089,7 +1105,7 @@ class BatDetect2API: ), ) - return cls( + api = cls( model_config=model_config, audio_config=audio_config, train_config=train_config, @@ -1108,6 +1124,11 @@ class BatDetect2API: output_transform=output_transform, ) + if compile_model: + api.compile() + + return api + @classmethod def from_checkpoint( cls, @@ -1118,6 +1139,7 @@ class BatDetect2API: inference_config: InferenceConfig | None = None, outputs_config: OutputsConfig | None = None, logging_config: AppLoggingConfig | None = None, + compile_model: bool = False, ) -> "BatDetect2API": """Build an API instance from a saved checkpoint. @@ -1138,6 +1160,8 @@ class BatDetect2API: Output config override. logging_config : AppLoggingConfig | None, optional Logging config override. + compile_model : bool, optional + If ``True``, compile the detector path after building the API. Returns ------- @@ -1221,7 +1245,7 @@ class BatDetect2API: transform=output_transform, ) - return cls( + api = cls( model_config=model_config, audio_config=audio_config, train_config=train_config, @@ -1240,6 +1264,11 @@ class BatDetect2API: output_transform=output_transform, ) + if compile_model: + api.compile() + + return api + def _set_trainable_parameters( self, trainable: Literal["all", "heads", "classifier_head", "size_head"], diff --git a/src/batdetect2/inference/batch.py b/src/batdetect2/inference/batch.py index d91df29..9c256c4 100644 --- a/src/batdetect2/inference/batch.py +++ b/src/batdetect2/inference/batch.py @@ -10,6 +10,7 @@ from batdetect2.inference.clips import get_clips_from_files from batdetect2.inference.config import InferenceConfig from batdetect2.inference.dataset import build_inference_loader from batdetect2.inference.lightning import InferenceModule +from batdetect2.models import compile_model from batdetect2.models.types import ModelProtocol from batdetect2.outputs import ( OutputsConfig, @@ -71,6 +72,9 @@ def run_batch_inference( batch_size=batch_size, ) + if inference_config.compile_model: + compile_model(model) + module = InferenceModule( model, output_transform=output_transform, diff --git a/src/batdetect2/inference/config.py b/src/batdetect2/inference/config.py index 1db715d..447f2a3 100644 --- a/src/batdetect2/inference/config.py +++ b/src/batdetect2/inference/config.py @@ -15,6 +15,7 @@ class ClipingConfig(BaseConfig): class InferenceConfig(BaseConfig): + compile_model: bool = False loader: InferenceLoaderConfig = Field( default_factory=InferenceLoaderConfig ) diff --git a/src/batdetect2/models/__init__.py b/src/batdetect2/models/__init__.py index 8564d05..1756ef3 100644 --- a/src/batdetect2/models/__init__.py +++ b/src/batdetect2/models/__init__.py @@ -100,6 +100,7 @@ __all__ = [ "ModelConfig", "build_model", "build_model_with_new_targets", + "compile_model", ] @@ -112,9 +113,6 @@ class ModelConfig(BaseConfig): Attributes ---------- - compile : bool - If ``True``, compile the model before training. Defaults to - ``False``. samplerate : int Expected input audio sample rate in Hz. Audio must be resampled to this rate before being passed to the model. Defaults to @@ -132,7 +130,6 @@ class ModelConfig(BaseConfig): ``PostprocessConfig()``. """ - compile: bool = False samplerate: int = Field(default=TARGET_SAMPLERATE_HZ, gt=0) architecture: BackboneConfig = Field(default_factory=UNetBackboneConfig) preprocess: PreprocessingConfig = Field( @@ -323,3 +320,15 @@ def build_model_with_new_targets( dimension_names=roi_mapper.dimension_names, config=model.get_config(), ) + + +def compile_model(model: ModelProtocol) -> ModelProtocol: + """Compile the detector path used by training and inference.""" + if not isinstance(model.detector, torch.nn.Module): + raise TypeError("Detector must be a torch.nn.Module to compile.") + + if getattr(model.detector, "_compiled_call_impl", None) is not None: + return model + + model.detector.compile() + return model diff --git a/src/batdetect2/train/config.py b/src/batdetect2/train/config.py index d12f332..8aec47d 100644 --- a/src/batdetect2/train/config.py +++ b/src/batdetect2/train/config.py @@ -41,6 +41,7 @@ class PLTrainerConfig(BaseConfig): class TrainingConfig(BaseConfig): + compile_model: bool = False precision: Literal["medium", "high"] | None = None train_loader: TrainLoaderConfig = Field(default_factory=TrainLoaderConfig) val_loader: ValLoaderConfig = Field(default_factory=ValLoaderConfig) diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index 0f4d8aa..f9430c3 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -16,7 +16,7 @@ from batdetect2.logging import ( LoggingCallback, build_logger, ) -from batdetect2.models import ModelConfig, build_model +from batdetect2.models import ModelConfig, build_model, compile_model from batdetect2.models.types import ModelProtocol from batdetect2.preprocess import PreprocessorProtocol, build_preprocessor from batdetect2.targets import ( @@ -208,17 +208,17 @@ def run_train( run_name=run_name, ) - if model_config.compile: - logger.info("Compiling model...") - module.compile() - if train_config.precision is not None: logger.info( - "Setting precision float precision to {}", + "Setting float32 matmul precision to {}", train_config.precision, ) torch.set_float32_matmul_precision(train_config.precision) + if train_config.compile_model: + logger.info("Compiling detector...") + compile_model(module.model) + logger.info("Starting main training loop...") trainer.fit( module, diff --git a/tests/conftest.py b/tests/conftest.py index 8a6a0a3..7b78eac 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,6 @@ import uuid from pathlib import Path -from typing import Callable, List, Optional +from typing import Any, Callable, List, Optional, cast from uuid import uuid4 import lightning as L @@ -15,6 +15,7 @@ from batdetect2.audio.clips import build_clipper from batdetect2.audio.types import AudioLoader, ClipperProtocol from batdetect2.data import DatasetConfig, load_dataset from batdetect2.data.annotations.batdetect2 import BatDetect2FilesAnnotations +from batdetect2.models.types import ModelProtocol from batdetect2.preprocess import build_preprocessor from batdetect2.preprocess.types import PreprocessorProtocol from batdetect2.targets import ( @@ -156,12 +157,15 @@ def generate_whistle(tmp_path: Path): offset = int((time - duration / 2) * samplerate) t = np.linspace(-duration / 2, duration / 2, frames, endpoint=False) - data = signal.gausspulse( - t, - fc=frequency, - bw=2 / (frequency * whistle_duration), + pulse = np.asarray( + signal.gausspulse( + t, + fc=frequency, + bw=2 / (frequency * whistle_duration), + ), + dtype=np.float64, ) - wave = (np.roll(data, offset) * np.iinfo(np.int16).max).astype( + wave = (np.roll(pulse, offset) * np.iinfo(np.int16).max).astype( np.int16 ) sf.write(str(path), wave, samplerate, subtype="PCM_16") @@ -363,6 +367,28 @@ def sample_audio_loader() -> AudioLoader: return build_audio_loader() +@pytest.fixture +def record_compiled_detector_calls( + monkeypatch: pytest.MonkeyPatch, +) -> Callable[[ModelProtocol], list[None]]: + def factory(model: ModelProtocol) -> list[None]: + compiled_calls: list[None] = [] + detector = cast(Any, model.detector) + original_call_impl = detector._call_impl + + def compile_detector() -> None: + def compiled_call(*args, **kwargs): + compiled_calls.append(None) + return original_call_impl(*args, **kwargs) + + detector._compiled_call_impl = compiled_call + + monkeypatch.setattr(detector, "compile", compile_detector) + return compiled_calls + + return factory + + @pytest.fixture def bat_tag() -> data.Tag: return data.Tag(key="class", value="bat") diff --git a/tests/test_api_v2/test_api_v2.py b/tests/test_api_v2/test_api_v2.py index dc1794d..0cc12ec 100644 --- a/tests/test_api_v2/test_api_v2.py +++ b/tests/test_api_v2/test_api_v2.py @@ -153,6 +153,40 @@ def test_process_spectrogram_rejects_batched_input( api_v2.process_spectrogram(spec) +def test_user_can_compile_api_detector( + api_v2: BatDetect2API, + example_audio_files: list[Path], + record_compiled_detector_calls, +) -> None: + compiled_calls = record_compiled_detector_calls(api_v2.model) + audio = api_v2.load_audio(example_audio_files[0]) + spec = api_v2.generate_spectrogram(audio) + + api_v2.compile() + api_v2.compile() + api_v2.process_spectrogram(spec) + + assert len(compiled_calls) == 1 + + +def test_api_from_config_compiles_detector_when_requested( + monkeypatch: pytest.MonkeyPatch, +) -> None: + compiled_models = [] + + def compile_model(model): + compiled_models.append(model) + return model + + monkeypatch.setattr("batdetect2.models.compile_model", compile_model) + + api = BatDetect2API.from_config( + compile_model=True, + ) + + assert compiled_models == [api.model] + + def test_user_can_read_top_class_and_other_class_scores( api_v2: BatDetect2API, example_audio_files: list[Path], diff --git a/tests/test_inference/test_batch.py b/tests/test_inference/test_batch.py index f835f5e..39751dd 100644 --- a/tests/test_inference/test_batch.py +++ b/tests/test_inference/test_batch.py @@ -3,6 +3,8 @@ from pathlib import Path import pytest from soundevent import data +from batdetect2.api_v2 import BatDetect2API +from batdetect2.inference import InferenceConfig from batdetect2.inference.batch import run_batch_inference from batdetect2.targets import build_roi_mapping, build_targets from batdetect2.train import load_model_from_checkpoint @@ -53,3 +55,28 @@ def test_run_batch_inference_matches_single_clip_inference( strict=True, ): assert_clip_detections_equal(batched, single) + + +def test_run_batch_inference_compiles_detector_when_config_requests_compile( + example_annotations: list[data.ClipAnnotation], + record_compiled_detector_calls, +) -> None: + api = BatDetect2API.from_config() + compiled_calls = record_compiled_detector_calls(api.model) + api.compile() + + predictions = run_batch_inference( + api.model, + [example_annotations[0].clip], + targets=api.targets, + roi_mapper=api.roi_mapper, + audio_loader=api.audio_loader, + preprocessor=api.preprocessor, + output_transform=api.output_transform, + inference_config=InferenceConfig(compile_model=True), + batch_size=1, + num_workers=0, + ) + + assert predictions + assert compiled_calls diff --git a/tests/test_train/test_lightning.py b/tests/test_train/test_lightning.py index 756329d..772ffed 100644 --- a/tests/test_train/test_lightning.py +++ b/tests/test_train/test_lightning.py @@ -49,6 +49,16 @@ def build_default_module( ) +def build_fast_train_config() -> TrainingConfig: + train_config = TrainingConfig() + train_config.trainer.limit_train_batches = 1 + train_config.trainer.limit_val_batches = 1 + train_config.trainer.log_every_n_steps = 1 + train_config.train_loader.batch_size = 1 + train_config.train_loader.augmentations.enabled = False + return train_config + + def test_can_initialize_default_module(): module = build_default_module() assert isinstance(module, L.LightningModule) @@ -271,19 +281,7 @@ def test_train_smoke_produces_loadable_checkpoint( sample_audio_loader: AudioLoader, ): # Given - train_config = TrainingConfig.model_validate( - { - "trainer": { - "limit_train_batches": 1, - "limit_val_batches": 1, - "log_every_n_steps": 1, - }, - "train_loader": { - "batch_size": 1, - "augmentations": {"enabled": False}, - }, - } - ) + train_config = build_fast_train_config() # When run_train( @@ -310,6 +308,73 @@ def test_train_smoke_produces_loadable_checkpoint( assert outputs is not None +@pytest.mark.slow +def test_run_train_compiles_detector_when_train_config_requests_compile( + tmp_path: Path, + example_annotations: list[data.ClipAnnotation], + record_compiled_detector_calls, +) -> None: + targets_config = TargetConfig() + targets = build_targets(targets_config) + roi_mapper = build_roi_mapping(targets_config.roi) + model = build_model( + ModelConfig(), + class_names=targets.class_names, + dimension_names=roi_mapper.dimension_names, + ) + train_config = build_fast_train_config() + train_config.compile_model = True + compiled_calls = record_compiled_detector_calls(model) + + module = run_train( + train_annotations=example_annotations[:1], + val_annotations=example_annotations[:1], + model=model, + targets=targets, + roi_mapper=roi_mapper, + targets_config=targets_config, + train_config=train_config, + num_epochs=1, + train_workers=0, + val_workers=0, + checkpoint_dir=tmp_path / "checkpoints", + log_dir=tmp_path / "logs", + seed=0, + ) + + assert ( + getattr(module.model.detector, "_compiled_call_impl", None) is not None + ) + assert compiled_calls + + +@pytest.mark.slow +def test_run_train_sets_float32_matmul_precision( + tmp_path: Path, + example_annotations: list[data.ClipAnnotation], +) -> None: + original_precision = torch.get_float32_matmul_precision() + train_config = build_fast_train_config() + train_config.precision = "high" + + try: + run_train( + train_annotations=example_annotations[:1], + val_annotations=example_annotations[:1], + train_config=train_config, + num_epochs=1, + train_workers=0, + val_workers=0, + checkpoint_dir=tmp_path / "checkpoints", + log_dir=tmp_path / "logs", + seed=0, + ) + + assert torch.get_float32_matmul_precision() == "high" + finally: + torch.set_float32_matmul_precision(original_precision) + + def test_build_training_module_uses_provided_model() -> None: targets = build_targets(TargetConfig()) roi_mapper = build_roi_mapping(TargetConfig().roi) From bad5d4f4fca1a8decf13514f381438789ec4ab0e Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Sat, 8 Aug 2026 12:04:19 +0100 Subject: [PATCH 08/10] test: cover compile idempotency paths --- tests/conftest.py | 21 ++++++++++++++------ tests/test_api_v2/test_api_v2.py | 27 ++++++++++++++++++++++--- tests/test_inference/test_batch.py | 32 +++++++++++++++++++++++++++--- tests/test_train/test_lightning.py | 7 ++++--- 4 files changed, 72 insertions(+), 15 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index 7b78eac..52286d4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,4 +1,5 @@ import uuid +from dataclasses import dataclass from pathlib import Path from typing import Any, Callable, List, Optional, cast from uuid import uuid4 @@ -32,6 +33,12 @@ from batdetect2.train.lightning import build_training_module from batdetect2.train.types import ClipLabeller +@dataclass +class DetectorCompileRecorder: + compile_count: int = 0 + call_count: int = 0 + + @pytest.fixture def example_data_dir() -> Path: pkg_dir = Path(__file__).parent.parent @@ -368,23 +375,25 @@ def sample_audio_loader() -> AudioLoader: @pytest.fixture -def record_compiled_detector_calls( +def record_detector_compilation( monkeypatch: pytest.MonkeyPatch, -) -> Callable[[ModelProtocol], list[None]]: - def factory(model: ModelProtocol) -> list[None]: - compiled_calls: list[None] = [] +) -> Callable[[ModelProtocol], DetectorCompileRecorder]: + def factory(model: ModelProtocol) -> DetectorCompileRecorder: + recorder = DetectorCompileRecorder() detector = cast(Any, model.detector) original_call_impl = detector._call_impl def compile_detector() -> None: + recorder.compile_count += 1 + def compiled_call(*args, **kwargs): - compiled_calls.append(None) + recorder.call_count += 1 return original_call_impl(*args, **kwargs) detector._compiled_call_impl = compiled_call monkeypatch.setattr(detector, "compile", compile_detector) - return compiled_calls + return recorder return factory diff --git a/tests/test_api_v2/test_api_v2.py b/tests/test_api_v2/test_api_v2.py index 0cc12ec..102eacf 100644 --- a/tests/test_api_v2/test_api_v2.py +++ b/tests/test_api_v2/test_api_v2.py @@ -156,9 +156,9 @@ def test_process_spectrogram_rejects_batched_input( def test_user_can_compile_api_detector( api_v2: BatDetect2API, example_audio_files: list[Path], - record_compiled_detector_calls, + record_detector_compilation, ) -> None: - compiled_calls = record_compiled_detector_calls(api_v2.model) + recorder = record_detector_compilation(api_v2.model) audio = api_v2.load_audio(example_audio_files[0]) spec = api_v2.generate_spectrogram(audio) @@ -166,7 +166,8 @@ def test_user_can_compile_api_detector( api_v2.compile() api_v2.process_spectrogram(spec) - assert len(compiled_calls) == 1 + assert recorder.compile_count == 1 + assert recorder.call_count == 1 def test_api_from_config_compiles_detector_when_requested( @@ -187,6 +188,26 @@ def test_api_from_config_compiles_detector_when_requested( assert compiled_models == [api.model] +def test_api_from_checkpoint_compiles_detector_when_requested( + tiny_checkpoint_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + compiled_models = [] + + def compile_model(model): + compiled_models.append(model) + return model + + monkeypatch.setattr("batdetect2.models.compile_model", compile_model) + + api = BatDetect2API.from_checkpoint( + tiny_checkpoint_path, + compile_model=True, + ) + + assert compiled_models == [api.model] + + def test_user_can_read_top_class_and_other_class_scores( api_v2: BatDetect2API, example_audio_files: list[Path], diff --git a/tests/test_inference/test_batch.py b/tests/test_inference/test_batch.py index 39751dd..6a2064d 100644 --- a/tests/test_inference/test_batch.py +++ b/tests/test_inference/test_batch.py @@ -59,10 +59,35 @@ def test_run_batch_inference_matches_single_clip_inference( def test_run_batch_inference_compiles_detector_when_config_requests_compile( example_annotations: list[data.ClipAnnotation], - record_compiled_detector_calls, + record_detector_compilation, ) -> None: api = BatDetect2API.from_config() - compiled_calls = record_compiled_detector_calls(api.model) + recorder = record_detector_compilation(api.model) + + predictions = run_batch_inference( + api.model, + [example_annotations[0].clip], + targets=api.targets, + roi_mapper=api.roi_mapper, + audio_loader=api.audio_loader, + preprocessor=api.preprocessor, + output_transform=api.output_transform, + inference_config=InferenceConfig(compile_model=True), + batch_size=1, + num_workers=0, + ) + + assert predictions + assert recorder.compile_count == 1 + assert recorder.call_count > 0 + + +def test_run_batch_inference_does_not_recompile_compiled_detector( + example_annotations: list[data.ClipAnnotation], + record_detector_compilation, +) -> None: + api = BatDetect2API.from_config() + recorder = record_detector_compilation(api.model) api.compile() predictions = run_batch_inference( @@ -79,4 +104,5 @@ def test_run_batch_inference_compiles_detector_when_config_requests_compile( ) assert predictions - assert compiled_calls + assert recorder.compile_count == 1 + assert recorder.call_count > 0 diff --git a/tests/test_train/test_lightning.py b/tests/test_train/test_lightning.py index 772ffed..feb9a97 100644 --- a/tests/test_train/test_lightning.py +++ b/tests/test_train/test_lightning.py @@ -312,7 +312,7 @@ def test_train_smoke_produces_loadable_checkpoint( def test_run_train_compiles_detector_when_train_config_requests_compile( tmp_path: Path, example_annotations: list[data.ClipAnnotation], - record_compiled_detector_calls, + record_detector_compilation, ) -> None: targets_config = TargetConfig() targets = build_targets(targets_config) @@ -324,7 +324,7 @@ def test_run_train_compiles_detector_when_train_config_requests_compile( ) train_config = build_fast_train_config() train_config.compile_model = True - compiled_calls = record_compiled_detector_calls(model) + recorder = record_detector_compilation(model) module = run_train( train_annotations=example_annotations[:1], @@ -345,7 +345,8 @@ def test_run_train_compiles_detector_when_train_config_requests_compile( assert ( getattr(module.model.detector, "_compiled_call_impl", None) is not None ) - assert compiled_calls + assert recorder.compile_count == 1 + assert recorder.call_count > 0 @pytest.mark.slow From c8f0f2bee12f525b99dc779642d749892f79905b Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Sat, 8 Aug 2026 12:18:55 +0100 Subject: [PATCH 09/10] docs: document runtime compilation options --- docs/source/reference/api.md | 5 +++++ .../configs/inference/inference-config.md | 13 +++++++++++++ .../configs/training/training-config.md | 16 ++++++++++++++++ src/batdetect2/api_v2.py | 3 +++ 4 files changed, 37 insertions(+) diff --git a/docs/source/reference/api.md b/docs/source/reference/api.md index fa12b25..1221ee8 100644 --- a/docs/source/reference/api.md +++ b/docs/source/reference/api.md @@ -15,6 +15,9 @@ Defined in `batdetect2.api_v2`. - `BatDetect2API.from_config(model_config=..., targets_config=..., ...)` - build a full model stack from config objects. +Both constructors accept `compile_model=True` to compile the detector after the +API is built. + ## Common tasks - Load a checkpoint and run prediction on one file. @@ -22,6 +25,8 @@ Defined in `batdetect2.api_v2`. - Save predictions in one of the supported output formats. - Evaluate a model on labelled data. - Fine-tune an existing checkpoint on new targets. +- Compile the detector explicitly with `BatDetect2API.compile()` when you want + to opt into PyTorch runtime compilation from Python. ## Generated reference diff --git a/docs/source/reference/configs/inference/inference-config.md b/docs/source/reference/configs/inference/inference-config.md index a3b5706..0e3a98b 100644 --- a/docs/source/reference/configs/inference/inference-config.md +++ b/docs/source/reference/configs/inference/inference-config.md @@ -7,6 +7,8 @@ Defined in `batdetect2.inference.config`. ## Top-level fields +- `compile_model` + - compile the detector before batch prediction. This is off by default. - `loader` - data-loader settings for inference. - `clipping` @@ -34,8 +36,19 @@ Override `InferenceConfig` when: - long recordings need different clipping behavior, - you want to tune batch size for your hardware, +- you want to opt into runtime model compilation for repeated predictions, - you need reproducible prediction settings across runs. +## Runtime compilation + +Set `compile_model: true` to compile the detector before batch inference. This +can help when you run repeated predictions with stable input shapes. For a +single short run, the compile step can cost more time than it saves. + +In Python, you can also compile explicitly with `BatDetect2API.compile()` or by +passing `compile_model=True` to `BatDetect2API.from_checkpoint(...)` or +`BatDetect2API.from_config(...)`. + ## Related pages - Tune inference clipping: diff --git a/docs/source/reference/configs/training/training-config.md b/docs/source/reference/configs/training/training-config.md index f5fdca4..5f843cf 100644 --- a/docs/source/reference/configs/training/training-config.md +++ b/docs/source/reference/configs/training/training-config.md @@ -7,6 +7,10 @@ Defined in `batdetect2.train.config`. ## Top-level fields +- `compile_model` + - compile the detector before training starts. This is off by default. +- `precision` + - optional float32 matrix multiplication precision setting passed to PyTorch. - `train_loader` - training data loading and clipping settings. - `val_loader` @@ -33,10 +37,22 @@ Use `TrainingConfig` when you want to change things like: - batch size, - augmentation, - optimiser and scheduler settings, +- runtime options such as model compilation and matrix multiplication + precision, - number of epochs, - validation frequency, - checkpoint behaviour. +## Runtime options + +Use `compile_model: true` to call `torch.compile` on the detector used during +training. This can help on longer runs with stable tensor shapes, but it may be +slower for short CPU-only experiments because PyTorch has to compile the graph +before it can reuse it. + +Use `precision` to set PyTorch's float32 matrix multiplication precision before +training starts. Supported values are `medium` and `high`. + Example files live under `example_data/configs/`, including `example_data/configs/training.yaml`. diff --git a/src/batdetect2/api_v2.py b/src/batdetect2/api_v2.py index 3928dc7..f78a2ee 100644 --- a/src/batdetect2/api_v2.py +++ b/src/batdetect2/api_v2.py @@ -241,6 +241,9 @@ class BatDetect2API: Training logger config override. logging_callbacks : Sequence[LoggingCallback[TrainLoggingContext]], optional Extra logging callbacks to run during training setup. + train_logger : Logger | None, optional + Pre-built Lightning logger to use for training. If omitted, one is + built from ``logger_config``. Returns ------- From 586e78814f0ca62f29f987773f4a4d8efef636ad Mon Sep 17 00:00:00 2001 From: mbsantiago Date: Sat, 8 Aug 2026 12:34:07 +0100 Subject: [PATCH 10/10] refactor: remove ambiguous training precision option --- .../configs/training/training-config.md | 8 +----- src/batdetect2/train/config.py | 3 --- src/batdetect2/train/train.py | 8 ------ tests/test_train/test_lightning.py | 27 ------------------- 4 files changed, 1 insertion(+), 45 deletions(-) diff --git a/docs/source/reference/configs/training/training-config.md b/docs/source/reference/configs/training/training-config.md index 5f843cf..bbc1ebb 100644 --- a/docs/source/reference/configs/training/training-config.md +++ b/docs/source/reference/configs/training/training-config.md @@ -9,8 +9,6 @@ Defined in `batdetect2.train.config`. - `compile_model` - compile the detector before training starts. This is off by default. -- `precision` - - optional float32 matrix multiplication precision setting passed to PyTorch. - `train_loader` - training data loading and clipping settings. - `val_loader` @@ -37,8 +35,7 @@ Use `TrainingConfig` when you want to change things like: - batch size, - augmentation, - optimiser and scheduler settings, -- runtime options such as model compilation and matrix multiplication - precision, +- runtime options such as model compilation, - number of epochs, - validation frequency, - checkpoint behaviour. @@ -50,9 +47,6 @@ training. This can help on longer runs with stable tensor shapes, but it may be slower for short CPU-only experiments because PyTorch has to compile the graph before it can reuse it. -Use `precision` to set PyTorch's float32 matrix multiplication precision before -training starts. Supported values are `medium` and `high`. - Example files live under `example_data/configs/`, including `example_data/configs/training.yaml`. diff --git a/src/batdetect2/train/config.py b/src/batdetect2/train/config.py index 8aec47d..ed1e201 100644 --- a/src/batdetect2/train/config.py +++ b/src/batdetect2/train/config.py @@ -1,5 +1,3 @@ -from typing import Literal - from pydantic import Field from batdetect2.core.configs import BaseConfig @@ -42,7 +40,6 @@ class PLTrainerConfig(BaseConfig): class TrainingConfig(BaseConfig): compile_model: bool = False - precision: Literal["medium", "high"] | None = None train_loader: TrainLoaderConfig = Field(default_factory=TrainLoaderConfig) val_loader: ValLoaderConfig = Field(default_factory=ValLoaderConfig) optimizer: OptimizerConfig = Field(default_factory=AdamOptimizerConfig) diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index f9430c3..e36c7ac 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -2,7 +2,6 @@ from collections.abc import Sequence from pathlib import Path from typing import Optional -import torch from lightning import Trainer, seed_everything from lightning.pytorch.loggers import Logger from loguru import logger @@ -208,13 +207,6 @@ def run_train( run_name=run_name, ) - if train_config.precision is not None: - logger.info( - "Setting float32 matmul precision to {}", - train_config.precision, - ) - torch.set_float32_matmul_precision(train_config.precision) - if train_config.compile_model: logger.info("Compiling detector...") compile_model(module.model) diff --git a/tests/test_train/test_lightning.py b/tests/test_train/test_lightning.py index feb9a97..35e8edc 100644 --- a/tests/test_train/test_lightning.py +++ b/tests/test_train/test_lightning.py @@ -349,33 +349,6 @@ def test_run_train_compiles_detector_when_train_config_requests_compile( assert recorder.call_count > 0 -@pytest.mark.slow -def test_run_train_sets_float32_matmul_precision( - tmp_path: Path, - example_annotations: list[data.ClipAnnotation], -) -> None: - original_precision = torch.get_float32_matmul_precision() - train_config = build_fast_train_config() - train_config.precision = "high" - - try: - run_train( - train_annotations=example_annotations[:1], - val_annotations=example_annotations[:1], - train_config=train_config, - num_epochs=1, - train_workers=0, - val_workers=0, - checkpoint_dir=tmp_path / "checkpoints", - log_dir=tmp_path / "logs", - seed=0, - ) - - assert torch.get_float32_matmul_precision() == "high" - finally: - torch.set_float32_matmul_precision(original_precision) - - def test_build_training_module_uses_provided_model() -> None: targets = build_targets(TargetConfig()) roi_mapper = build_roi_mapping(TargetConfig().roi)