Hyperparameter Optimization Pipelines and Supporting Tools

minerva/__init__.py

@@ -0,0 +1,10 @@
+
+import minerva
+import minerva.analysis
+import minerva.callbacks
+import minerva.data
+import minerva.losses
+import minerva.models
+import minerva.pipelines
+import minerva.transforms
+import minerva.utils

minerva/analysis/metrics/transformed_metrics.py

@@ -0,0 +1,191 @@
+import warnings
+from typing import Optional
+
+import torch
+from torchmetrics import Metric
+
+
+class CroppedMetric(Metric):
+    def __init__(
+        self,
+        target_h_size: int,
+        target_w_size: int,
+        metric: Metric,
+        dist_sync_on_step: bool = False,
+    ):
+        """
+        Initializes a new instance of CroppedMetric.
+
+        Parameters
+        ----------
+            target_h_size: int
+                The target height size.
+            target_w_size: int
+                The target width size.
+            dist_sync_on_step: bool, optional
+                Whether to synchronize metric state across processes at each step.
+                Defaults to False.
+        """
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.metric = metric
+        self.target_h_size = target_h_size
+        self.target_w_size = target_w_size
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor):
+        """
+        Updates the metric state with the predictions and targets.
+
+        Parameters
+        ----------
+            preds: torch.Tensor
+                The predicted tensor.
+            target:
+                torch.Tensor The target tensor.
+        """
+
+        preds = self.crop(preds)
+        target = self.crop(target)
+        self.metric.update(preds, target)
+
+    def compute(self) -> float:
+        """
+        Computes the cropped metric.
+
+        Returns:
+            float: The cropped metric.
+        """
+        return self.metric.compute()
+
+    def crop(self, x: torch.Tensor) -> torch.Tensor:
+        """crops the input tensor to the target size.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The cropped tensor.
+        """
+        h, w = x.shape[-2:]
+        start_h = (h - self.target_h_size) // 2
+        start_w = (w - self.target_w_size) // 2
+        end_h = start_h + self.target_h_size
+        end_w = start_w + self.target_w_size
+
+        return x[..., start_h:end_h, start_w:end_w]
+
+
+class ResizedMetric(Metric):
+    def __init__(
+        self,
+        target_h_size: Optional[int],
+        target_w_size: Optional[int],
+        metric: Metric,
+        keep_aspect_ratio: bool = False,
+        dist_sync_on_step: bool = False,
+    ):
+        """
+        Initializes a new instance of ResizeMetric.
+
+        Parameters
+        ----------
+            target_h_size: int
+                The target height size.
+            target_w_size: int
+                The target width size.
+            dist_sync_on_step: bool, optional
+                Whether to synchronize metric state across processes at each step.
+                Defaults to False.
+        """
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+
+        if target_h_size is None and target_w_size is None:
+            raise ValueError(
+                "At least one of target_h_size or target_w_size must be provided."
+            )
+
+        if (
+            target_h_size is not None and target_w_size is None
+        ) and keep_aspect_ratio is False:
+            warnings.warn(
+                "A target_w_size is not provided, but keep_aspect_ratio is set to False. keep_aspect_ratio will be set to True. If you want to resize the image to a specific width, please provide a target_w_size."
+            )
+            keep_aspect_ratio = True
+
+        if (
+            target_w_size is not None and target_h_size is None
+        ) and keep_aspect_ratio is False:
+            warnings.warn(
+                "A target_h_size is not provided, but keep_aspect_ratio is set to False. keep_aspect_ratio will be set to True. If you want to resize the image to a specific height, please provide a target_h_size."
+            )
+            keep_aspect_ratio = True
+
+        self.metric = metric
+        self.target_h_size = target_h_size
+        self.target_w_size = target_w_size
+        self.keep_aspect_ratio = keep_aspect_ratio
+
+    def update(self, preds: torch.Tensor, target: torch.Tensor):
+        """
+        Updates the metric state with the predictions and targets.
+
+        Parameters
+        ----------
+            preds: torch.Tensor
+                The predicted tensor.
+            target:
+                torch.Tensor The target tensor.
+        """
+
+        preds = self.resize(preds)
+        target = self.resize(target)
+        self.metric.update(preds, target)
+
+    def compute(self) -> float:
+        """
+        Computes the resized metric.
+
+        Returns:
+            float: The resized metric.
+        """
+        return self.metric.compute()
+
+    def resize(self, x: torch.Tensor) -> torch.Tensor:
+        """Resizes the input tensor to the target size.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            The input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            The resized tensor.
+        """
+        h, w = x.shape[-2:]
+
+        target_h_size = self.target_h_size
+        target_w_size = self.target_w_size
+        if self.keep_aspect_ratio:
+            if self.target_h_size is None:
+                scale = target_w_size / w
+                target_h_size = int(h * scale)
+            elif self.target_w_size is None:
+                scale = target_h_size / h
+                target_w_size = int(w * scale)
+        type_convert = False
+        if "LongTensor" in x.type():
+            x = x.to(torch.uint8)
+            type_convert = True
+
+        return (
+            torch.nn.functional.interpolate(x, size=(target_h_size, target_w_size))
+            if not type_convert
+            else torch.nn.functional.interpolate(
+                x, size=(target_h_size, target_w_size)
+            ).to(torch.long)
+        )

