nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/data/balanced_batch_sampler.py

@@ -0,0 +1,132 @@
+import heapq
+from functools import cached_property
+from logging import getLogger
+from typing import Any, Protocol, runtime_checkable
+
+import numpy as np
+import torch
+import torch.distributed
+from lightning_fabric.utilities.distributed import _DatasetSamplerWrapper
+from torch.utils.data import BatchSampler, Dataset, DistributedSampler
+from typing_extensions import override
+
+log = getLogger(__name__)
+
+
+def _all_gather(tensor: torch.Tensor, device: torch.device | None = None):
+    gathered = [
+        torch.zeros_like(tensor, device=device)
+        for _ in range(torch.distributed.get_world_size())
+    ]
+    _ = torch.distributed.all_gather(gathered, tensor)
+    return gathered
+
+
+# @numba.njit
+def _balanced_partition(sizes: np.ndarray, num_parts: int):
+    """
+    Greedily partition the given set by always inserting
+    the largest element into the smallest partition.
+    """
+    sort_idx = np.argsort(-sizes)  # Sort in descending order
+    heap = []
+    for idx in sort_idx[:num_parts]:
+        heap.append((sizes[idx], [idx]))
+    heapq.heapify(heap)
+    for idx in sort_idx[num_parts:]:
+        smallest_part = heapq.heappop(heap)
+        new_size = smallest_part[0] + sizes[idx]
+        new_idx = smallest_part[1] + [idx]
+        heapq.heappush(heap, (new_size, new_idx))
+    idx_balanced = [part[1] for part in heap]
+    return idx_balanced
+
+
+@runtime_checkable
+class DatasetWithSizes(Protocol):
+    def data_sizes(self, indices: list[int]) -> np.ndarray: ...
+
+
+class BalancedBatchSampler(BatchSampler):
+    @staticmethod
+    def _ensure_supported(dataset: Any):
+        if not isinstance(dataset, Dataset):
+            raise ValueError(
+                "BalancedBatchSampler requires a dataset that implements `__getitem__`"
+            )
+
+        if not isinstance(dataset, DatasetWithSizes):
+            raise ValueError(
+                "BalancedBatchSampler requires a dataset that implements `data_sizes`"
+            )
+
+        log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
+        return dataset
+
+    @staticmethod
+    def _unwrap_dataset(dataset: Dataset) -> Dataset:
+        if isinstance(dataset, _DatasetSamplerWrapper):
+            if (data_source := getattr(dataset._sampler, "data_source", None)) is None:
+                raise ValueError("Could not unwrap dataset from _DatasetSamplerWrapper")
+            return data_source
+        return dataset
+
+    @property
+    def distributed_sampler(self):
+        if not isinstance(self.sampler, DistributedSampler):
+            raise ValueError(
+                f"Sampler must be a DistributedSampler, got {type(self.sampler)}"
+            )
+        return self.sampler
+
+    @cached_property
+    def dataset(self):
+        return self._ensure_supported(
+            self._unwrap_dataset(self.distributed_sampler.dataset)
+        )
+
+    def __init__(
+        self,
+        sampler: DistributedSampler,
+        *,
+        batch_size: int,
+        device: torch.device,
+        drop_last: bool = False,
+    ):
+        super().__init__(sampler, batch_size, drop_last=drop_last)
+
+        self._device = device
+
+        log.info(
+            f"Created BalancedBatchSampler with {sampler=}, {batch_size=}, {drop_last=}"
+        )
+
+    @staticmethod
+    def _dist_enabled():
+        return torch.distributed.is_available() and torch.distributed.is_initialized()
+
+    @override
+    def __iter__(self):
+        if not self._dist_enabled():
+            yield from super().__iter__()
+            return
+
+        for batch_idx in super().__iter__():
+            sizes = self.dataset.data_sizes(batch_idx)
+            idx_sizes = torch.stack(
+                [
+                    torch.tensor(batch_idx, device=self._device),
+                    torch.tensor(sizes, device=self._device),
+                ]
+            )
+            idx_sizes_all = _all_gather(idx_sizes, device=self._device)
+            idx_sizes_all = torch.cat(idx_sizes_all, dim=-1).cpu()
+            idx_all = idx_sizes_all[0]
+            sizes_all = idx_sizes_all[1]
+
+            local_idx_balanced = _balanced_partition(
+                sizes_all.numpy(), num_parts=self.distributed_sampler.num_replicas
+            )
+            # Since DistributedSampler pads the last batch
+            # this should always have an entry for each replica.
+            yield idx_all[local_idx_balanced[self.distributed_sampler.rank]].tolist()

