nshtrainer 1.0.0b11__tar.gz → 1.0.0b13__tar.gz

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 1.0.0b11
+Version: 1.0.0b13
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "1.0.0-beta11"
+version = "1.0.0-beta13"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

nshtrainer-1.0.0b13/src/nshtrainer/callbacks/lr_monitor.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Literal
+
+from lightning.pytorch.callbacks import LearningRateMonitor
+
+from .base import CallbackConfigBase
+
+
+class LearningRateMonitorConfig(CallbackConfigBase):
+    logging_interval: Literal["step", "epoch"] | None = None
+    """
+    Set to 'epoch' or 'step' to log 'lr' of all optimizers at the same interval, set to None to log at individual interval according to the 'interval' key of each scheduler. Defaults to None.
+    """
+
+    log_momentum: bool = False
+    """
+    Option to also log the momentum values of the optimizer, if the optimizer has the 'momentum' or 'betas' attribute. Defaults to False.
+    """
+
+    log_weight_decay: bool = False
+    """
+    Option to also log the weight decay values of the optimizer. Defaults to False.
+    """
+
+    def create_callbacks(self, trainer_config):
+        yield LearningRateMonitor(
+            logging_interval=self.logging_interval,
+            log_momentum=self.log_momentum,
+            log_weight_decay=self.log_weight_decay,
+        )

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/__init__.py

@@ -132,10 +132,8 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         GradientClippingConfig as GradientClippingConfig,
     )
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
     )
     from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
     from nshtrainer.util._environment_info import (
@@ -325,6 +323,10 @@ else:
             ).LastCheckpointStrategyConfig
         if name == "LeakyReLUNonlinearityConfig":
             return importlib.import_module("nshtrainer.nn").LeakyReLUNonlinearityConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LinearWarmupCosineDecayLRSchedulerConfig":
             return importlib.import_module(
                 "nshtrainer.lr_scheduler"
@@ -333,8 +335,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MLPConfig":
             return importlib.import_module("nshtrainer.nn").MLPConfig
         if name == "MetricConfig":
@@ -349,10 +349,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).OptimizationConfig
         if name == "OptimizerConfigBase":
             return importlib.import_module("nshtrainer.optimizer").OptimizerConfigBase
         if name == "PReLUConfig":
@@ -373,10 +369,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.lr_scheduler"
             ).ReduceLROnPlateauConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/callbacks/__init__.py

@@ -62,6 +62,9 @@ if TYPE_CHECKING:
         CheckpointMetadata as CheckpointMetadata,
     )
     from nshtrainer.callbacks.early_stopping import MetricConfig as MetricConfig
+    from nshtrainer.callbacks.lr_monitor import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
 else:
 
     def __getattr__(name):
@@ -115,6 +118,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.callbacks"
@@ -167,6 +174,7 @@ from . import ema as ema
 from . import finite_checks as finite_checks
 from . import gradient_skipping as gradient_skipping
 from . import log_epoch as log_epoch
+from . import lr_monitor as lr_monitor
 from . import norm_logging as norm_logging
 from . import print_table as print_table
 from . import rlp_sanity_checks as rlp_sanity_checks

nshtrainer-1.0.0b13/src/nshtrainer/configs/callbacks/lr_monitor/__init__.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from typing import TYPE_CHECKING
+
+# Config/alias imports
+
+if TYPE_CHECKING:
+    from nshtrainer.callbacks.lr_monitor import CallbackConfigBase as CallbackConfigBase
+    from nshtrainer.callbacks.lr_monitor import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "CallbackConfigBase":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).CallbackConfigBase
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.callbacks.lr_monitor"
+            ).LearningRateMonitorConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+# Submodule exports

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/trainer/__init__.py

@@ -9,6 +9,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from nshtrainer.trainer import TrainerConfig as TrainerConfig
     from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
+    from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
     from nshtrainer.trainer._config import (
         BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
     )
@@ -39,20 +40,21 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
     from nshtrainer.trainer._config import (
         LogEpochCallbackConfig as LogEpochCallbackConfig,
     )
     from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
     from nshtrainer.trainer._config import MetricConfig as MetricConfig
     from nshtrainer.trainer._config import (
-        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
     )
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
-    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
         RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
     )