minerva/callbacks/HyperSearchCallbacks.py

@@ -0,0 +1,108 @@
+import os
+import shutil
+import tempfile
+from pathlib import Path
+
+import lightning.pytorch as L
+from ray import train
+from ray._private.usage.usage_lib import TagKey, record_extra_usage_tag
+from ray.train import Checkpoint
+
+
+class TrainerReportOnIntervalCallback(L.Callback):
+
+    CHECKPOINT_NAME = "checkpoint.ckpt"
+
+    def __init__(self, interval: int = 1) -> None:
+        super().__init__()
+        self.trial_name = train.get_context().get_trial_name()
+        self.local_rank = train.get_context().get_local_rank()
+        self.tmpdir_prefix = Path(tempfile.gettempdir(), self.trial_name).as_posix()
+        self.interval = interval
+        self.step = 0
+        if os.path.isdir(self.tmpdir_prefix) and self.local_rank == 0:
+            shutil.rmtree(self.tmpdir_prefix)
+
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_train_epoch_end(
+        self, trainer: L.Trainer, pl_module: L.LightningModule
+    ) -> None:
+
+        # Fetch metrics
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+
+        # (Optional) Add customized metrics
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        tmpdir = Path(self.tmpdir_prefix, str(trainer.current_epoch)).as_posix()
+        os.makedirs(tmpdir, exist_ok=True)
+
+        if self.step % self.interval == 0:
+
+            # Save checkpoint to local
+            ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+            trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+            # Report to train session
+            checkpoint = Checkpoint.from_directory(tmpdir)
+            train.report(metrics=metrics, checkpoint=checkpoint)
+        else:
+            train.report(metrics=metrics)
+
+        # Add a barrier to ensure all workers finished reporting here
+        trainer.strategy.barrier()
+
+        if self.local_rank == 0:
+            shutil.rmtree(tmpdir)
+
+        self.step += 1
+
+
+class TrainerReportKeepOnlyLastCallback(L.Callback):
+
+    CHECKPOINT_NAME = "checkpoint.ckpt"
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.trial_name = train.get_context().get_trial_name()
+        self.local_rank = train.get_context().get_local_rank()
+        self.tmpdir_prefix = Path(tempfile.gettempdir(), self.trial_name).as_posix()
+        if os.path.isdir(self.tmpdir_prefix) and self.local_rank == 0:
+            shutil.rmtree(self.tmpdir_prefix)
+
+        record_extra_usage_tag(TagKey.TRAIN_LIGHTNING_RAYTRAINREPORTCALLBACK, "1")
+
+    def on_train_epoch_end(
+        self, trainer: L.Trainer, pl_module: L.LightningModule
+    ) -> None:
+        # Fetch metrics
+        metrics = trainer.callback_metrics
+        metrics = {k: v.item() for k, v in metrics.items()}
+
+        # (Optional) Add customized metrics
+        metrics["epoch"] = trainer.current_epoch
+        metrics["step"] = trainer.global_step
+
+        tmpdir = Path(self.tmpdir_prefix, "last").as_posix()
+        os.makedirs(tmpdir, exist_ok=True)
+
+        # Delete previous checkpoint
+        if os.path.isdir(tmpdir):
+            shutil.rmtree(tmpdir)
+
+        # Save checkpoint to local
+        ckpt_path = Path(tmpdir, self.CHECKPOINT_NAME).as_posix()
+        trainer.save_checkpoint(ckpt_path, weights_only=False)
+
+        # Report to train session
+        checkpoint = Checkpoint.from_directory(tmpdir)
+        train.report(metrics=metrics, checkpoint=checkpoint)
+
+        # Add a barrier to ensure all workers finished reporting here
+        trainer.strategy.barrier()
+
+        if self.local_rank == 0:
+            shutil.rmtree(tmpdir)

