example_data/config.yaml

@@ -109,7 +109,7 @@ train:
     sigma: 3
 
   trainer:
-    max_epochs: 5
+    max_epochs: 40
 
   dataloaders:
     train:
@@ -136,9 +136,9 @@ train:
       weight: 0.1
 
   logger:
-    name: mlflow
-    tracking_uri: http://10.20.20.211:9000
-    log_model: true
+    logger_type: csv
+    # save_dir: outputs/log/
+    # name: logs
 
   augmentations:
     enabled: true

pyproject.toml

@@ -23,7 +23,7 @@ dependencies = [
   "tqdm>=4.66.2",
   "cf-xarray>=0.9.0",
   "onnx>=1.16.0",
-  "lightning[extra]==2.5.0",
+  "lightning[extra]>=2.2.2",
   "tensorboard>=2.16.2",
   "omegaconf>=2.3.0",
   "pyyaml>=6.0.2",

src/batdetect2/cli/train.py

@@ -26,9 +26,6 @@ __all__ = ["train_command"]
 @click.option("--config-field", type=str)
 @click.option("--train-workers", type=int)
 @click.option("--val-workers", type=int)
-@click.option("--experiment-name", type=str)
-@click.option("--run-name", type=str)
-@click.option("--seed", type=int)
 @click.option(
     "-v",
     "--verbose",
@@ -43,11 +40,8 @@ def train_command(
     log_dir: Optional[Path] = None,
     config: Optional[Path] = None,
     config_field: Optional[str] = None,
-    seed: Optional[int] = None,
     train_workers: int = 0,
     val_workers: int = 0,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
     verbose: int = 0,
 ):
     logger.remove()
@@ -93,9 +87,6 @@ def train_command(
         model_path=model_path,
         train_workers=train_workers,
         val_workers=val_workers,
-        experiment_name=experiment_name,
         log_dir=log_dir,
         checkpoint_dir=ckpt_dir,
-        seed=seed,
-        run_name=run_name,
     )

src/batdetect2/train/callbacks.py

@@ -68,7 +68,7 @@ class ValidationMetrics(Callback):
             n_examples=4,
         ):
             plotter(
-                f"examples/{class_name}",
+                f"images/{class_name}_examples",
                 fig,
                 pl_module.global_step,
             )

src/batdetect2/train/logging.py

@@ -1,16 +1,5 @@
 import io
-from pathlib import Path
-from typing import (
-    Annotated,
-    Any,
-    Dict,
-    Generic,
-    Literal,
-    Optional,
-    Protocol,
-    TypeVar,
-    Union,
-)
+from typing import Annotated, Any, Literal, Optional, Union
 
 import numpy as np
 from lightning.pytorch.loggers import Logger, MLFlowLogger, TensorBoardLogger
@@ -20,34 +9,39 @@ from soundevent import data
 
 from batdetect2.configs import BaseConfig
 
-DEFAULT_LOGS_DIR: Path = Path("outputs") / "logs"
+DEFAULT_LOGS_DIR: str = "outputs"
 
 
-class BaseLoggerConfig(BaseConfig):
-    log_dir: Path = DEFAULT_LOGS_DIR
-    experiment_name: Optional[str] = None
+class DVCLiveConfig(BaseConfig):
+    logger_type: Literal["dvclive"] = "dvclive"
+    dir: str = DEFAULT_LOGS_DIR
     run_name: Optional[str] = None
-
-
-class DVCLiveConfig(BaseLoggerConfig):
-    name: Literal["dvclive"] = "dvclive"
     prefix: str = ""
     log_model: Union[bool, Literal["all"]] = False
     monitor_system: bool = False
 
 
-class CSVLoggerConfig(BaseLoggerConfig):
-    name: Literal["csv"] = "csv"
+class CSVLoggerConfig(BaseConfig):
+    logger_type: Literal["csv"] = "csv"
+    save_dir: str = DEFAULT_LOGS_DIR
+    name: Optional[str] = "logs"
+    version: Optional[str] = None
     flush_logs_every_n_steps: int = 100
 
 
-class TensorBoardLoggerConfig(BaseLoggerConfig):
-    name: Literal["tensorboard"] = "tensorboard"
+class TensorBoardLoggerConfig(BaseConfig):
+    logger_type: Literal["tensorboard"] = "tensorboard"
+    save_dir: str = DEFAULT_LOGS_DIR
+    name: Optional[str] = "logs"
+    version: Optional[str] = None
     log_graph: bool = False
 
 
-class MLFlowLoggerConfig(BaseLoggerConfig):
-    name: Literal["mlflow"] = "mlflow"
+class MLFlowLoggerConfig(BaseConfig):
+    logger_type: Literal["mlflow"] = "mlflow"
+    experiment_name: str = "default"
+    run_name: Optional[str] = None
+    save_dir: Optional[str] = "./mlruns"
     tracking_uri: Optional[str] = None
     tags: Optional[dict[str, Any]] = None
     log_model: bool = False
@@ -60,28 +54,13 @@ LoggerConfig = Annotated[
         TensorBoardLoggerConfig,
         MLFlowLoggerConfig,
     ],
-    Field(discriminator="name"),
+    Field(discriminator="logger_type"),
 ]
 
 
