nshtrainer 1.0.0b11__py3-none-any.whl → 1.0.0b13__py3-none-any.whl

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/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/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/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/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/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/data/datamodule.py

@@ -1,10 +1,12 @@
 from __future__ import annotations
 
 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, cast
 
 import nshconfig as C
+import torch
 from lightning.pytorch import LightningDataModule
 from typing_extensions import Never, TypeVar, deprecated, override
 
@@ -55,3 +57,68 @@ class LightningDataModuleBase(
             )
         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/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

nshtrainer/trainer/_config.py

@@ -9,7 +9,6 @@ from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from pathlib import Path
 from typing import (
-    TYPE_CHECKING,
     Annotated,
     Any,
     ClassVar,
@@ -41,11 +40,13 @@ from ..callbacks import (
     CallbackConfig,
     EarlyStoppingCallbackConfig,
     LastCheckpointCallbackConfig,
+    NormLoggingCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
+from ..callbacks.lr_monitor import LearningRateMonitorConfig
 from ..callbacks.rlp_sanity_checks import RLPSanityChecksCallbackConfig
 from ..callbacks.shared_parameters import SharedParametersCallbackConfig
 from ..loggers import (
@@ -54,6 +55,7 @@ from ..loggers import (
     TensorboardLoggerConfig,
     WandbLoggerConfig,
 )
+from ..loggers._base import BaseLoggerConfig
 from ..loggers.actsave import ActSaveLoggerConfig
 from ..metrics._config import MetricConfig
 from ..profiler import ProfilerConfig
@@ -62,103 +64,6 @@ from ..util._environment_info import EnvironmentConfig
 log = logging.getLogger(__name__)
 
 
-class LoggingConfig(CallbackConfigBase):
-    enabled: bool = True
-    """Enable experiment tracking."""
-
-    loggers: Sequence[LoggerConfig] = [
-        WandbLoggerConfig(),
-        CSVLoggerConfig(),
-        TensorboardLoggerConfig(),
-    ]
-    """Loggers to use for experiment tracking."""
-
-    log_lr: bool | Literal["step", "epoch"] = True
-    """If enabled, will register a `LearningRateMonitor` callback to log the learning rate to the logger."""
-    log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
-    """If enabled, will log the fractional epoch number to the logger."""
-
-    actsave_logger: ActSaveLoggerConfig | None = None
-    """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
-
-    @property
-    def wandb(self):
-        return next(
-            (
-                logger
-                for logger in self.loggers
-                if isinstance(logger, WandbLoggerConfig)
-            ),
-            None,
-        )
-
-    @property
-    def csv(self):
-        return next(
-            (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
-            None,
-        )
-
-    @property
-    def tensorboard(self):
-        return next(
-            (
-                logger
-                for logger in self.loggers
-                if isinstance(logger, TensorboardLoggerConfig)
-            ),
-            None,
-        )
-
-    def create_loggers(self, trainer_config: TrainerConfig):
-        """
-        Constructs and returns a list of loggers based on the provided root configuration.
-
-        Args:
-            trainer_config (TrainerConfig): The root configuration object.
-
-        Returns:
-            list[Logger]: A list of constructed loggers.
-        """
-        if not self.enabled:
-            return
-
-        for logger_config in sorted(
-            self.loggers,
-            key=lambda x: x.priority,
-            reverse=True,
-        ):
-            if not logger_config.enabled:
-                continue
-            if (logger := logger_config.create_logger(trainer_config)) is None:
-                continue
-            yield logger
-
-        # If the actsave_metrics is enabled, add the ActSave logger
-        if self.actsave_logger:
-            yield self.actsave_logger.create_logger(trainer_config)
-
-    @override
-    def create_callbacks(self, trainer_config):
-        if self.log_lr:
-            from lightning.pytorch.callbacks import LearningRateMonitor
-
-            logging_interval: str | None = None
-            if isinstance(self.log_lr, str):
-                logging_interval = self.log_lr
-
-            yield LearningRateMonitor(logging_interval=logging_interval)
-
-        if self.log_epoch:
-            yield from self.log_epoch.create_callbacks(trainer_config)
-
-        for logger in self.loggers:
-            if not logger or not isinstance(logger, CallbackConfigBase):
-                continue
-
-            yield from logger.create_callbacks(trainer_config)
-
-
 class GradientClippingConfig(C.Config):
     enabled: bool = True
     """Enable gradient clipping."""
@@ -168,32 +73,6 @@ class GradientClippingConfig(C.Config):
     """Norm type to use for gradient clipping."""
 
 
-class OptimizationConfig(CallbackConfigBase):
-    log_grad_norm: bool | str | float = False
-    """If enabled, will log the gradient norm (averaged across all model parameters) to the logger."""
-    log_grad_norm_per_param: bool | str | float = False
-    """If enabled, will log the gradient norm for each model parameter to the logger."""
-
-    log_param_norm: bool | str | float = False
-    """If enabled, will log the parameter norm (averaged across all model parameters) to the logger."""
-    log_param_norm_per_param: bool | str | float = False
-    """If enabled, will log the parameter norm for each model parameter to the logger."""
-
-    gradient_clipping: GradientClippingConfig | None = None
-    """Gradient clipping configuration, or None to disable gradient clipping."""
-
-    @override
-    def create_callbacks(self, trainer_config):
-        from ..callbacks.norm_logging import NormLoggingCallbackConfig
-
-        yield from NormLoggingCallbackConfig(
-            log_grad_norm=self.log_grad_norm,
-            log_grad_norm_per_param=self.log_grad_norm_per_param,
-            log_param_norm=self.log_param_norm,
-            log_param_norm_per_param=self.log_param_norm_per_param,
-        ).create_callbacks(trainer_config)
-
-
 TPlugin = TypeVar(
     "TPlugin",
     Precision,
@@ -253,15 +132,6 @@ StrategyLiteral: TypeAlias = Literal[
 ]
 
 
-class ReproducibilityConfig(C.Config):
-    deterministic: bool | Literal["warn"] | None = None
-    """
-    If ``True``, sets whether PyTorch operations must use deterministic algorithms.
-        Set to ``"warn"`` to use deterministic algorithms whenever possible, throwing warnings on operations
-        that don't support deterministic mode. If not set, defaults to ``False``. Default: ``None``.
-    """
-
-
 CheckpointCallbackConfig: TypeAlias = Annotated[
     BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
@@ -635,14 +505,34 @@ class TrainerConfig(C.Config):
     hf_hub: HuggingFaceHubConfig = HuggingFaceHubConfig()
     """Hugging Face Hub configuration options."""
 
-    logging: LoggingConfig = LoggingConfig()
-    """Logging/experiment tracking (e.g., WandB) configuration options."""
+    loggers: Sequence[LoggerConfig] = [
+        WandbLoggerConfig(),
+        CSVLoggerConfig(),
+        TensorboardLoggerConfig(),
+    ]
+    """Loggers to use for experiment tracking."""
 
-    optimizer: OptimizationConfig = OptimizationConfig()
-    """Optimization configuration options."""
+    actsave_logger: ActSaveLoggerConfig | None = None
+    """If enabled, will automatically save logged metrics using ActSave (if nshutils is installed)."""
 
-    reproducibility: ReproducibilityConfig = ReproducibilityConfig()
-    """Reproducibility configuration options."""
+    lr_monitor: LearningRateMonitorConfig | None = LearningRateMonitorConfig()
+    """Learning rate monitoring configuration options."""
+
+    log_epoch: LogEpochCallbackConfig | None = LogEpochCallbackConfig()
+    """If enabled, will log the fractional epoch number to the logger."""
+
+    gradient_clipping: GradientClippingConfig | None = None
+    """Gradient clipping configuration, or None to disable gradient clipping."""
+
+    log_norms: NormLoggingCallbackConfig | None = None
+    """Norm logging configuration options."""
+
+    deterministic: bool | Literal["warn"] | None = None
+    """
+    If ``True``, sets whether PyTorch operations must use deterministic algorithms.
+        Set to ``"warn"`` to use deterministic algorithms whenever possible, throwing warnings on operations
+        that don't support deterministic mode. If not set, defaults to ``False``. Default: ``None``.
+    """
 
     reduce_lr_on_plateau_sanity_checking: RLPSanityChecksCallbackConfig | None = (
         RLPSanityChecksCallbackConfig()
@@ -857,27 +747,87 @@ class TrainerConfig(C.Config):
     set_float32_matmul_precision: Literal["medium", "high", "highest"] | None = None
     """If enabled, will set the torch float32 matmul precision to the specified value. Useful for faster training on Ampere+ GPUs."""
 
+    @property
+    def wandb_logger(self):
+        return next(
+            (
+                logger
+                for logger in self.loggers
+                if isinstance(logger, WandbLoggerConfig)
+            ),
+            None,
+        )
+
+    @property
+    def csv_logger(self):
+        return next(
+            (logger for logger in self.loggers if isinstance(logger, CSVLoggerConfig)),
+            None,
+        )
+
+    @property
+    def tensorboard_logger(self):
+        return next(
+            (
+                logger
+                for logger in self.loggers
+                if isinstance(logger, TensorboardLoggerConfig)
+            ),
+            None,
+        )
+
     def _nshtrainer_all_callback_configs(self) -> Iterable[CallbackConfigBase | None]:
         yield self.early_stopping
         yield self.checkpoint_saving
-        yield self.logging
-        yield self.optimizer
+        yield self.lr_monitor
+        yield from (
+            logger_config
+            for logger_config in self.loggers
+            if logger_config is not None
+            and isinstance(logger_config, CallbackConfigBase)
+        )
+        yield self.log_epoch
+        yield self.log_norms
         yield self.hf_hub
         yield self.shared_parameters
         yield self.reduce_lr_on_plateau_sanity_checking
         yield self.auto_set_debug_flag
         yield from self.callbacks
 
+    def _nshtrainer_all_logger_configs(self) -> Iterable[BaseLoggerConfig | None]:
+        yield from self.loggers
+        yield self.actsave_logger
+
     # region Helper Methods
+    def fast_dev_run_(self, value: int | bool = True, /):
+        """
+        Enables fast_dev_run mode for the trainer.
+        This will run the training loop for a specified number of batches,
+        if an integer is provided, or for a single batch if True is provided.
+        """
+        self.fast_dev_run = value
+        return self
+
     def with_fast_dev_run(self, value: int | bool = True, /):
         """
         Enables fast_dev_run mode for the trainer.
         This will run the training loop for a specified number of batches,
         if an integer is provided, or for a single batch if True is provided.
         """
-        config = copy.deepcopy(self)
-        config.fast_dev_run = value
-        return config
+        return copy.deepcopy(self).fast_dev_run_(value)
+
+    def project_root_(self, project_root: str | Path | os.PathLike):
+        """
+        Set the project root directory for the trainer.
+
+        Args:
+            project_root (Path): The base directory to use.
+
+        Returns:
+            self: The current instance of the class.
+        """
+        self.directory.project_root = Path(project_root)
+        return self
 
     def with_project_root(self, project_root: str | Path | os.PathLike):
         """
@@ -889,9 +839,7 @@ class TrainerConfig(C.Config):
         Returns:
             self: The current instance of the class.
         """
-        config = copy.deepcopy(self)
-        config.directory.project_root = Path(project_root)
-        return config
+        return copy.deepcopy(self).project_root_(project_root)
 
     def reset_run(
         self,

nshtrainer/trainer/trainer.py

@@ -2,28 +2,19 @@ from __future__ import annotations
 
 import logging
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
-from lightning.fabric.utilities.cloud_io import _load as pl_load
-from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 from lightning.pytorch import LightningModule
 from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.core.saving import (
-    _default_map_location,
-    load_hparams_from_tags_csv,
-    load_hparams_from_yaml,
-)
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.trainer.states import TrainerFn
-from lightning.pytorch.utilities.migration import pl_legacy_patch
-from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
@@ -79,7 +70,7 @@ class Trainer(LightningTrainer):
         kwargs_ctor: LightningTrainerKwargs,
     ):
         kwargs: LightningTrainerKwargs = {
-            "deterministic": hparams.reproducibility.deterministic,
+            "deterministic": hparams.deterministic,
             "fast_dev_run": hparams.fast_dev_run,
             "max_epochs": hparams.max_epochs,
             "min_epochs": hparams.min_epochs,
@@ -218,7 +209,7 @@ class Trainer(LightningTrainer):
             _update_kwargs(detect_anomaly=detect_anomaly)
 
         if (
-            grad_clip_config := hparams.optimizer.gradient_clipping
+            grad_clip_config := hparams.gradient_clipping
         ) is not None and grad_clip_config.enabled:
             # kwargs["gradient_clip_algorithm"] = grad_clip_config.algorithm
             # kwargs["gradient_clip_val"] = grad_clip_config.value
@@ -248,17 +239,14 @@ class Trainer(LightningTrainer):
                 ]
             )
 
-        if not hparams.logging.enabled:
-            log.critical(f"Disabling logger because {hparams.logging.enabled=}.")
-            kwargs["logger"] = False
-        else:
-            _update_kwargs(
-                logger=[
-                    logger
-                    for logger in hparams.logging.create_loggers(hparams)
-                    if logger is not None
-                ]
-            )
+        _update_kwargs(
+            logger=[
+                logger
+                for logger_config in hparams._nshtrainer_all_logger_configs()
+                if logger_config is not None
+                and (logger := logger_config.create_logger(hparams)) is not None
+            ]
+        )
 
         if hparams.auto_determine_num_nodes:
             # When num_nodes is auto, we need to detect the number of nodes.
@@ -473,62 +461,46 @@ class Trainer(LightningTrainer):
         _callback._call_on_checkpoint_saved(self, filepath, metadata_path)
 
     @classmethod
-    def load_from_checkpoint(
+    def hparams_from_checkpoint(
         cls,
-        checkpoint_path: _PATH | IO,
-        map_location: _MAP_LOCATION_TYPE = None,
-        hparams_file: _PATH | None = None,
-        **kwargs: Any,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[TrainerConfig], TrainerConfig] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
     ):
-        loaded = _load_from_checkpoint(
-            checkpoint_path,
-            map_location=map_location,
-            hparams_file=hparams_file,
-            **kwargs,
-        )
-        return loaded
-
-
-def _load_from_checkpoint(
-    checkpoint_path: _PATH | IO,
-    map_location: _MAP_LOCATION_TYPE = None,
-    hparams_file: _PATH | None = None,
-    **kwargs: Any,
-):
-    map_location = map_location or _default_map_location
-    with pl_legacy_patch():
-        checkpoint = pl_load(checkpoint_path, map_location=map_location)
-
-    # convert legacy checkpoints to the new format
-    checkpoint = _pl_migrate_checkpoint(
-        checkpoint,
-        checkpoint_path=(
-            checkpoint_path if isinstance(checkpoint_path, (str, Path)) else None
-        ),
-    )
-
-    if hparams_file is not None:
-        extension = str(hparams_file).split(".")[-1]
-        if extension.lower() == "csv":
-            hparams = load_hparams_from_tags_csv(hparams_file)
-        elif extension.lower() in ("yml", "yaml"):
-            hparams = load_hparams_from_yaml(hparams_file)
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
         else:
-            raise ValueError(".csv, .yml or .yaml is required for `hparams_file`")
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
 
-        # overwrite hparams by the given file
-        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
+        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)
 
-    # for past checkpoint need to add the new key
-    checkpoint.setdefault(Trainer.CHECKPOINT_HYPER_PARAMS_KEY, {})
-    # override the hparams with values that were passed in
-    checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
 
-    # load the hparams
-    hparams = Trainer.hparams_cls().model_validate(
-        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY]
-    )
+        return hparams
 
-    # create the trainer
-    trainer = Trainer(hparams)
-    return trainer
+    @classmethod
+    def from_checkpoint(
+        cls,
+        path: str | Path,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[TrainerConfig], TrainerConfig] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        hparams = cls.hparams_from_checkpoint(
+            path,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+        return cls(hparams)

{nshtrainer-1.0.0b11.dist-info → nshtrainer-1.0.0b13.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-1.0.0b13.dist-info}/RECORD

@@ -22,6 +22,7 @@ nshtrainer/callbacks/finite_checks.py,sha256=iCiKQ5i9RckkzcPeCHzC3hkg3AlW3ESuWtF
 nshtrainer/callbacks/gradient_skipping.py,sha256=k5qNaNeileZ_5YFad4ssfLplMxMKeKFhPcY8-QVmLek,3464
 nshtrainer/callbacks/interval.py,sha256=UCzUzt3XCFVyQyCWL9lOrStkkxesvduNOYk8yMrGTTk,8116
 nshtrainer/callbacks/log_epoch.py,sha256=Wr-Ksxsynsqu_zyB_zoiPLjnWv-ksC3xPekY6iyN-P8,1396
+nshtrainer/callbacks/lr_monitor.py,sha256=IyFZoXaxJoTBSkdLu1iEZ1qI8_UFNJwafR_xTVPZXXU,1050
 nshtrainer/callbacks/norm_logging.py,sha256=C44Mvt73gqQEpCFd0j3qYg6NY7sL2jm3X1qJVY_XLfI,6329
 nshtrainer/callbacks/print_table.py,sha256=WIgfzVSfAfS3_8kUuX-nWJOGWBEmtNlejypuoJQViPY,2884
 nshtrainer/callbacks/rlp_sanity_checks.py,sha256=kWl2dYOXn2L8k6ub_012jNkqOxtyea1yr1qWRNG6UW4,9990
@@ -29,13 +30,13 @@ nshtrainer/callbacks/shared_parameters.py,sha256=33eRzifNj6reKbvmGuam1hUofo3sD4J
 nshtrainer/callbacks/timer.py,sha256=BB-M7tV4QNYOwY_Su6j9P7IILxVRae_upmDq4qsxiao,4670
 nshtrainer/callbacks/wandb_upload_code.py,sha256=PTqNE1QB5U8NR5zhbiQZrmQuugX2UV7B12UdMpo9aV0,2353
 nshtrainer/callbacks/wandb_watch.py,sha256=tTTcFzxd2Ia9xu8tCogQ5CLJZBq1ne5JlpGVE75vKYs,2976
-nshtrainer/configs/__init__.py,sha256=0TczWa5OFRKOGKHgabeB7VUMxPpD0RCgDR6AvdAD-tI,22721
+nshtrainer/configs/__init__.py,sha256=Vyf_gn7u3s9ET4Yszf6SILtqvpIGiJ4X5RJfmW-FK6I,22293
 nshtrainer/configs/_checkpoint/__init__.py,sha256=vuiBbd4VzCo7lRyhyTUArEQeWwJkewvNPKDxBJiUHoY,2719
 nshtrainer/configs/_checkpoint/loader/__init__.py,sha256=hdLpypoEkES1MTaTHAdGFJnSoZzgx_8NzAKbK143SyI,2399
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=X9KxpcoHQbJp6-MTGvp4pct-MYHaHcl82s9yqZ5KiSk,867
 nshtrainer/configs/_directory/__init__.py,sha256=mTUoSz-DSsvI2M98cqu2Z2x215oM0sLyljh_5rVexvQ,1029
 nshtrainer/configs/_hf_hub/__init__.py,sha256=3HGCGhRb7NhOuLeskGqbYNuS9c81oOUbX6ibyF3XiCY,1063
-nshtrainer/configs/callbacks/__init__.py,sha256=-dHN8NZdCaNUy_isnlh779FZ1w9_WkOkv6VSN_-86jM,7316
+nshtrainer/configs/callbacks/__init__.py,sha256=xgCa98EmqU8cHxlJa-64Cc4c_0fS0Cz2iVac4edL_yc,7657
 nshtrainer/configs/callbacks/actsave/__init__.py,sha256=AkVWS9vCcDJFpPUpyc7i9cjaFZU2kKxDyFDqakMZA-E,809
 nshtrainer/configs/callbacks/base/__init__.py,sha256=OdtHDMkYC_ioCEAkg7bSQi3o7e2t5WHPcFjavXdfdTA,602
 nshtrainer/configs/callbacks/checkpoint/__init__.py,sha256=DE-JTtY4NdaP3mgWduearFYMvy1tswWRBWMde06RzQc,2700
@@ -50,6 +51,7 @@ nshtrainer/configs/callbacks/ema/__init__.py,sha256=KlPGdJWjYTKLdpl-VnN4BYY2sA_L
 nshtrainer/configs/callbacks/finite_checks/__init__.py,sha256=V6Owp05XdIk3EO67AMVGdwbT4-D86QRuvqWM2gu5Xpw,949
 nshtrainer/configs/callbacks/gradient_skipping/__init__.py,sha256=RhIfJFq-x_sWYrWVGaVEBeT8uUFYjFgt0Ug8pPgpJSg,981
 nshtrainer/configs/callbacks/log_epoch/__init__.py,sha256=4jePzjE3bVxaI7hQrcWW5SrKT5MrFyplJZwK8bQHbGI,900
+nshtrainer/configs/callbacks/lr_monitor/__init__.py,sha256=iC8U0oWC75JzPUMRoGWkC8WkMuLbF9-zuN_yQlByycY,916
 nshtrainer/configs/callbacks/norm_logging/__init__.py,sha256=8M9hcGpEuQadfgTR4-YL4TWeyxZjg0s84x420B03-aE,941
 nshtrainer/configs/callbacks/print_table/__init__.py,sha256=Ni47iS2mIzwGu8XuHfUY5BJKawUO_2TyJMZ62QBpEW0,961
 nshtrainer/configs/callbacks/rlp_sanity_checks/__init__.py,sha256=OgPywk8Z9y_dnq_liH2PPWuQSpUlQ_Q2-q99HDN9Leg,977
@@ -78,8 +80,8 @@ nshtrainer/configs/profiler/_base/__init__.py,sha256=smDYCplrI5B38XJcNZ462ZeTo9l
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=VCCvbzhEeOcdZ0Unvk_anAcmbQuGogTQhK_bXs5RG9U,892
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=hJ90ym5ElI-BY_XS3VSLjcgQWfV0Pp1MdzTU6Qi8MFg,884
 nshtrainer/configs/profiler/simple/__init__.py,sha256=18V64kKYrJeSCrPmY3wYnshEISaf7xmrfw2Ny-6P3uE,859
-nshtrainer/configs/trainer/__init__.py,sha256=Gf5RizrVL84NWfHnagVCRHtXiD6x0UA4N7vhYypluTk,7916
-nshtrainer/configs/trainer/_config/__init__.py,sha256=BeQ5t_9d6rx6SbSu4ZqD9eitLCQRpTOVOnSxT0LCrlM,7806
+nshtrainer/configs/trainer/__init__.py,sha256=QLCDVxVg1Ig-wgUW5r8I1FdPdbYz9-gse17s3R69Fw0,8019
+nshtrainer/configs/trainer/_config/__init__.py,sha256=YEcliai8jLOoB53lxT5BZIR3NzfoLb4x3VGbMakFVo4,7909
 nshtrainer/configs/trainer/checkpoint_connector/__init__.py,sha256=pSu79zOFFWvqjI3SkHWl13H8ZNJFTc6a5J1r2KnfUKM,667
 nshtrainer/configs/trainer/trainer/__init__.py,sha256=P-Y2DOZZcJtEdPjGEKCxq5R3JSzKhUUoidkSvO_cfKI,797
 nshtrainer/configs/util/__init__.py,sha256=ZcmEqg2OWKKcPBqzDG1SnuaAMgR4_c0jog-Xg6QTUzc,4555
@@ -89,7 +91,7 @@ nshtrainer/configs/util/config/dtype/__init__.py,sha256=NCXMVO-EUz3JvPmlDci72O9Z
 nshtrainer/configs/util/config/duration/__init__.py,sha256=8llT1MCKQpsdNldN5h5Wo0GjUuRn28Sxw2FTXTNKBpM,1060
 nshtrainer/data/__init__.py,sha256=K4i3Tw4g9EOK2zlMMbidi99y0SyI4w8P7_XUf1n42Ts,260
 nshtrainer/data/balanced_batch_sampler.py,sha256=r1cBKRXKHD8E1Ax6tj-FUbE-z1qpbO58mQ9VrK9uLnc,5481
-nshtrainer/data/datamodule.py,sha256=NC0y7JAOvYCei-yPGUOeIB7MkHVcWUQE-5dWC0Hvpxo,1824
+nshtrainer/data/datamodule.py,sha256=lSOgH32nysJWa6Y7ba1QyOdUV0DVVdO98qokP8wigjk,4138
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
 nshtrainer/loggers/__init__.py,sha256=11X6D_lF0cHAkxkYsVZY3Q3r30Fq0FUi9heeb5RD870,570
 nshtrainer/loggers/_base.py,sha256=nw4AZzJP3Z-fljgQlgq7FkuMkPmYKTsXj7OfJJSmtXI,811
@@ -104,7 +106,7 @@ nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=vXH5S26ESHO_LPPqW8aDC3S5N
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=kR9OJLZXIaZpG8UpHG3EQwOL36HB8yPk6afYjoIM0XM,1324
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
-nshtrainer/model/base.py,sha256=qJXzu1w4mqF1eVwRBahEuJwJTIVSAZGZaV1vigwrg0Y,6817
+nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,10356
 nshtrainer/model/mixins/callback.py,sha256=Ea_legORzs0N078j0N9RJivDVeWH5KtXDpdJS75IwIo,3098
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=LQDJJbiv30PlWX6rTT_EhjNBNfUFfcvGz5sX4MnOCzI,5330
@@ -119,13 +121,12 @@ nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
-nshtrainer/scripts/find_packages.py,sha256=-oNnSNPp3pujCVgManW_LFlJcnvhrtHgvUJ4W88-6-o,1460
 nshtrainer/trainer/__init__.py,sha256=MmoydVS6aYeav7zgDAUHxAQrV_PMQsbnZTCuPnLH9Wk,128
-nshtrainer/trainer/_config.py,sha256=5HC2Gs7YbJK9gMcKm4cQ78B3PxcOvwQ372Glhf4fzbo,34232
+nshtrainer/trainer/_config.py,sha256=2S6Qhwn724n_jgGhWVI64Wi_pHKjU1ggoY4sxq-_SlA,32309
 nshtrainer/trainer/_runtime_callback.py,sha256=T3epaj1YeIN0R8CS2cg5HNJIB21TyaD_PVNNOPJ6nJs,4200
 nshtrainer/trainer/checkpoint_connector.py,sha256=pC1tTDcq0p6sAsoTmAbwINW49IfqupMMtnE9-AKdTUw,2824
 nshtrainer/trainer/signal_connector.py,sha256=YMJf6vTnW0JcnBkuYikm9x_9XscaokrCEzCn4THOGao,10776
-nshtrainer/trainer/trainer.py,sha256=JLcfdPqlOaAmKJAE3w9xIEai7Xyyl3JhDyMPsZ8fyv8,20692
+nshtrainer/trainer/trainer.py,sha256=kIXh_25jDJSGcwEyLjbvqWN0P5B35VBJLXOwXqUGqF4,19759
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/_useful_types.py,sha256=dwZokFkIe7M5i2GR3nQ9A1lhGw06DMAFfH5atyquqSA,8000
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
@@ -138,6 +139,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b11.dist-info/METADATA,sha256=6IsguMlupAvNHaGUyo3m1-TzLvwycrpucABLwbzyj_Q,937
-nshtrainer-1.0.0b11.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-1.0.0b11.dist-info/RECORD,,
+nshtrainer-1.0.0b13.dist-info/METADATA,sha256=9PQNipTw68KmSV_7Kt4fK_KtlYKSaKBcvvkBZrwWFtY,937
+nshtrainer-1.0.0b13.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-1.0.0b13.dist-info/RECORD,,

nshtrainer/scripts/find_packages.py

@@ -1,52 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import ast
-import glob
-import sys
-from pathlib import Path
-
-
-def get_imports(file_path: Path):
-    with open(file_path, "r") as file:
-        try:
-            tree = ast.parse(file.read())
-        except SyntaxError:
-            print(f"Syntax error in file: {file_path}", file=sys.stderr)
-            return set()
-
-    imports = set()
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Import):
-            for alias in node.names:
-                imports.add(alias.name.split(".")[0])
-        elif isinstance(node, ast.ImportFrom):
-            if node.level == 0 and node.module:  # Absolute import
-                imports.add(node.module.split(".")[0])
-    return imports
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Find unique Python packages used in files."
-    )
-    parser.add_argument("glob_pattern", help="Glob pattern to match files")
-    parser.add_argument(
-        "--exclude-std", action="store_true", help="Exclude Python standard libraries"
-    )
-    args = parser.parse_args()
-
-    all_imports = set()
-    for file_path in glob.glob(args.glob_pattern, recursive=True):
-        all_imports.update(get_imports(Path(file_path)))
-
-    if args.exclude_std:
-        std_libs = set(sys.stdlib_module_names)
-        all_imports = all_imports - std_libs
-
-    for package in sorted(all_imports):
-        print(package)
-
-
-if __name__ == "__main__":
-    main()