minerva/data/datasets/supervised_dataset.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 import numpy as np
 
@@ -15,7 +15,7 @@ class SupervisedReconstructionDataset(SimpleDataset):
     Usually, both input and target data have the same shape.
 
     This dataset is useful for supervised tasks such as image reconstruction,
-    segmantic segmentation, and object detection, where the input data is the
+    semantic segmentation, and object detection, where the input data is the
     original data and the target is a mask or a segmentation map.
 
     Examples
@@ -45,7 +45,12 @@ class SupervisedReconstructionDataset(SimpleDataset):
         ```
     """
 
-    def __init__(self, readers: List[_Reader], transforms: Optional[_Transform] = None):
+    def __init__(
+        self,
+        readers: List[_Reader],
+        transforms: Optional[_Transform] = None,
+        support_context_transforms: bool = False,
+    ):
         """A simple dataset class for supervised reconstruction tasks.
 
         Parameters
@@ -62,12 +67,13 @@ def __init__(self, readers: List[_Reader], transforms: Optional[_Transform] = No
             AssertionError: If the number of readers is not exactly 2.
         """
         super().__init__(readers, transforms)
+        self.support_context_transforms = support_context_transforms
 
         assert (
             len(self.readers) == 2
         ), "SupervisedReconstructionDataset requires exactly 2 readers"
 
-    def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
+    def __getitem__(self, index: int) -> Tuple[Any, Any]:
         """Load data from sources and apply specified transforms. The same
         transform is applied to both input and target data.
 
@@ -78,10 +84,29 @@ def __getitem__(self, index: int) -> Tuple[np.ndarray, np.ndarray]:
 
         Returns
         -------
-        Tuple[np.ndarray, np.ndarray]
-            A tuple containing two numpy arrays representing the data.
+        Tuple[Any, Any]
+            A tuple containing two elements: the input data and the target data.
 
         """
-        data = super().__getitem__(index)
-
-        return (data[0], data[1])
+        if not self.support_context_transforms:
+            data = super().__getitem__(index)
+
+            return (data[0], data[1])
+        else:
+
+            data = []
+
+            # For each reader and transform, read the data and apply the transform.
+            # Then, append the transformed data to the list of data.
+            for reader, transform in zip(reversed(self.readers), self.transforms):
+                sample = reader[index]
+                # Apply the transform if it is not None
+                if transform is not None:
+                    sample = transform(sample)
+                data.append(sample)
+            # Return the list of transformed data or a single sample if return_single
+            # is True and there is only one reader.
+            if self.return_single:
+                return data[1]
+            else:
+                return tuple(reversed(data))