nshtrainer 0.1.0__py3-none-any.whl

nshtrainer/callbacks/print_table.py

@@ -0,0 +1,90 @@
+import copy
+import fnmatch
+import importlib.util
+import logging
+from typing import Literal
+
+import torch
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+class PrintTableMetricsCallback(Callback):
+    """Prints a table with the metrics in columns on every epoch end."""
+
+    def __init__(
+        self,
+        metric_patterns: list[str] | None = None,
+    ) -> None:
+        self.metrics: list = []
+        self.rich_available = importlib.util.find_spec("rich") is not None
+        self.metric_patterns = metric_patterns
+
+        if not self.rich_available:
+            log.warning(
+                "rich is not installed. Please install it to use PrintTableMetricsCallback."
+            )
+
+    @override
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.rich_available:
+            return
+
+        metrics_dict = copy.copy(trainer.callback_metrics)
+        # Filter metrics based on the patterns
+        if self.metric_patterns is not None:
+            metrics_dict = {
+                key: value
+                for key, value in metrics_dict.items()
+                if any(
+                    fnmatch.fnmatch(key, pattern) for pattern in self.metric_patterns
+                )
+            }
+        self.metrics.append(metrics_dict)
+
+        from rich.console import Console
+
+        console = Console()
+        table = self.create_metrics_table()
+        console.print(table)
+
+    def create_metrics_table(self):
+        from rich.table import Table
+
+        table = Table(show_header=True, header_style="bold magenta")
+
+        # Add columns to the table based on the keys in the first metrics dictionary
+        for key in self.metrics[0].keys():
+            table.add_column(key)
+
+        # Add rows to the table based on the metrics dictionaries
+        for metric_dict in self.metrics:
+            values: list[str] = []
+            for value in metric_dict.values():
+                if torch.is_tensor(value):
+                    value = float(value.item())
+                values.append(str(value))
+            table.add_row(*values)
+
+        return table
+
+
+class PrintTableMetricsConfig(CallbackConfigBase):
+    """Configuration class for PrintTableMetricsCallback."""
+
+    name: Literal["print_table_metrics"] = "print_table_metrics"
+
+    enabled: bool = True
+    """Whether to enable the callback or not."""
+
+    metric_patterns: list[str] | None = None
+    """List of patterns to filter the metrics to be displayed. If None, all metrics are displayed."""
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield PrintTableMetricsCallback(metric_patterns=self.metric_patterns)

nshtrainer/callbacks/throughput_monitor.py

@@ -0,0 +1,56 @@
+from logging import getLogger
+from typing import Any, Literal, Protocol, TypedDict, cast, runtime_checkable
+
+from typing_extensions import NotRequired, override
+
+from ._throughput_monitor_callback import ThroughputMonitor as _ThroughputMonitor
+from .base import CallbackConfigBase
+
+log = getLogger(__name__)
+
+
+class ThroughputMonitorBatchStats(TypedDict):
+    batch_size: int
+    length: NotRequired[int | None]
+
+
+@runtime_checkable
+class SupportsThroughputMonitorModuleProtocol(Protocol):
+    def throughput_monitor_batch_stats(
+        self, batch: Any
+    ) -> ThroughputMonitorBatchStats: ...
+
+
+class ThroughputMonitor(_ThroughputMonitor):
+    def __init__(self, window_size: int = 100) -> None:
+        super().__init__(cast(Any, None), cast(Any, None), window_size=window_size)
+
+    @override
+    def setup(self, trainer, pl_module, stage):
+        if not isinstance(pl_module, SupportsThroughputMonitorModuleProtocol):
+            raise RuntimeError(
+                "The model does not implement `throughput_monitor_batch_stats`. "
+                "Please either implement this method, or do not use the `ThroughputMonitor` callback."
+            )
+
+        def batch_size_fn(batch):
+            return pl_module.throughput_monitor_batch_stats(batch)["batch_size"]
+
+        def length_fn(batch):
+            return pl_module.throughput_monitor_batch_stats(batch).get("length")
+
+        self.batch_size_fn = batch_size_fn
+        self.length_fn = length_fn
+
+        return super().setup(trainer, pl_module, stage)
+
+
+class ThroughputMonitorConfig(CallbackConfigBase):
+    name: Literal["throughput_monitor"] = "throughput_monitor"
+
+    window_size: int = 100
+    """Number of batches to use for a rolling average."""
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield ThroughputMonitor(window_size=self.window_size)