@@ -75,6 +77,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).ActSaveLoggerConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).BaseLoggerConfig
         if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +125,28 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer.trainer._config").MetricConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
+            ).NormLoggingCallbackConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
+            ).OnExceptionCheckpointCallbackConfig
         if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).RLPSanityChecksCallbackConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/src/nshtrainer/configs/trainer/_config/__init__.py

@@ -8,6 +8,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from nshtrainer.trainer._config import ActSaveLoggerConfig as ActSaveLoggerConfig
+    from nshtrainer.trainer._config import BaseLoggerConfig as BaseLoggerConfig
     from nshtrainer.trainer._config import (
         BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
     )
@@ -38,20 +39,21 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import (
+        LearningRateMonitorConfig as LearningRateMonitorConfig,
+    )
     from nshtrainer.trainer._config import (
         LogEpochCallbackConfig as LogEpochCallbackConfig,
     )
     from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
-    from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
     from nshtrainer.trainer._config import MetricConfig as MetricConfig
     from nshtrainer.trainer._config import (
-        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+        NormLoggingCallbackConfig as NormLoggingCallbackConfig,
     )
-    from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
-    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
-        ReproducibilityConfig as ReproducibilityConfig,
+        OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
     )
+    from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
     from nshtrainer.trainer._config import (
         RLPSanityChecksCallbackConfig as RLPSanityChecksCallbackConfig,
     )
@@ -75,6 +77,10 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).ActSaveLoggerConfig
+        if name == "BaseLoggerConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).BaseLoggerConfig
         if name == "BestCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +125,28 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LastCheckpointCallbackConfig
+        if name == "LearningRateMonitorConfig":
+            return importlib.import_module(
+                "nshtrainer.trainer._config"
+            ).LearningRateMonitorConfig
         if name == "LogEpochCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).LogEpochCallbackConfig
-        if name == "LoggingConfig":
-            return importlib.import_module("nshtrainer.trainer._config").LoggingConfig
         if name == "MetricConfig":
             return importlib.import_module("nshtrainer.trainer._config").MetricConfig
-        if name == "OnExceptionCheckpointCallbackConfig":
+        if name == "NormLoggingCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OnExceptionCheckpointCallbackConfig
-        if name == "OptimizationConfig":
+            ).NormLoggingCallbackConfig
+        if name == "OnExceptionCheckpointCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
-            ).OptimizationConfig
+            ).OnExceptionCheckpointCallbackConfig
         if name == "RLPSanityChecksCallbackConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).RLPSanityChecksCallbackConfig
-        if name == "ReproducibilityConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).ReproducibilityConfig
         if name == "SanityCheckingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"

nshtrainer-1.0.0b13/src/nshtrainer/data/datamodule.py

@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import Any, Generic, cast
+
+import nshconfig as C
+import torch
+from lightning.pytorch import LightningDataModule
+from typing_extensions import Never, TypeVar, deprecated, override
+
+from ..model.mixins.callback import CallbackRegistrarModuleMixin
+from ..model.mixins.debug import _DebugModuleMixin
+
+THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
+
+
+class LightningDataModuleBase(
+    _DebugModuleMixin,
+    CallbackRegistrarModuleMixin,
+    LightningDataModule,
+    ABC,
+    Generic[THparams],
+):
+    @property
+    @override
+    def hparams(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(THparams, super().hparams)
+
+    @property
+    @override
+    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+        hparams = cast(THparams, super().hparams_initial)
+        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
+
+    @property
+    @deprecated("Use `hparams` instead")
+    def config(self):
+        return cast(Never, self.hparams)
+
+    @classmethod
+    @abstractmethod
+    def hparams_cls(cls) -> type[THparams]: ...
+
+    @override
+    def __init__(self, hparams: THparams | Mapping[str, Any]):
+        super().__init__()
+
+        # Validate and save hyperparameters
+        hparams_cls = self.hparams_cls()
+        if isinstance(hparams, Mapping):
+            hparams = hparams_cls.model_validate(hparams)
+        elif not isinstance(hparams, hparams_cls):
+            raise TypeError(
+                f"Expected hparams to be either a Mapping or an instance of {hparams_cls}, got {type(hparams)}"
+            )
+        hparams = hparams.model_deep_validate()
+        self.save_hyperparameters(hparams)
+
+    @override
+    @classmethod
+    def load_from_checkpoint(cls, *args, **kwargs) -> Never:
+        raise ValueError("This method is not supported. Use `from_checkpoint` instead.")
+
+    @classmethod
+    def hparams_from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
+
+        if (hparams := ckpt.get(cls.CHECKPOINT_HYPER_PARAMS_KEY)) is None:
+            raise ValueError(
+                f"The checkpoint does not contain hyperparameters. It must contain the key '{cls.CHECKPOINT_HYPER_PARAMS_KEY}'."
+            )
+        if update_hparams_dict is not None:
+            hparams = update_hparams_dict(hparams)
+
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
+
+        return hparams
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        map_location: torch.serialization.MAP_LOCATION = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        # Load checkpoint
+        if isinstance(ckpt_or_path, Mapping):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location=map_location)
+
+        # Load hyperparameters from checkpoint
+        hparams = cls.hparams_from_checkpoint(
+            ckpt,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+
+        # Load datamodule from checkpoint
+        datamodule = cls(hparams)
+        if datamodule.__class__.__qualname__ in ckpt:
+            datamodule.load_state_dict(ckpt[datamodule.__class__.__qualname__])
+
+        return datamodule

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b13}/src/nshtrainer/model/base.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
+from pathlib import Path
 from typing import Any, Generic, Literal, cast
 
 import nshconfig as C
