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..bbc1ebb 100644 --- a/docs/source/reference/configs/training/training-config.md +++ b/docs/source/reference/configs/training/training-config.md @@ -7,6 +7,8 @@ Defined in `batdetect2.train.config`. ## Top-level fields +- `compile_model` + - compile the detector before training starts. This is off by default. - `train_loader` - training data loading and clipping settings. - `val_loader` @@ -33,10 +35,18 @@ Use `TrainingConfig` when you want to change things like: - batch size, - augmentation, - optimiser and scheduler settings, +- runtime options such as model compilation, - 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. + Example files live under `example_data/configs/`, including `example_data/configs/training.yaml`. 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 d2c0a8f..f78a2ee 100644 --- a/src/batdetect2/api_v2.py +++ b/src/batdetect2/api_v2.py @@ -8,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 @@ -152,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, @@ -192,6 +206,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. @@ -226,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 ------- @@ -255,6 +273,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 @@ -979,6 +998,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. @@ -1004,6 +1024,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 ------- @@ -1086,7 +1108,7 @@ class BatDetect2API: ), ) - return cls( + api = cls( model_config=model_config, audio_config=audio_config, train_config=train_config, @@ -1105,6 +1127,11 @@ class BatDetect2API: output_transform=output_transform, ) + if compile_model: + api.compile() + + return api + @classmethod def from_checkpoint( cls, @@ -1115,6 +1142,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. @@ -1135,6 +1163,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 ------- @@ -1218,7 +1248,7 @@ class BatDetect2API: transform=output_transform, ) - return cls( + api = cls( model_config=model_config, audio_config=audio_config, train_config=train_config, @@ -1237,6 +1267,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 ee96d93..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", ] @@ -319,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 94584cb..ed1e201 100644 --- a/src/batdetect2/train/config.py +++ b/src/batdetect2/train/config.py @@ -39,6 +39,7 @@ class PLTrainerConfig(BaseConfig): class TrainingConfig(BaseConfig): + compile_model: bool = False 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/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[ diff --git a/src/batdetect2/train/train.py b/src/batdetect2/train/train.py index f59138b..e36c7ac 100644 --- a/src/batdetect2/train/train.py +++ b/src/batdetect2/train/train.py @@ -15,7 +15,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 ( @@ -71,6 +71,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) @@ -164,7 +165,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, @@ -206,6 +207,10 @@ def run_train( run_name=run_name, ) + 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..52286d4 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,6 +1,7 @@ import uuid +from dataclasses import dataclass 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 +16,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 ( @@ -31,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 @@ -156,12 +164,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 +374,30 @@ def sample_audio_loader() -> AudioLoader: return build_audio_loader() +@pytest.fixture +def record_detector_compilation( + monkeypatch: pytest.MonkeyPatch, +) -> 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): + recorder.call_count += 1 + return original_call_impl(*args, **kwargs) + + detector._compiled_call_impl = compiled_call + + monkeypatch.setattr(detector, "compile", compile_detector) + return recorder + + 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..102eacf 100644 --- a/tests/test_api_v2/test_api_v2.py +++ b/tests/test_api_v2/test_api_v2.py @@ -153,6 +153,61 @@ 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_detector_compilation, +) -> None: + recorder = record_detector_compilation(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 recorder.compile_count == 1 + assert recorder.call_count == 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_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 f835f5e..6a2064d 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,54 @@ 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_detector_compilation, +) -> None: + api = BatDetect2API.from_config() + 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( + 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 diff --git a/tests/test_train/test_lightning.py b/tests/test_train/test_lightning.py index 756329d..35e8edc 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,47 @@ 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_detector_compilation, +) -> 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 + recorder = record_detector_compilation(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 recorder.compile_count == 1 + assert recorder.call_count > 0 + + def test_build_training_module_uses_provided_model() -> None: targets = build_targets(TargetConfig()) roi_mapper = build_roi_mapping(TargetConfig().roi)