nshtrainer/callbacks/timer.py

@@ -0,0 +1,157 @@
+import logging
+import time
+from typing import Any, Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from lightning.pytorch.utilities.types import STEP_OUTPUT
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = logging.getLogger(__name__)
+
+
+class EpochTimer(Callback):
+    def __init__(self):
+        super().__init__()
+
+        self._start_time: dict[str, float] = {}
+        self._elapsed_time: dict[str, float] = {}
+        self._total_batches: dict[str, int] = {}
+
+    @override
+    def on_train_epoch_start(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._start_time["train"] = time.monotonic()
+        self._total_batches["train"] = 0
+
+    @override
+    def on_train_batch_end(
+        self,
+        trainer: "Trainer",
+        pl_module: "LightningModule",
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+    ) -> None:
+        self._total_batches["train"] += 1
+
+    @override
+    def on_train_epoch_end(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._elapsed_time["train"] = time.monotonic() - self._start_time["train"]
+        if trainer.is_global_zero:
+            self._log_epoch_info("train")
+
+    @override
+    def on_validation_epoch_start(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._start_time["val"] = time.monotonic()
+        self._total_batches["val"] = 0
+
+    @override
+    def on_validation_batch_end(
+        self,
+        trainer: "Trainer",
+        pl_module: "LightningModule",
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._total_batches["val"] += 1
+
+    @override
+    def on_validation_epoch_end(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._elapsed_time["val"] = time.monotonic() - self._start_time["val"]
+        if trainer.is_global_zero:
+            self._log_epoch_info("val")
+
+    @override
+    def on_test_epoch_start(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._start_time["test"] = time.monotonic()
+        self._total_batches["test"] = 0
+
+    @override
+    def on_test_batch_end(
+        self,
+        trainer: "Trainer",
+        pl_module: "LightningModule",
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._total_batches["test"] += 1
+
+    @override
+    def on_test_epoch_end(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._elapsed_time["test"] = time.monotonic() - self._start_time["test"]
+        if trainer.is_global_zero:
+            self._log_epoch_info("test")
+
+    @override
+    def on_predict_epoch_start(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._start_time["predict"] = time.monotonic()
+        self._total_batches["predict"] = 0
+
+    @override
+    def on_predict_batch_end(
+        self,
+        trainer: "Trainer",
+        pl_module: "LightningModule",
+        outputs: STEP_OUTPUT,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._total_batches["predict"] += 1
+
+    @override
+    def on_predict_epoch_end(
+        self, trainer: "Trainer", pl_module: "LightningModule"
+    ) -> None:
+        self._elapsed_time["predict"] = time.monotonic() - self._start_time["predict"]
+        if trainer.is_global_zero:
+            self._log_epoch_info("predict")
+
+    def _log_epoch_info(self, stage: str) -> None:
+        if (elapsed_time := self._elapsed_time.get(stage)) is None:
+            return
+        total_batches = self._total_batches[stage]
+        log.critical(
+            f"Epoch {stage.capitalize()} Summary: Elapsed Time: {elapsed_time:.2f} seconds | "
+            f"Total Batches: {total_batches}"
+        )
+
+    @override
+    def state_dict(self) -> dict[str, Any]:
+        return {
+            "elapsed_time": self._elapsed_time,
+            "total_batches": self._total_batches,
+        }
+
+    @override
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        self._elapsed_time = state_dict["elapsed_time"]
+        self._total_batches = state_dict["total_batches"]
+
+
+class EpochTimerConfig(CallbackConfigBase):
+    name: Literal["epoch_timer"] = "epoch_timer"
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield EpochTimer()

nshtrainer/callbacks/wandb_watch.py

@@ -0,0 +1,103 @@
+from logging import getLogger
+from typing import Literal, Protocol, cast, runtime_checkable
+
+import torch.nn as nn
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks.callback import Callback
+from lightning.pytorch.loggers import WandbLogger
+from typing_extensions import override
+
+from .base import CallbackConfigBase
+
+log = getLogger(__name__)
+
+
+@runtime_checkable
+class _HasWandbLogModuleProtocol(Protocol):
+    def wandb_log_module(self) -> nn.Module | None: ...
+
+
+class WandbWatchCallback(Callback):
+    def __init__(self, config: "WandbWatchConfig"):
+        super().__init__()
+
+        self.config = config
+
+    @override
+    def on_fit_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self._on_start(trainer, pl_module)
+
+    @override
+    def on_validation_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self._on_start(trainer, pl_module)
+
+    @override
+    def on_test_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self._on_start(trainer, pl_module)
+
+    @override
+    def on_predict_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        self._on_start(trainer, pl_module)
+
+    def _on_start(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        # If not enabled, return
+        if not self.config:
+            return
+
+        # If we're in fast_dev_run, don't watch the model
+        if getattr(trainer, "fast_dev_run", False):
+            return
+
+        if (
+            logger := next(
+                (
+                    logger
+                    for logger in trainer.loggers
+                    if isinstance(logger, WandbLogger)
+                ),
+                None,
+            )
+        ) is None:
+            log.warning("Could not find wandb logger or module to log")
+            return
+
+        if getattr(pl_module, "_model_watched", False):
+            return
+
+        # Get which module to log
+        if (
+            not isinstance(pl_module, _HasWandbLogModuleProtocol)
+            or (module := pl_module.wandb_log_module()) is None
+        ):
+            module = cast(nn.Module, pl_module)
+
+        logger.watch(
+            module,
+            log=cast(str, self.config.log),
+            log_freq=self.config.log_freq,
+            log_graph=self.config.log_graph,
+        )
+        setattr(pl_module, "_model_watched", True)
+
+
+class WandbWatchConfig(CallbackConfigBase):
+    name: Literal["finite_checks"] = "finite_checks"
+
+    enabled: bool = True
+    """Enable watching the model for wandb."""
+
+    log: str | None = None
+    """Log type for wandb."""
+
+    log_graph: bool = True
+    """Whether to log the graph for wandb."""
+
+    log_freq: int = 100
+    """Log frequency for wandb."""
+
+    def __bool__(self):
+        return self.enabled
+
+    @override
+    def construct_callbacks(self, root_config):
+        yield WandbWatchCallback(self)

nshtrainer/config.py

@@ -0,0 +1,289 @@
+from collections.abc import Mapping, MutableMapping
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from pydantic import BaseModel, ConfigDict
+from pydantic import Field as Field
+from pydantic import PrivateAttr as PrivateAttr
+from typing_extensions import deprecated, override
+
+from ._config.missing import MISSING, validate_no_missing_values
+from ._config.missing import AllowMissing as AllowMissing
+from ._config.missing import MissingField as MissingField
+
+_MutableMappingBase = MutableMapping[str, Any]
+if TYPE_CHECKING:
+    _MutableMappingBase = object
+
+
+_DraftConfigContextSentinel = object()
+
+
+class TypedConfig(BaseModel, _MutableMappingBase):
+    _is_draft_config: bool = PrivateAttr(default=False)
+    """
+    Whether this config is a draft config or not.
+
+    Draft configs are configs that are not yet fully validated.
+    They allow for a nicer API when creating configs, e.g.:
+
+        ```python
+        config = MyConfig.draft()
+
+        # Set some values
+        config.a = 10
+        config.b = "hello"
+
+        # Finalize the config
+        config = config.finalize()
+        ```
+    """
+
+    repr_diff_only: ClassVar[bool] = True
+    """
+    If `True`, the repr methods will only show values for fields that are different from the default.
+    """
+
+    MISSING: ClassVar[Any] = MISSING
+    """
+    Alias for the `MISSING` constant.
+    """
+
+    model_config: ClassVar[ConfigDict] = ConfigDict(
+        # By default, Pydantic will throw a warning if a field starts with "model_",
+        # so we need to disable that warning (beacuse "model_" is a popular prefix for ML).
+        protected_namespaces=(),
+        validate_assignment=True,
+        validate_return=True,
+        validate_default=True,
+        strict=True,
+        revalidate_instances="always",
+        arbitrary_types_allowed=True,
+        extra="ignore",
+        validation_error_cause=True,
+        use_attribute_docstrings=True,
+    )
+
+    def __draft_pre_init__(self):
+        """Called right before a draft config is finalized."""
+        pass
+
+    def __post_init__(self):
+        """Called after the final config is validated."""
+        pass
+
+    @classmethod
+    @deprecated("Use `model_validate` instead.")
+    def from_dict(cls, model_dict: Mapping[str, Any]):
+        return cls.model_validate(model_dict)
+
+    def model_deep_validate(self, strict: bool = True):
+        """
+        Validate the config and all of its sub-configs.
+
+        Args:
+            config: The config to validate.
+            strict: Whether to validate the config strictly.
+        """
+        config_dict = self.model_dump(round_trip=True)
+        config = self.model_validate(config_dict, strict=strict)
+
+        # Make sure that this is not a draft config
+        if config._is_draft_config:
+            raise ValueError("Draft configs are not valid. Call `finalize` first.")
+
+        return config
+
+    @classmethod
+    def draft(cls, **kwargs):
+        config = cls.model_construct_draft(**kwargs)
+        return config
+
+    def finalize(self, strict: bool = True):
+        # This must be a draft config, otherwise we raise an error
+        if not self._is_draft_config:
+            raise ValueError("Finalize can only be called on drafts.")
+
+        # First, we call `__draft_pre_init__` to allow the config to modify itself a final time
+        self.__draft_pre_init__()
+
+        # Then, we dump the config to a dict and then re-validate it
+        return self.model_deep_validate(strict=strict)
+
+    @override
+    def model_post_init(self, __context: Any) -> None:
+        super().model_post_init(__context)
+
+        # Call the `__post_init__` method if this is not a draft config
+        if __context is _DraftConfigContextSentinel:
+            return
+
+        self.__post_init__()
+
+        # After `_post_init__` is called, we perform the final round of validation
+        self.model_post_init_validate()
+
+    def model_post_init_validate(self):
+        validate_no_missing_values(self)
+
+    @classmethod
+    def model_construct_draft(cls, _fields_set: set[str] | None = None, **values: Any):
+        """
+        NOTE: This is a copy of the `model_construct` method from Pydantic's `Model` class,
+            with the following changes:
+            - The `model_post_init` method is called with the `_DraftConfigContext` context.
+            - The `_is_draft_config` attribute is set to `True` in the `values` dict.
+
+        Creates a new instance of the `Model` class with validated data.
+
+        Creates a new model setting `__dict__` and `__pydantic_fields_set__` from trusted or pre-validated data.
+        Default values are respected, but no other validation is performed.
+
+        !!! note
+            `model_construct()` generally respects the `model_config.extra` setting on the provided model.
+            That is, if `model_config.extra == 'allow'`, then all extra passed values are added to the model instance's `__dict__`
+            and `__pydantic_extra__` fields. If `model_config.extra == 'ignore'` (the default), then all extra passed values are ignored.
+            Because no validation is performed with a call to `model_construct()`, having `model_config.extra == 'forbid'` does not result in
+            an error if extra values are passed, but they will be ignored.
+
+        Args:
+            _fields_set: The set of field names accepted for the Model instance.
+            values: Trusted or pre-validated data dictionary.
+
+        Returns:
+            A new instance of the `Model` class with validated data.
+        """
+
+        values["_is_draft_config"] = True
+
+        m = cls.__new__(cls)
+        fields_values: dict[str, Any] = {}
+        fields_set = set()
+
+        for name, field in cls.model_fields.items():
+            if field.alias and field.alias in values:
+                fields_values[name] = values.pop(field.alias)
+                fields_set.add(name)
+            elif name in values:
+                fields_values[name] = values.pop(name)
+                fields_set.add(name)
+            elif not field.is_required():
+                fields_values[name] = field.get_default(call_default_factory=True)
+        if _fields_set is None:
+            _fields_set = fields_set
+
+        _extra: dict[str, Any] | None = None
+        if cls.model_config.get("extra") == "allow":
+            _extra = {}
+            for k, v in values.items():
+                _extra[k] = v
+        object.__setattr__(m, "__dict__", fields_values)
+        object.__setattr__(m, "__pydantic_fields_set__", _fields_set)
+        if not cls.__pydantic_root_model__:
+            object.__setattr__(m, "__pydantic_extra__", _extra)
+
+        if cls.__pydantic_post_init__:
+            m.model_post_init(_DraftConfigContextSentinel)
+            # update private attributes with values set
+            if (
+                hasattr(m, "__pydantic_private__")
+                and m.__pydantic_private__ is not None
+            ):
+                for k, v in values.items():
+                    if k in m.__private_attributes__:
+                        m.__pydantic_private__[k] = v
+
+        elif not cls.__pydantic_root_model__:
+            # Note: if there are any private attributes, cls.__pydantic_post_init__ would exist
+            # Since it doesn't, that means that `__pydantic_private__` should be set to None
+            object.__setattr__(m, "__pydantic_private__", None)
+
+        return m
+
+    @override
+    def __repr_args__(self):
+        # If `repr_diff_only` is `True`, we only show the fields that are different from the default.
+        if not self.repr_diff_only:
+            yield from super().__repr_args__()
+            return
+
+        # First, we get the default values for all fields.
+        default_values = self.model_construct_draft()
+
+        # Then, we compare the default values with the current values.
+        for k, v in super().__repr_args__():
+            if k is None:
+                yield k, v
+                continue
+
+            # If there is no default value or the value is different from the default, we yield it.
+            if not hasattr(default_values, k) or getattr(default_values, k) != v:
+                yield k, v
+                continue
+
+            # Otherwise, we can skip this field.
+
+    # region MutableMapping implementation
+    if not TYPE_CHECKING:
+        # This is mainly so the config can be used with lightning's hparams
+        #   transparently and without any issues.
+
+        @property
+        def _ll_dict(self):
+            return self.model_dump()
+
+        # We need to make sure every config class
+        #   is a MutableMapping[str, Any] so that it can be used
+        #   with lightning's hparams.
+        @override
+        def __getitem__(self, key: str):
+            # Key can be of the format "a.b.c"
+            #   so we need to split it into a list of keys.
+            [first_key, *rest_keys] = key.split(".")
+            value = self._ll_dict[first_key]
+
+            for key in rest_keys:
+                if isinstance(value, Mapping):
+                    value = value[key]
+                else:
+                    value = getattr(value, key)
+
+            return value
+
+        @override
+        def __setitem__(self, key: str, value: Any):
+            # Key can be of the format "a.b.c"
+            #   so we need to split it into a list of keys.
+            [first_key, *rest_keys] = key.split(".")
+            if len(rest_keys) == 0:
+                self._ll_dict[first_key] = value
+                return
+
+            # We need to traverse the keys until we reach the last key
+            #   and then set the value
+            current_value = self._ll_dict[first_key]
+            for key in rest_keys[:-1]:
+                if isinstance(current_value, Mapping):
+                    current_value = current_value[key]
+                else:
+                    current_value = getattr(current_value, key)
+
+            # Set the value
+            if isinstance(current_value, MutableMapping):
+                current_value[rest_keys[-1]] = value
+            else:
+                setattr(current_value, rest_keys[-1], value)
+
+        @override
+        def __delitem__(self, key: str):
+            # This is unsupported for this class
+            raise NotImplementedError
+
+        @override
+        def __iter__(self):
+            return iter(self._ll_dict)
+
+        @override
+        def __len__(self):
+            return len(self._ll_dict)
+
+    # endregion

nshtrainer/data/__init__.py

@@ -0,0 +1,4 @@
+from . import transform as dataset_transform
+from .balanced_batch_sampler import BalancedBatchSampler as BalancedBatchSampler
+
+_ = dataset_transform