nshtrainer 1.0.0b12__tar.gz → 1.0.0b14__tar.gz

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b14}/PKG-INFO

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

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b14}/pyproject.toml

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

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b14}/src/nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py

@@ -96,8 +96,8 @@ class OnExceptionCheckpointCallback(_OnExceptionCheckpoint):
     def on_exception(self, trainer: LightningTrainer, *args: Any, **kwargs: Any):
         # Monkey-patch the strategy instance to make the barrier operation a no-op.
         # We do this because `save_checkpoint` calls `barrier`. This is okay in most
-        # cases, but when we want to save a checkpoint in the case of an exception,
-        # `barrier` causes a deadlock. So we monkey-patch the strategy instance to
-        # make the barrier operation a no-op.
+        #   cases, but when we want to save a checkpoint in the case of an exception,
+        #   `barrier` causes a deadlock. So we monkey-patch the strategy instance to
+        #   make the barrier operation a no-op.
         with _monkey_patch_disable_barrier(trainer):
             return super().on_exception(trainer, *args, **kwargs)

nshtrainer-1.0.0b14/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.0b12 → nshtrainer-1.0.0b14}/src/nshtrainer/configs/__init__.py

@@ -9,19 +9,7 @@ from typing import TYPE_CHECKING
 if TYPE_CHECKING:
     from nshtrainer import MetricConfig as MetricConfig
     from nshtrainer import TrainerConfig as TrainerConfig
-    from nshtrainer._checkpoint.loader import (
-        BestCheckpointStrategyConfig as BestCheckpointStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import (
-        CheckpointLoadingStrategyConfig as CheckpointLoadingStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import CheckpointMetadata as CheckpointMetadata
-    from nshtrainer._checkpoint.loader import (
-        LastCheckpointStrategyConfig as LastCheckpointStrategyConfig,
-    )
-    from nshtrainer._checkpoint.loader import (
-        UserProvidedPathCheckpointStrategyConfig as UserProvidedPathCheckpointStrategyConfig,
-    )
+    from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
     from nshtrainer._directory import DirectoryConfig as DirectoryConfig
     from nshtrainer._hf_hub import CallbackConfigBase as CallbackConfigBase
     from nshtrainer._hf_hub import (
@@ -122,9 +110,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -132,10 +117,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 (
@@ -201,21 +184,13 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).BestCheckpointCallbackConfig
-        if name == "BestCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).BestCheckpointStrategyConfig
         if name == "CSVLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").CSVLoggerConfig
         if name == "CallbackConfigBase":
             return importlib.import_module("nshtrainer._hf_hub").CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointMetadata":
             return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
+                "nshtrainer._checkpoint.metadata"
             ).CheckpointMetadata
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
@@ -319,12 +294,12 @@ else:
             return importlib.import_module(
                 "nshtrainer.callbacks"
             ).LastCheckpointCallbackConfig
-        if name == "LastCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).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 +308,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 +322,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 +342,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"
@@ -411,10 +376,6 @@ else:
             return importlib.import_module("nshtrainer.loggers").TensorboardLoggerConfig
         if name == "TrainerConfig":
             return importlib.import_module("nshtrainer").TrainerConfig
-        if name == "UserProvidedPathCheckpointStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).UserProvidedPathCheckpointStrategyConfig
         if name == "WandbLoggerConfig":
             return importlib.import_module("nshtrainer.loggers").WandbLoggerConfig
         if name == "WandbUploadCodeCallbackConfig":
@@ -431,10 +392,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CheckpointCallbackConfig
-        if name == "CheckpointLoadingStrategyConfig":
-            return importlib.import_module(
-                "nshtrainer._checkpoint.loader"
-            ).CheckpointLoadingStrategyConfig
         if name == "DurationConfig":
             return importlib.import_module("nshtrainer.util.config").DurationConfig
         if name == "LRSchedulerConfig":

nshtrainer-1.0.0b14/src/nshtrainer/configs/_checkpoint/__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._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
+    from nshtrainer._checkpoint.metadata import EnvironmentConfig as EnvironmentConfig
+else:
+
+    def __getattr__(name):
+        import importlib
+
+        if name in globals():
+            return globals()[name]
+        if name == "CheckpointMetadata":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).CheckpointMetadata
+        if name == "EnvironmentConfig":
+            return importlib.import_module(
+                "nshtrainer._checkpoint.metadata"
+            ).EnvironmentConfig
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+
+# Submodule exports
+from . import metadata as metadata

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b14}/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.0b14/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.0b12 → nshtrainer-1.0.0b14}/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,
     )
@@ -17,9 +18,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -39,20 +37,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 +74,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"
@@ -85,10 +88,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +118,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"
@@ -176,5 +173,4 @@ else:
 
 # Submodule exports
 from . import _config as _config
-from . import checkpoint_connector as checkpoint_connector
 from . import trainer as trainer

{nshtrainer-1.0.0b12 → nshtrainer-1.0.0b14}/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,
     )
@@ -16,9 +17,6 @@ if TYPE_CHECKING:
     from nshtrainer.trainer._config import (
         CheckpointCallbackConfig as CheckpointCallbackConfig,
     )
-    from nshtrainer.trainer._config import (
-        CheckpointLoadingConfig as CheckpointLoadingConfig,
-    )
     from nshtrainer.trainer._config import (
         CheckpointSavingConfig as CheckpointSavingConfig,
     )
@@ -38,20 +36,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 +74,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"
@@ -85,10 +88,6 @@ else:
             return importlib.import_module(
                 "nshtrainer.trainer._config"
             ).CallbackConfigBase
-        if name == "CheckpointLoadingConfig":
-            return importlib.import_module(
-                "nshtrainer.trainer._config"
-            ).CheckpointLoadingConfig
         if name == "CheckpointSavingConfig":
             return importlib.import_module(
                 "nshtrainer.trainer._config"
@@ -119,30 +118,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.0b12 → nshtrainer-1.0.0b14}/src/nshtrainer/data/datamodule.py

@@ -8,8 +8,6 @@ from typing import Any, Generic, cast
 import nshconfig as C
 import torch
 from lightning.pytorch import LightningDataModule
-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 ..model.mixins.callback import CallbackRegistrarModuleMixin