nshtrainer/data/transform.py

@@ -0,0 +1,67 @@
+import copy
+from collections.abc import Callable
+from typing import Any, cast
+
+from typing_extensions import TypeVar
+
+TDataset = TypeVar("TDataset", infer_variance=True)
+
+
+def transform(
+    dataset: TDataset,
+    transform: Callable[[Any], Any],
+    *,
+    deepcopy: bool = False,
+) -> TDataset:
+    """
+    Wraps a dataset with a transform function.
+
+    Args:
+        dataset: The dataset to wrap.
+        transform: The transform function to apply to each item.
+        deepcopy: Whether to deep copy each item before applying the transform.
+    """
+
+    import wrapt
+
+    class _TransformedDataset(wrapt.ObjectProxy):
+        def __getitem__(self, idx):
+            nonlocal deepcopy, transform
+
+            data = self.__wrapped__.__getitem__(idx)
+            if deepcopy:
+                data = copy.deepcopy(data)
+            data = transform(data)
+            return data
+
+    return cast(TDataset, _TransformedDataset(dataset))
+
+
+def transform_with_index(
+    dataset: TDataset,
+    transform: Callable[[Any, int], Any],
+    *,
+    deepcopy: bool = False,
+) -> TDataset:
+    """
+    Wraps a dataset with a transform function that takes an index, in addition to the item.
+
+    Args:
+        dataset: The dataset to wrap.
+        transform: The transform function to apply to each item.
+        deepcopy: Whether to deep copy each item before applying the transform.
+    """
+
+    import wrapt
+
+    class _TransformedWithIndexDataset(wrapt.ObjectProxy):
+        def __getitem__(self, idx: int):
+            nonlocal deepcopy, transform
+
+            data = self.__wrapped__.__getitem__(idx)
+            if deepcopy:
+                data = copy.deepcopy(data)
+            data = transform(data, idx)
+            return data
+
+    return cast(TDataset, _TransformedWithIndexDataset(dataset))

nshtrainer/lr_scheduler/__init__.py

@@ -0,0 +1,18 @@
+from typing import Annotated, TypeAlias
+
+from ..config import Field
+from ._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from ._base import LRSchedulerMetadata as LRSchedulerMetadata
+from .linear_warmup_cosine import (
+    LinearWarmupCosineAnnealingLR as LinearWarmupCosineAnnealingLR,
+)
+from .linear_warmup_cosine import (
+    LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
+)
+from .reduce_lr_on_plateau import ReduceLROnPlateau as ReduceLROnPlateau
+from .reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
+
+LRSchedulerConfig: TypeAlias = Annotated[
+    LinearWarmupCosineDecayLRSchedulerConfig | ReduceLROnPlateauConfig,
+    Field(discriminator="name"),
+]

nshtrainer/lr_scheduler/_base.py

