Merge pull request #75 from macaodha/feat/compile-model

Add runtime model compilation options

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
 

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:

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`.
 

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 \

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"],

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,

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
     )

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

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)

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[

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,

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(
+        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")

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],

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

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)