@@ -10,11 +11,13 @@ import torch
 import torch.distributed
 from lightning.pytorch import LightningModule
 from lightning.pytorch.profilers import PassThroughProfiler, Profiler
+from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 from typing_extensions import Never, TypeVar, deprecated, override
 
 from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
 from .mixins.callback import CallbackModuleMixin
-from .mixins.debug import _DebugModuleMixin, _trainer
+from .mixins.debug import _DebugModuleMixin
 from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
@@ -241,3 +244,98 @@ class LightningModuleBase(
         loss = sum((0.0 * v).sum() for v in self.parameters() if v.requires_grad)
         loss = cast(torch.Tensor, loss)
         return loss
+
+    @override
+    @classmethod
+    def load_from_checkpoint(cls, *args, **kwargs) -> Never:
+        raise ValueError("This method is not supported. Use `from_checkpoint` instead.")
+
+    @classmethod
+    def hparams_from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
+
+        if (hparams := ckpt.get(cls.CHECKPOINT_HYPER_PARAMS_KEY)) is None:
+            raise ValueError(
+                f"The checkpoint does not contain hyperparameters. It must contain the key '{cls.CHECKPOINT_HYPER_PARAMS_KEY}'."
+            )
+        if update_hparams_dict is not None:
+            hparams = update_hparams_dict(hparams)
+
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
+
+        return hparams
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        map_location: torch.serialization.MAP_LOCATION = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        # Load checkpoint
+        if isinstance(ckpt_or_path, Mapping):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location=map_location)
+
+        # Load hyperparameters from checkpoint
+        hparams = cls.hparams_from_checkpoint(
+            ckpt,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+
+        # Load model from checkpoint
+        model = cls(hparams)
+
+        # Load model state from checkpoint
+        if (
+            model._strict_loading is not None
+            and strict is not None
+            and strict != model.strict_loading
+        ):
+            raise ValueError(
+                f"You set `.load_from_checkpoint(..., strict={strict!r})` which is in conflict with"
+                f" `{cls.__name__}.strict_loading={model.strict_loading!r}. Please set the same value for both of them."
+            )
+        strict = model.strict_loading if strict is None else strict
+
+        if is_overridden("configure_model", model):
+            model.configure_model()
+
+        # give model a chance to load something
+        model.on_load_checkpoint(ckpt)
+
+        # load the state_dict on the model automatically
+
+        keys = model.load_state_dict(ckpt["state_dict"], strict=strict)
+
+        if not strict:
+            if keys.missing_keys:
+                rank_zero_warn(
+                    f"Found keys that are in the model state dict but not in the checkpoint: {keys.missing_keys}"
+                )
+            if keys.unexpected_keys:
+                rank_zero_warn(
+                    f"Found keys that are not in the model state dict but in the checkpoint: {keys.unexpected_keys}"
+                )
+
+        return model