diff --git a/README.md b/README.md index fc807019..3a270ed0 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,7 @@ flowchart LR - **Annotation Tools**: Create, upload, and manage annotations (segmentations, labels, measurements) with ease - **Experiment Tracking**: Seamless support for experiment management via MLflow integration - **One-line Trainers**: Train segmentation, classification, and detection models with built-in PyTorch Lightning trainers, skipping the dataset class, training loop, and logging setup +- **Model Benchmarking**: Compare several trainers against the same dataset and split, and get a ranked leaderboard of their performance - **DICOM Support**: Native handling of DICOM files, including powerful anonymization capabilities during upload to protect patient privacy - **Multi-format Support**: Robust support for a wide range of medical imaging formats: PNG, JPEG, NIfTI (NIfTI/NRRD), DICOMs and more diff --git a/datamint/dataset/base.py b/datamint/dataset/base.py index ac2ef95f..c6887a23 100644 --- a/datamint/dataset/base.py +++ b/datamint/dataset/base.py @@ -1201,15 +1201,23 @@ def collate_fn(batch: list[dict]) -> dict: if all(s == shapes[0] for s in shapes): collated[key] = torch.stack(values) else: - _LOGGER.warning(f"Different shapes for {key}: {shapes}") - collated[key] = values + raise ValueError( + f"Cannot collate batch: tensors under key '{key}' have mismatched " + f"shapes {shapes}. This usually means the dataset/trainer has no " + "resize transform configured for a project with variable-resolution " + "resources (e.g. pass image_size=... to the trainer)." + ) elif isinstance(values[0], np.ndarray): shapes = [a.shape for a in values] if all(s == shapes[0] for s in shapes): collated[key] = np.stack(values) else: - _LOGGER.warning(f"Different shapes for {key}: {shapes}") - collated[key] = values + raise ValueError( + f"Cannot collate batch: arrays under key '{key}' have mismatched " + f"shapes {shapes}. This usually means the dataset/trainer has no " + "resize transform configured for a project with variable-resolution " + "resources (e.g. pass image_size=... to the trainer)." + ) else: collated[key] = values diff --git a/datamint/lightning/__init__.py b/datamint/lightning/__init__.py index d1762ab5..74fa8bc8 100644 --- a/datamint/lightning/__init__.py +++ b/datamint/lightning/__init__.py @@ -1,5 +1,6 @@ """Datamint Lightning integration.""" +from .benchmark import Benchmark from .datamodule import DatamintDataModule from .trainers import ( BaseTrainer, @@ -20,6 +21,7 @@ __all__ = [ "BaseTrainer", + "Benchmark", "ClassificationTrainer", "DatamintDataModule", "DeepLabV3PlusTrainer", diff --git a/datamint/lightning/benchmark.py b/datamint/lightning/benchmark.py new file mode 100644 index 00000000..4e2eaa9a --- /dev/null +++ b/datamint/lightning/benchmark.py @@ -0,0 +1,270 @@ +"""Benchmark: compare several trainers against one shared dataset. + +Not a new training system, each trainer runs its own complete +:meth:`~datamint.lightning.trainers.base_trainer.BaseTrainer.fit` pipeline. +``Benchmark`` guarantees every trainer sees the same dataset and split, and +collects their results into a single ranked leaderboard. +""" + +from __future__ import annotations + +import importlib +import logging +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Any + +import pandas as pd +import yaml + +from datamint.lightning.trainers.base_trainer import BaseTrainer +from datamint.lightning.trainers.classification_trainer import ClassificationTrainer +from datamint.lightning.trainers.detection_trainer import DetectionTrainer +from datamint.lightning.trainers.seg2d_trainer import SemanticSegmentation2DTrainer +from datamint.lightning.trainers.vol_seg_trainer import VolumeSegmentationTrainer + +if TYPE_CHECKING: + from datamint.dataset.base import DatamintBaseDataset + +_LOGGER = logging.getLogger(__name__) + +TrainerSpec = tuple[type[BaseTrainer], dict[str, Any]] + +# Common task-family ancestors. A trainer that doesn't descend from any of +# these (e.g. NNUNetTrainer, which subclasses BaseTrainer directly) falls +# back to being its own family -- it can only be benchmarked against other +# instances of that same class. +_TASK_FAMILIES: tuple[type[BaseTrainer], ...] = ( + ClassificationTrainer, + SemanticSegmentation2DTrainer, + VolumeSegmentationTrainer, + DetectionTrainer, +) + + +def _task_family(trainer_cls: type[BaseTrainer]) -> type[BaseTrainer]: + """Return the task-family root for *trainer_cls*, or the class itself if none match.""" + for family in _TASK_FAMILIES: + if issubclass(trainer_cls, family): + return family + return trainer_cls + + +def _to_scalar(value: Any) -> float | None: + if value is None: + return None + if hasattr(value, 'item'): + return value.item() + return value + + +def _utc_now_isoformat() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _class_dotted_path(cls: type) -> str: + return f"{cls.__module__}.{cls.__qualname__}" + + +def _resolve_trainer_class(dotted_path: str) -> type[BaseTrainer]: + module_path, _, class_name = dotted_path.rpartition('.') + if not module_path: + raise ValueError(f"Invalid trainer class path {dotted_path!r}; expected 'module.ClassName'.") + module = importlib.import_module(module_path) + cls = getattr(module, class_name) + if not (isinstance(cls, type) and issubclass(cls, BaseTrainer)): + raise TypeError(f"{dotted_path!r} does not resolve to a BaseTrainer subclass.") + return cls + + +class Benchmark: + """Train several trainers sequentially against one shared dataset/split and rank them. + + Every trainer is constructed with the same dataset and the + same ``split_as_of_timestamp``, so they all resolve the identical + project-scoped split assignments. This requires the dataset's project to + already have split assignments. + + All trainers must share one task-family ancestor (classification, + 2-D segmentation, volume segmentation, or detection). + A trainer that doesn't descend from any of those (e.g. ``NNUNetTrainer``) + is treated as its own family and can only be + benchmarked against other instances of that same class. + + Args: + dataset: Pre-built dataset shared across every trainer. + trainers: ``(TrainerClass, kwargs)`` pairs. Each ``kwargs`` dict must + include a unique, explicit ``model_name``. + split_as_of_timestamp: Timestamp pinning the project-scoped split + assignments every trainer resolves against. When omitted, + captured once (UTC now) at construction. + main_metric: Bare metric key (e.g. ``'accuracy'``, ``'dice'``) shared + by every trainer's ``val/`` and ``test/`` logged metric names. + Used as the leaderboard's sort column. When omitted, falls back + to the first trainer's ``_monitor_metric()``. + main_metric_mode: ``'max'`` (default) or ``'min'``. Sort direction + for *main_metric*. Ignored when main_metric is omitted (the + mode from ``_monitor_metric()`` is used instead). + **shared_kwargs: Extra kwargs forwarded to every trainer's + constructor (e.g. ``max_epochs=``). Per-trainer kwargs win on + conflict. + + Example:: + + from datamint import ImageDataset + from datamint.lightning import Benchmark + from datamint.lightning.trainers import UNetPPTrainer, DeepLabV3PlusTrainer + + dataset = ImageDataset(project='BUS_Segmentation', return_segmentations=True) + + bench = Benchmark( + dataset=dataset, + trainers=[ + (UNetPPTrainer, {'encoder_name': 'resnet34', 'model_name': 'unetpp_r34'}), + (DeepLabV3PlusTrainer, {'model_name': 'deeplabv3plus'}), + ], + main_metric='dice', + ) + leaderboard = bench.run() + + bench.save_config('benchmark.yaml') + # later: Benchmark.load_from_file('benchmark.yaml', dataset=dataset).run() + """ + + def __init__( + self, + dataset: DatamintBaseDataset, + trainers: list[TrainerSpec], + split_as_of_timestamp: str | None = None, + main_metric: str | None = None, + main_metric_mode: str = 'max', + **shared_kwargs: Any, + ) -> None: + if not trainers: + raise ValueError("'trainers' must contain at least one (TrainerClass, kwargs) pair.") + if main_metric_mode not in ('max', 'min'): + raise ValueError(f"main_metric_mode must be 'max' or 'min', got {main_metric_mode!r}.") + + self.dataset = dataset + self._specs: list[TrainerSpec] = list(trainers) + self.split_as_of_timestamp = split_as_of_timestamp or _utc_now_isoformat() + self.main_metric = main_metric + self.main_metric_mode = main_metric_mode + self._shared_kwargs = shared_kwargs + + self._validate_same_task_family() + self._validate_unique_model_names() + + def _validate_same_task_family(self) -> None: + families = {_task_family(cls) for cls, _ in self._specs} + if len(families) > 1: + names = sorted(f.__name__ for f in families) + raise ValueError( + f"Cannot benchmark trainers from different task families: {names}. " + "All trainers in a Benchmark must share a common task-family ancestor." + ) + + def _validate_unique_model_names(self) -> None: + names = [kwargs.get('model_name') for _, kwargs in self._specs] + if any(name is None for name in names): + raise ValueError("Every trainer spec needs an explicit 'model_name' for benchmarking.") + if len(set(names)) != len(names): + dupes = sorted({name for name in names if names.count(name) > 1}) + raise ValueError( + f"Duplicate model_name(s) across trainer specs: {dupes}. " + "Each trainer needs a unique model_name." + ) + + @staticmethod + def _find_checkpoint_callback(lightning_trainer: Any) -> Any | None: + for callback in lightning_trainer.callbacks: + if hasattr(callback, 'registered_model_info'): + return callback + return None + + def _metric_name_and_mode(self, trainer: BaseTrainer) -> tuple[str, str]: + if self.main_metric is not None: + return f"val/{self.main_metric}", self.main_metric_mode + return trainer._monitor_metric() + + def run(self) -> pd.DataFrame: + """Train every trainer sequentially and return a leaderboard DataFrame.""" + rows: list[dict[str, Any]] = [] + test_col = val_col = None + mode = self.main_metric_mode + + for trainer_cls, kwargs in self._specs: + model_name = kwargs['model_name'] + _LOGGER.info("Benchmark: training '%s' (%s)...", model_name, trainer_cls.__name__) + + trainer = trainer_cls( + dataset=self.dataset, + split_as_of_timestamp=self.split_as_of_timestamp, + **{**self._shared_kwargs, **kwargs}, + ) + result = trainer.fit() + + if result.get('paused'): + raise RuntimeError( + f"Training for '{model_name}' was interrupted before completion " + f"(run_id={result['run_id']}); Benchmark does not support resuming mid-run." + ) + + val_col, mode = self._metric_name_and_mode(trainer) + test_col = f"test/{val_col.removeprefix('val/')}" + + test_results = result['test_results'][0] if result['test_results'] else {} + checkpoint_cb = self._find_checkpoint_callback(trainer._lightning_trainer) + registered_info = getattr(checkpoint_cb, 'registered_model_info', None) + + rows.append({ + 'model_name': model_name, + 'trainer_class': trainer_cls.__name__, + val_col: _to_scalar(getattr(checkpoint_cb, 'best_model_score', None)), + test_col: test_results.get(test_col), + 'run_id': trainer._lightning_trainer.logger.run_id, + 'registered_version': registered_info.version if registered_info is not None else None, + }) + + leaderboard = pd.DataFrame(rows) + sort_col = test_col if leaderboard[test_col].notna().any() else val_col + return leaderboard.sort_values(sort_col, ascending=(mode == 'min')).reset_index(drop=True) + + def save_config(self, path: str) -> None: + """Save this benchmark's trainer specs and settings as a YAML file. """ + + config = { + 'split_as_of_timestamp': self.split_as_of_timestamp, + 'main_metric': self.main_metric, + 'main_metric_mode': self.main_metric_mode, + 'shared_kwargs': self._shared_kwargs, + 'trainers': [ + {'trainer_class': _class_dotted_path(cls), 'kwargs': kwargs} + for cls, kwargs in self._specs + ], + } + with open(path, 'w') as f: + yaml.safe_dump(config, f, sort_keys=False) + + @classmethod + def load_from_file(cls, path: str, dataset: DatamintBaseDataset) -> Benchmark: + """Reconstruct a :class:`Benchmark` from a file saved by :meth:`save_config`. + + Args: + path: Path to the YAML file. + dataset: Freshly-built dataset to benchmark against. + """ + with open(path) as f: + config = yaml.safe_load(f) + + trainers = [ + (_resolve_trainer_class(spec['trainer_class']), spec['kwargs']) + for spec in config['trainers'] + ] + return cls( + dataset=dataset, + trainers=trainers, + split_as_of_timestamp=config.get('split_as_of_timestamp'), + main_metric=config.get('main_metric'), + main_metric_mode=config.get('main_metric_mode', 'max'), + **config.get('shared_kwargs', {}), + ) diff --git a/docs/source/datamint.lightning_api.rst b/docs/source/datamint.lightning_api.rst index 707b127e..f57110cb 100644 --- a/docs/source/datamint.lightning_api.rst +++ b/docs/source/datamint.lightning_api.rst @@ -121,6 +121,17 @@ Specialized Trainers :undoc-members: :show-inheritance: +Benchmark +--------- + +Compares several trainers sequentially against one shared dataset and split, +and returns a ranked leaderboard. See :doc:`trainer_api` for a walkthrough. + +.. automodule:: datamint.lightning.benchmark + :members: + :undoc-members: + :show-inheritance: + Lightning Modules ----------------- diff --git a/docs/source/trainer_api.rst b/docs/source/trainer_api.rst index ededec3f..b6ae35f1 100644 --- a/docs/source/trainer_api.rst +++ b/docs/source/trainer_api.rst @@ -334,6 +334,64 @@ or wrap the final model in a ``DatamintModel`` afterwards. Segmentation batches expose masks in ``batch["segmentations"]`` and include the background channel at index 0. +Benchmarking Several Trainers +------------------------------ + +``Benchmark`` runs several existing trainers sequentially against one shared dataset and split, then +returns a ranked leaderboard. + +.. code-block:: python + + from datamint import ImageDataset + from datamint.lightning import Benchmark + from datamint.lightning.trainers import UNetPPTrainer, DeepLabV3PlusTrainer, TransUNetTrainer + + dataset = ImageDataset(project="BUSI_Segmentation", return_segmentations=True) + + bench = Benchmark( + dataset=dataset, + trainers=[ + (UNetPPTrainer, {"encoder_name": "resnet34", "model_name": "unetpp_r34"}), + (DeepLabV3PlusTrainer, {"model_name": "deeplabv3plus"}), + (TransUNetTrainer, {"model_name": "transunet"}), + ], + main_metric="dice", + main_metric_mode="max", + max_epochs=20, # shared kwarg, forwarded to every trainer + ) + + leaderboard = bench.run() + + bench.save_config("benchmark.yaml") + # later, possibly in a different process/session: + # Benchmark.load_from_file("benchmark.yaml", dataset=dataset).run() + +``leaderboard`` is a ``pandas.DataFrame``, one row per trainer, sorted by +``test/{main_metric}`` (falling back to ``val/{main_metric}`` if the test column is +entirely missing), with columns ``model_name``, ``trainer_class``, +``val/{main_metric}``, ``test/{main_metric}``, ``run_id``, and +``registered_version`` -- so a specific trained model stays unambiguously +loadable afterward via ``runs://model``. + +A few rules ``Benchmark`` enforces upfront (before training anything): + +- All trainers must share one task-family ancestor: classification, 2-D + segmentation, volume segmentation, or detection. A trainer outside those + families (e.g. ``NNUNetTrainer``) can only be benchmarked against other + instances of that same class. +- Every trainer spec needs an explicit, unique ``model_name`` -- this keeps + each trainer's registered model separate in the MLflow Model Registry + instead of silently stacking versions under one shared name. +- The dataset's project must already have split assignments. ``Benchmark`` + pins one ``split_as_of_timestamp`` at construction (or an explicit override) + and forwards it to every trainer, so all of them resolve the identical + split assignments even if someone edits them mid-benchmark. + +``save_config()``/``load_from_file()`` round-trip the benchmark definition +through YAML, independent of any live dataset object. The dataset itself is +never serialized: ``load_from_file(path, dataset=...)`` requires a +freshly-built dataset, mirroring the constructor's own requirement. + Related Examples ----------------