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)