-T = TypeVar("T", bound=LoggerConfig, contravariant=True)
-
-
-class LoggerBuilder(Protocol, Generic[T]):
-    def __call__(
-        self,
-        config: T,
-        log_dir: Optional[Path] = None,
-        experiment_name: Optional[str] = None,
-        run_name: Optional[str] = None,
-    ) -> Logger: ...
-
-
 def create_dvclive_logger(
     config: DVCLiveConfig,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    log_dir: Optional[data.PathLike] = None,
 ) -> Logger:
     try:
         from dvclive.lightning import DVCLiveLogger  # type: ignore
@@ -93,11 +72,8 @@ def create_dvclive_logger(
         ) from error
 
     return DVCLiveLogger(
-        dir=log_dir if log_dir is not None else config.log_dir,
-        run_name=run_name if run_name is not None else config.run_name,
-        experiment=experiment_name
-        if experiment_name is not None
-        else config.experiment_name,
+        dir=log_dir if log_dir is not None else config.dir,
+        run_name=config.run_name,
         prefix=config.prefix,
         log_model=config.log_model,
         monitor_system=config.monitor_system,
@@ -106,58 +82,28 @@ def create_dvclive_logger(
 
 def create_csv_logger(
     config: CSVLoggerConfig,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    log_dir: Optional[data.PathLike] = None,
 ) -> Logger:
     from lightning.pytorch.loggers import CSVLogger
 
-    if log_dir is None:
-        log_dir = Path(config.log_dir)
-
-    if run_name is None:
-        run_name = config.run_name
-
-    if experiment_name is None:
-        experiment_name = config.experiment_name
-
-    name = run_name
-
-    if run_name is not None and experiment_name is not None:
-        name = str(Path(experiment_name) / run_name)
-
     return CSVLogger(
-        save_dir=str(log_dir),
-        name=name,
+        save_dir=str(log_dir) if log_dir is not None else config.save_dir,
+        name=config.name,
+        version=config.version,
         flush_logs_every_n_steps=config.flush_logs_every_n_steps,
     )
 
 
 def create_tensorboard_logger(
     config: TensorBoardLoggerConfig,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    log_dir: Optional[data.PathLike] = None,
 ) -> Logger:
     from lightning.pytorch.loggers import TensorBoardLogger
 
-    if log_dir is None:
-        log_dir = Path(config.log_dir)
-
-    if run_name is None:
-        run_name = config.run_name
-
-    if experiment_name is None:
-        experiment_name = config.experiment_name
-
-    name = run_name
-
-    if run_name is not None and experiment_name is not None:
-        name = str(Path(experiment_name) / run_name)
-
     return TensorBoardLogger(
-        save_dir=str(log_dir),
-        name=name,
+        save_dir=str(log_dir) if log_dir is not None else config.save_dir,
+        name=config.name,
+        version=config.version,
         log_graph=config.log_graph,
     )
 
@@ -165,8 +111,6 @@ def create_tensorboard_logger(
 def create_mlflow_logger(
     config: MLFlowLoggerConfig,
     log_dir: Optional[data.PathLike] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
 ) -> Logger:
     try:
         from lightning.pytorch.loggers import MLFlowLogger
@@ -177,25 +121,17 @@ def create_mlflow_logger(
             "or `uv add mlflow`"
         ) from error
 
-    if experiment_name is None:
-        experiment_name = config.experiment_name or "Default"
-
-    if log_dir is None:
-        log_dir = config.log_dir
-
     return MLFlowLogger(
-        experiment_name=experiment_name
-        if experiment_name is not None
-        else config.experiment_name,
-        run_name=run_name if run_name is not None else config.run_name,
-        save_dir=str(log_dir),
+        experiment_name=config.experiment_name,
+        run_name=config.run_name,
+        save_dir=str(log_dir) if log_dir is not None else config.save_dir,
         tracking_uri=config.tracking_uri,
         tags=config.tags,
         log_model=config.log_model,
     )
 
 
-LOGGER_FACTORY: Dict[str, LoggerBuilder] = {
+LOGGER_FACTORY = {
     "dvclive": create_dvclive_logger,
     "csv": create_csv_logger,
     "tensorboard": create_tensorboard_logger,
@@ -205,9 +141,7 @@ LOGGER_FACTORY: Dict[str, LoggerBuilder] = {
 
 def build_logger(
     config: LoggerConfig,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    log_dir: Optional[data.PathLike] = None,
 ) -> Logger:
     """
     Creates a logger instance from a validated Pydantic config object.
@@ -216,19 +150,14 @@ def build_logger(
         "Building logger with config: \n{}",
         lambda: config.to_yaml_string(),
     )
-    logger_type = config.name
+    logger_type = config.logger_type
 
     if logger_type not in LOGGER_FACTORY:
         raise ValueError(f"Unknown logger type: {logger_type}")
 
     creation_func = LOGGER_FACTORY[logger_type]
 
-    return creation_func(
-        config,
-        log_dir=log_dir,
-        experiment_name=experiment_name,
-        run_name=run_name,
-    )
+    return creation_func(config, log_dir=log_dir)
 
 
 def get_image_plotter(logger: Logger):
@@ -244,8 +173,8 @@ def get_image_plotter(logger: Logger):
         def plot_figure(name, figure, step):
             image = _convert_figure_to_image(figure)
             return logger.experiment.log_image(
-                logger.run_id,
-                image,
+                run_id=logger.run_id,
+                image=image,
                 key=name,
                 step=step,
             )

src/batdetect2/train/train.py

@@ -1,9 +1,8 @@
 from collections.abc import Sequence
-from pathlib import Path
 from typing import List, Optional
 
 import torch
-from lightning import Trainer, seed_everything
+from lightning import Trainer
 from lightning.pytorch.callbacks import Callback, ModelCheckpoint
 from loguru import logger
 from soundevent import data
@@ -46,8 +45,6 @@ __all__ = [
     "train",
 ]
 
-DEFAULT_CHECKPOINT_DIR: Path = Path("outputs") / "checkpoints"
-
 
 def train(
     train_annotations: Sequence[data.ClipAnnotation],
@@ -56,15 +53,9 @@ def train(
     model_path: Optional[data.PathLike] = None,
     train_workers: Optional[int] = None,
     val_workers: Optional[int] = None,
-    checkpoint_dir: Optional[Path] = None,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
-    seed: Optional[int] = None,
+    checkpoint_dir: Optional[data.PathLike] = None,
+    log_dir: Optional[data.PathLike] = None,
 ):
-    if seed is not None:
-        seed_everything(seed)
-
     config = config or FullTrainingConfig()
 
     targets = build_targets(config.targets)
@@ -116,8 +107,6 @@ def train(
         targets=targets,
         checkpoint_dir=checkpoint_dir,
         log_dir=log_dir,
-        experiment_name=experiment_name,
-        run_name=run_name,
     )
 
     logger.info("Starting main training loop...")
@@ -145,32 +134,17 @@ def build_trainer_callbacks(
     targets: TargetProtocol,
     preprocessor: PreprocessorProtocol,
     config: EvaluationConfig,
-    checkpoint_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    checkpoint_dir: Optional[data.PathLike] = None,
 ) -> List[Callback]:
     if checkpoint_dir is None:
-        checkpoint_dir = DEFAULT_CHECKPOINT_DIR
-
-    filename = "best-{epoch:02d}-{val_loss:.0f}"
-
-    if run_name is not None:
-        filename = f"run_{run_name}_{filename}"
-
-    if experiment_name is not None:
-        filename = f"experiment_{experiment_name}_{filename}"
-
-    model_checkpoint = ModelCheckpoint(
-        dirpath=str(checkpoint_dir),
-        save_top_k=1,
-        filename=filename,
-        monitor="total_loss/val",
-    )
-
-    model_checkpoint.CHECKPOINT_EQUALS_CHAR = "_"  # type: ignore
+        checkpoint_dir = "outputs/checkpoints"
 
     return [
-        model_checkpoint,
+        ModelCheckpoint(
+            dirpath=str(checkpoint_dir),
+            save_top_k=1,
+            monitor="total_loss/val",
+        ),
         ValidationMetrics(
             metrics=[
                 DetectionAveragePrecision(),
@@ -188,22 +162,15 @@ def build_trainer_callbacks(
 def build_trainer(
     conf: FullTrainingConfig,
     targets: TargetProtocol,
-    checkpoint_dir: Optional[Path] = None,
-    log_dir: Optional[Path] = None,
-    experiment_name: Optional[str] = None,
-    run_name: Optional[str] = None,
+    checkpoint_dir: Optional[data.PathLike] = None,
+    log_dir: Optional[data.PathLike] = None,
 ) -> Trainer:
     trainer_conf = conf.train.trainer
     logger.opt(lazy=True).debug(
         "Building trainer with config: \n{config}",
         config=lambda: trainer_conf.to_yaml_string(exclude_none=True),
     )
-    train_logger = build_logger(
-        conf.train.logger,
-        log_dir=log_dir,
-        experiment_name=experiment_name,
-        run_name=run_name,
-    )
+    train_logger = build_logger(conf.train.logger, log_dir=log_dir)
 
     train_logger.log_hyperparams(
         conf.model_dump(
@@ -220,7 +187,6 @@ def build_trainer(
             config=conf.evaluation,
             preprocessor=build_preprocessor(conf.preprocess),
             checkpoint_dir=checkpoint_dir,
-            experiment_name=train_logger.name,
         ),
     )