@@ -0,0 +1,101 @@
+import math
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, ClassVar, Literal, TypeAlias
+
+from lightning.pytorch.utilities.types import (
+    LRSchedulerConfigType,
+    LRSchedulerTypeUnion,
+)
+from torch.optim import Optimizer
+from typing_extensions import NotRequired, TypedDict
+
+from ..config import TypedConfig
+
+if TYPE_CHECKING:
+    from ..model.base import LightningModuleBase
+
+
+class LRSchedulerMetadata(TypedDict):
+    interval: Literal["epoch", "step"]
+    """Interval to update the learning rate."""
+
+    name: NotRequired[str | None]
+    """Name of the learning rate scheduler. Default is `None`."""
+
+    frequency: NotRequired[int]
+    """Frequency to update the learning rate. Default is `1`."""
+
+    reduce_on_plateau: NotRequired[bool]
+    """Whether to reduce the learning rate on plateau. Default is `False`."""
+
+    monitor: NotRequired[str | None]
+    """Value to monitor for reducing the learning rate on plateau. Required if `reduce_on_plateau` is `True`.
+    Default is `None`."""
+
+    strict: NotRequired[bool]
+    """Whether to enforce that the monitor exists for reducing the learning rate on plateau. Default is `True`."""
+
+
+class LRSchedulerConfigBase(TypedConfig, ABC):
+    Metadata: ClassVar[TypeAlias] = LRSchedulerMetadata
+
+    @abstractmethod
+    def metadata(self) -> LRSchedulerMetadata: ...
+
+    @abstractmethod
+    def create_scheduler_impl(
+        self,
+        optimizer: Optimizer,
+        lightning_module: "LightningModuleBase",
+        lr: float,
+    ) -> LRSchedulerTypeUnion | LRSchedulerConfigType: ...
+
+    def create_scheduler(
+        self,
+        optimizer: Optimizer,
+        lightning_module: "LightningModuleBase",
+        lr: float,
+    ) -> LRSchedulerConfigType:
+        # Create the scheduler.
+        scheduler = self.create_scheduler_impl(optimizer, lightning_module, lr)
+
+        # If the scheduler is not a `LRSchedulerConfigType`, then make it one.
+        if not isinstance(scheduler, Mapping):
+            scheduler = LRSchedulerConfigType(scheduler=scheduler)
+
+        # Update the scheduler config with the metadata (if not already present).
+        metadata = self.metadata()
+        # - `interval` has to be present.
+        if scheduler.get("interval") is None:
+            scheduler["interval"] = metadata["interval"]
+        # - `name`
+        if scheduler.get("name") is None and "name" in metadata:
+            scheduler["name"] = metadata["name"]
+        # - `frequency`
+        if scheduler.get("frequency") is None and "frequency" in metadata:
+            scheduler["frequency"] = metadata["frequency"]
+        # - `reduce_on_plateau`
+        if (
+            scheduler.get("reduce_on_plateau") is None
+            and "reduce_on_plateau" in metadata
+        ):
+            scheduler["reduce_on_plateau"] = metadata["reduce_on_plateau"]
+        # - `monitor`
+        if scheduler.get("monitor") is None and "monitor" in metadata:
+            scheduler["monitor"] = metadata["monitor"]
+        # - `strict`
+        if scheduler.get("strict") is None and "strict" in metadata:
+            scheduler["strict"] = metadata["strict"]  # type: ignore
+
+        return scheduler
+
+    def compute_num_steps_per_epoch(
+        self, lightning_module: "LightningModuleBase"
+    ) -> int:
+        trainer = lightning_module.trainer
+        # Use the Lightning trainer to convert the epoch-based values to step-based values
+        _ = trainer.estimated_stepping_batches
+        # ^ This is a hack to trigger the computation of the estimated stepping batches
+        #   and make sure that the `trainer.num_training_batches` attribute is set.
+        return math.ceil(trainer.num_training_batches / trainer.accumulate_grad_batches)

nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -0,0 +1,138 @@
+import math
+import warnings
+from typing import Literal
+
+from torch.optim import Optimizer
+from torch.optim.lr_scheduler import LRScheduler
+from typing_extensions import override
+
+from ..config import Field
+from ._base import LRSchedulerConfigBase
+
+
+class LinearWarmupCosineAnnealingLR(LRScheduler):
+    _get_lr_called_within_step: bool
+
+    def __init__(
+        self,
+        optimizer: Optimizer,
+        warmup_epochs: int,
+        max_epochs: int,
+        warmup_start_lr: float = 0.0,
+        eta_min: float = 0.0,
+        last_epoch: int = -1,
+        should_restart: bool = True,
+    ) -> None:
+        self.warmup_epochs = warmup_epochs
+        self.max_epochs = max_epochs
+        self.warmup_start_lr = warmup_start_lr
+        self.eta_min = eta_min
+        self.should_restart = should_restart
+
+        super().__init__(optimizer, last_epoch)
+
+    @override
+    def get_lr(self) -> list[float]:  # pyright: ignore[reportIncompatibleMethodOverride]
+        if not self._get_lr_called_within_step:
+            warnings.warn(
+                "To get the last learning rate computed by the scheduler, "
+                "please use `get_last_lr()`.",
+                UserWarning,
+            )
+
+        if self.last_epoch == 0:
+            return [self.warmup_start_lr] * len(self.base_lrs)
+        if self.last_epoch < self.warmup_epochs:
+            return [
+                group["lr"]
+                + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
+                for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            ]
+        if self.last_epoch == self.warmup_epochs:
+            return self.base_lrs
+
+        if not self.should_restart and self.last_epoch >= self.max_epochs:
+            return [self.eta_min] * len(self.base_lrs)
+
+        if (self.last_epoch - 1 - self.max_epochs) % (
+            2 * (self.max_epochs - self.warmup_epochs)
+        ) == 0:
+            return [
+                group["lr"]
+                + (base_lr - self.eta_min)
+                * (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs)))
+                / 2
+                for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
+            ]
+
+        return [
+            (
+                1
+                + math.cos(
+                    math.pi
+                    * (self.last_epoch - self.warmup_epochs)
+                    / (self.max_epochs - self.warmup_epochs)
+                )
+            )
+            / (
+                1
+                + math.cos(
+                    math.pi
+                    * (self.last_epoch - self.warmup_epochs - 1)
+                    / (self.max_epochs - self.warmup_epochs)
+                )
+            )
+            * (group["lr"] - self.eta_min)
+            + self.eta_min
+            for group in self.optimizer.param_groups
+        ]
+
+
+class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
+    name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
+
+    warmup_epochs: int = Field(ge=0)
+    r"""The number of epochs for the linear warmup phase.
+    The learning rate is linearly increased from `warmup_start_lr` to the initial learning rate over this number of epochs."""
+
+    max_epochs: int = Field(gt=0)
+    r"""The total number of epochs.
+    The learning rate is decayed to `min_lr` over this number of epochs."""
+
+    warmup_start_lr_factor: float = 0.0
+    r"""The initial learning rate for the linear warmup phase, as a factor of the initial learning rate.
+    The learning rate is linearly increased from this value to the initial learning rate over `warmup_epochs` epochs."""
+
+    min_lr_factor: float = 0.0
+    r"""The minimum learning rate, as a factor of the initial learning rate.
+    The learning rate is decayed to this value over `max_epochs` epochs."""
+
+    annealing: bool = False
+    r"""Whether to restart the learning rate schedule after `max_epochs` epochs.
+    If `False`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be set to `min_lr` for all subsequent epochs.
+    If `True`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be increased back to the initial learning rate over `max_epochs` epochs, and so on (this is called a cosine annealing schedule)."""
+
+    @override
+    def metadata(self) -> LRSchedulerConfigBase.Metadata:
+        return {
+            "interval": "step",
+        }
+
+    @override
+    def create_scheduler_impl(self, optimizer, lightning_module, lr):
+        num_steps_per_epoch = self.compute_num_steps_per_epoch(lightning_module)
+        warmup_steps = self.warmup_epochs * num_steps_per_epoch
+        max_steps = self.max_epochs * num_steps_per_epoch
+        warmup_start_lr = self.warmup_start_lr_factor * lr
+        min_lr = self.min_lr_factor * lr
+
+        # Create the scheduler
+        scheduler = LinearWarmupCosineAnnealingLR(
+            optimizer=optimizer,
+            warmup_epochs=warmup_steps,
+            max_epochs=max_steps,
+            warmup_start_lr=warmup_start_lr,
+            eta_min=min_lr,
+            should_restart=self.annealing,
+        )
+        return scheduler

nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -0,0 +1,73 @@
+from typing import TYPE_CHECKING, Literal, cast
+
+from torch.optim.lr_scheduler import ReduceLROnPlateau
+from typing_extensions import override
+
+from ll.lr_scheduler._base import LRSchedulerMetadata
+
+from ..model.config import MetricConfig
+from ._base import LRSchedulerConfigBase
+
+if TYPE_CHECKING:
+    from ..model.base import BaseConfig
+
+
+class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
+    """Reduce learning rate when a metric has stopped improving."""
+
+    name: Literal["reduce_lr_on_plateau"] = "reduce_lr_on_plateau"
+
+    metric: MetricConfig | None = None
+    """Metric to monitor.
+    If not provided, the primary metric of the runner will be used."""
+
+    patience: int = 10
+    r"""Number of epochs with no improvement after which learning rate will be reduced."""
+
+    factor: float = 0.1
+    r"""Factor by which the learning rate will be reduced. new_lr = lr * factor."""
+
+    min_lr: float | list[float] = 0.0
+    r"""A scalar or a list of scalars. A lower bound on the learning rate of all param groups or each group respectively."""
+
+    eps: float = 1.0e-8
+    r"""Minimal decay applied to lr. If the difference between new and old lr is smaller than eps, the update is ignored."""
+
+    cooldown: int = 0
+    r"""Number of epochs to wait before resuming normal operation after lr has been reduced."""
+
+    threshold: float = 1.0e-4
+    r"""Threshold for measuring the new optimum, to only focus on significant changes."""
+
+    threshold_mode: str = "rel"
+    r"""One of `rel`, `abs`. In `rel` mode, dynamic_threshold = best * (1 + threshold) in 'max' mode or best * (1 - threshold) in `min` mode. In `abs` mode, dynamic_threshold = best + threshold in `max` mode or best - threshold in `min` mode. Default: 'rel'."""
+
+    @override
+    def create_scheduler_impl(self, optimizer, lightning_module, lr):
+        if (metric := self.metric) is None:
+            lm_config = cast("BaseConfig", lightning_module.config)
+            assert (
+                metric := lm_config.primary_metric
+            ) is not None, "Primary metric must be provided if metric is not specified."
+
+        lr_scheduler = ReduceLROnPlateau(
+            optimizer,
+            mode=metric.mode,
+            factor=self.factor,
+            patience=self.patience,
+            threshold=self.threshold,
+            threshold_mode=self.threshold_mode,
+            cooldown=self.cooldown,
+            min_lr=self.min_lr,
+            eps=self.eps,
+        )
+        return {
+            "scheduler": lr_scheduler,
+            "monitor": metric.validation_monitor,
+        }
+
+    @override
+    def metadata(self) -> LRSchedulerMetadata:
+        return {
+            "interval": "epoch",
+        }

nshtrainer/model/__init__.py

@@ -0,0 +1,44 @@
+from typing_extensions import TypeAlias
+
+from .base import Base as Base
+from .base import LightningDataModuleBase as LightningDataModuleBase
+from .base import LightningModuleBase as LightningModuleBase
+from .config import ActSaveConfig as ActSaveConfig
+from .config import BaseConfig as BaseConfig
+from .config import BaseLoggerConfig as BaseLoggerConfig
+from .config import BaseProfilerConfig as BaseProfilerConfig
+from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
+from .config import CheckpointSavingConfig as CheckpointSavingConfig
+from .config import DirectoryConfig as DirectoryConfig
+from .config import EarlyStoppingConfig as EarlyStoppingConfig
+from .config import (
+    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
+)
+from .config import EnvironmentConfig as EnvironmentConfig
+from .config import (
+    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+)
+from .config import (
+    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
+)
+from .config import GradientClippingConfig as GradientClippingConfig
+from .config import (
+    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
+)
+from .config import LoggingConfig as LoggingConfig
+from .config import MetricConfig as MetricConfig
+from .config import ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig
+from .config import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)
+from .config import OptimizationConfig as OptimizationConfig
+from .config import PrimaryMetricConfig as PrimaryMetricConfig
+from .config import PythonLogging as PythonLogging
+from .config import ReproducibilityConfig as ReproducibilityConfig
+from .config import RunnerConfig as RunnerConfig
+from .config import SanityCheckingConfig as SanityCheckingConfig
+from .config import SeedConfig as SeedConfig
+from .config import TrainerConfig as TrainerConfig
+from .config import WandbWatchConfig as WandbWatchConfig
+
+ConfigList: TypeAlias = list[tuple[BaseConfig, type[LightningModuleBase]]]