nshtrainer 0.32.0__tar.gz → 0.33.0__tar.gz

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/PKG-INFO

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

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/pyproject.toml

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

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/src/nshtrainer/callbacks/debug_flag.py

@@ -5,8 +5,6 @@ from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks import Callback
 from typing_extensions import override
 
-from nshtrainer.model.config import BaseConfig
-
 from .base import CallbackConfigBase
 
 if TYPE_CHECKING:

nshtrainer-0.33.0/src/nshtrainer/config.py

@@ -0,0 +1,135 @@
+from nshconfig._config import Config as Config
+from nshsnap._config import SnapshotConfig as SnapshotConfig
+
+from nshtrainer._checkpoint.loader import (
+    CheckpointLoadingConfig as CheckpointLoadingConfig,
+)
+from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
+from nshtrainer._directory import DirectoryConfig as DirectoryConfig
+from nshtrainer._hf_hub import (
+    HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig,
+)
+from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
+from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
+from nshtrainer.callbacks.base import CallbackConfigBase as CallbackConfigBase
+from nshtrainer.callbacks.checkpoint._base import (
+    BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.checkpoint.best_checkpoint import (
+    BestCheckpointCallbackConfig as BestCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.checkpoint.last_checkpoint import (
+    LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import (
+    OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
+)
+from nshtrainer.callbacks.debug_flag import (
+    DebugFlagCallbackConfig as DebugFlagCallbackConfig,
+)
+from nshtrainer.callbacks.directory_setup import (
+    DirectorySetupConfig as DirectorySetupConfig,
+)
+from nshtrainer.callbacks.early_stopping import (
+    EarlyStoppingConfig as EarlyStoppingConfig,
+)
+from nshtrainer.callbacks.ema import EMAConfig as EMAConfig
+from nshtrainer.callbacks.finite_checks import FiniteChecksConfig as FiniteChecksConfig
+from nshtrainer.callbacks.gradient_skipping import (
+    GradientSkippingConfig as GradientSkippingConfig,
+)
+from nshtrainer.callbacks.norm_logging import NormLoggingConfig as NormLoggingConfig
+from nshtrainer.callbacks.print_table import (
+    PrintTableMetricsConfig as PrintTableMetricsConfig,
+)
+from nshtrainer.callbacks.rlp_sanity_checks import (
+    RLPSanityChecksConfig as RLPSanityChecksConfig,
+)
+from nshtrainer.callbacks.shared_parameters import (
+    SharedParametersConfig as SharedParametersConfig,
+)
+from nshtrainer.callbacks.throughput_monitor import (
+    ThroughputMonitorConfig as ThroughputMonitorConfig,
+)
+from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
+from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
+from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
+from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
+from nshtrainer.loggers.tensorboard import (
+    TensorboardLoggerConfig as TensorboardLoggerConfig,
+)
+from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
+from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from nshtrainer.lr_scheduler.linear_warmup_cosine import (
+    DurationConfig as DurationConfig,
+)
+from nshtrainer.lr_scheduler.linear_warmup_cosine import (
+    LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig,
+)
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
+    ReduceLROnPlateauConfig as ReduceLROnPlateauConfig,
+)
+from nshtrainer.metrics._config import MetricConfig as MetricConfig
+from nshtrainer.model.config import BaseConfig as BaseConfig
+from nshtrainer.nn.mlp import MLPConfig as MLPConfig
+from nshtrainer.nn.nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
+from nshtrainer.nn.nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import (
+    LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
+from nshtrainer.nn.nonlinearity import NonlinearityConfig as NonlinearityConfig
+from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
+from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import (
+    SigmoidNonlinearityConfig as SigmoidNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
+from nshtrainer.nn.nonlinearity import (
+    SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import (
+    SoftplusNonlinearityConfig as SoftplusNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import (
+    SoftsignNonlinearityConfig as SoftsignNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import (
+    SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import (
+    SwishNonlinearityConfig as SwishNonlinearityConfig,
+)
+from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from nshtrainer.optimizer import AdamWConfig as AdamWConfig
+from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
+from nshtrainer.profiler._base import BaseProfilerConfig as BaseProfilerConfig
+from nshtrainer.profiler.advanced import (
+    AdvancedProfilerConfig as AdvancedProfilerConfig,
+)
+from nshtrainer.profiler.pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
+from nshtrainer.profiler.simple import SimpleProfilerConfig as SimpleProfilerConfig
+from nshtrainer.trainer._config import CallbackConfig as CallbackConfig
+from nshtrainer.trainer._config import (
+    CheckpointCallbackConfig as CheckpointCallbackConfig,
+)
+from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
+from nshtrainer.trainer._config import GradientClippingConfig as GradientClippingConfig
+from nshtrainer.trainer._config import LoggerConfig as LoggerConfig
+from nshtrainer.trainer._config import LoggingConfig as LoggingConfig
+from nshtrainer.trainer._config import OptimizationConfig as OptimizationConfig
+from nshtrainer.trainer._config import ProfilerConfig as ProfilerConfig
+from nshtrainer.trainer._config import ReproducibilityConfig as ReproducibilityConfig
+from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
+from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
+from nshtrainer.util._environment_info import (
+    EnvironmentClassInformationConfig as EnvironmentClassInformationConfig,
+)
+from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
+from nshtrainer.util._environment_info import (
+    EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig,
+)
+from nshtrainer.util._environment_info import (
+    EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig,
+)

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/src/nshtrainer/loggers/wandb.py

@@ -7,8 +7,8 @@ from lightning.pytorch import Callback, LightningModule, Trainer
 from packaging import version
 from typing_extensions import override
 
-from ..callbacks import WandbWatchConfig
 from ..callbacks.base import CallbackConfigBase
+from ..callbacks.wandb_watch import WandbWatchConfig
 from ._base import BaseLoggerConfig
 
 if TYPE_CHECKING:

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/src/nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -6,7 +6,7 @@ from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from typing_extensions import override
 
-from ..util.config import Duration
+from ..util.config import DurationConfig
 from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
 
@@ -91,11 +91,11 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
 class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
     name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
 
-    warmup_duration: Duration
+    warmup_duration: DurationConfig
     r"""The duration for the linear warmup phase.
     The learning rate is linearly increased from `warmup_start_lr` to the initial learning rate over this duration."""
 
-    max_duration: Duration
+    max_duration: DurationConfig
     r"""The total duration.
     The learning rate is decayed to `min_lr` over this duration."""
 

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/src/nshtrainer/profiler/__init__.py

@@ -9,5 +9,5 @@ from .simple import SimpleProfilerConfig as SimpleProfilerConfig
 
 ProfilerConfig: TypeAlias = Annotated[
     SimpleProfilerConfig | AdvancedProfilerConfig | PyTorchProfilerConfig,
-    C.Discriminator("name"),
+    C.Field(discriminator="name"),
 ]

nshtrainer-0.33.0/src/nshtrainer/util/config/__init__.py

@@ -0,0 +1,4 @@
+from . import duration as duration
+from .duration import DurationConfig as DurationConfig
+from .duration import EpochsConfig as EpochsConfig
+from .duration import StepsConfig as StepsConfig

{nshtrainer-0.32.0 → nshtrainer-0.33.0}/src/nshtrainer/util/config/duration.py

@@ -4,7 +4,7 @@ from typing import Annotated, Literal
 import nshconfig as C
 
 
-class Steps(C.Config):
+class StepsConfig(C.Config):
     kind: Literal["steps"] = "steps"
 
     value: Annotated[int, C.Field(ge=0)]
@@ -14,7 +14,7 @@ class Steps(C.Config):
         return self
 
 
-class Epochs(C.Config):
+class EpochsConfig(C.Config):
     kind: Literal["epochs"] = "epochs"
 
     value: Annotated[int | float, C.Field(ge=0)]
@@ -25,7 +25,7 @@ class Epochs(C.Config):
         if not isinstance(value, int):
             value = int(math.ceil(value))
 
-        return Steps(value=value)
+        return StepsConfig(value=value)
 
 
-Duration = Annotated[Steps | Epochs, C.Field(discriminator="kind")]
+DurationConfig = Annotated[StepsConfig | EpochsConfig, C.Field(discriminator="kind")]

nshtrainer-0.32.0/src/nshtrainer/config.py

@@ -1,67 +0,0 @@
-from nshconfig._config import Config as Config
-from nshsnap._config import SnapshotConfig as SnapshotConfig
-from nshtrainer._checkpoint.loader import CheckpointLoadingConfig as CheckpointLoadingConfig
-from nshtrainer._checkpoint.metadata import CheckpointMetadata as CheckpointMetadata
-from nshtrainer._directory import DirectoryConfig as DirectoryConfig
-from nshtrainer._hf_hub import HuggingFaceHubAutoCreateConfig as HuggingFaceHubAutoCreateConfig
-from nshtrainer._hf_hub import HuggingFaceHubConfig as HuggingFaceHubConfig
-from nshtrainer.callbacks.actsave import ActSaveConfig as ActSaveConfig
-from nshtrainer.callbacks.base import CallbackConfigBase as CallbackConfigBase
-from nshtrainer.callbacks.checkpoint._base import BaseCheckpointCallbackConfig as BaseCheckpointCallbackConfig
-from nshtrainer.callbacks.checkpoint.best_checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
-from nshtrainer.callbacks.checkpoint.last_checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
-from nshtrainer.callbacks.checkpoint.on_exception_checkpoint import OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig
-from nshtrainer.callbacks.directory_setup import DirectorySetupConfig as DirectorySetupConfig
-from nshtrainer.callbacks.early_stopping import EarlyStoppingConfig as EarlyStoppingConfig
-from nshtrainer.callbacks.ema import EMAConfig as EMAConfig
-from nshtrainer.callbacks.finite_checks import FiniteChecksConfig as FiniteChecksConfig
-from nshtrainer.callbacks.gradient_skipping import GradientSkippingConfig as GradientSkippingConfig
-from nshtrainer.callbacks.norm_logging import NormLoggingConfig as NormLoggingConfig
-from nshtrainer.callbacks.print_table import PrintTableMetricsConfig as PrintTableMetricsConfig
-from nshtrainer.callbacks.rlp_sanity_checks import RLPSanityChecksConfig as RLPSanityChecksConfig
-from nshtrainer.callbacks.shared_parameters import SharedParametersConfig as SharedParametersConfig
-from nshtrainer.callbacks.throughput_monitor import ThroughputMonitorConfig as ThroughputMonitorConfig
-from nshtrainer.callbacks.timer import EpochTimerConfig as EpochTimerConfig
-from nshtrainer.callbacks.wandb_watch import WandbWatchConfig as WandbWatchConfig
-from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
-from nshtrainer.loggers.csv import CSVLoggerConfig as CSVLoggerConfig
-from nshtrainer.loggers.tensorboard import TensorboardLoggerConfig as TensorboardLoggerConfig
-from nshtrainer.loggers.wandb import WandbLoggerConfig as WandbLoggerConfig
-from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
-from nshtrainer.lr_scheduler.linear_warmup_cosine import LinearWarmupCosineDecayLRSchedulerConfig as LinearWarmupCosineDecayLRSchedulerConfig
-from nshtrainer.lr_scheduler.reduce_lr_on_plateau import ReduceLROnPlateauConfig as ReduceLROnPlateauConfig
-from nshtrainer.metrics._config import MetricConfig as MetricConfig
-from nshtrainer.model.config import BaseConfig as BaseConfig
-from nshtrainer.nn.mlp import MLPConfig as MLPConfig
-from nshtrainer.nn.nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
-from nshtrainer.nn.nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
-from nshtrainer.nn.nonlinearity import PReLUConfig as PReLUConfig
-from nshtrainer.nn.nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SiLUNonlinearityConfig as SiLUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SwiGLUNonlinearityConfig as SwiGLUNonlinearityConfig
-from nshtrainer.nn.nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
-from nshtrainer.nn.nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
-from nshtrainer.optimizer import AdamWConfig as AdamWConfig
-from nshtrainer.optimizer import OptimizerConfigBase as OptimizerConfigBase
-from nshtrainer.profiler._base import BaseProfilerConfig as BaseProfilerConfig
-from nshtrainer.profiler.advanced import AdvancedProfilerConfig as AdvancedProfilerConfig
-from nshtrainer.profiler.pytorch import PyTorchProfilerConfig as PyTorchProfilerConfig
-from nshtrainer.profiler.simple import SimpleProfilerConfig as SimpleProfilerConfig
-from nshtrainer.trainer._config import CheckpointSavingConfig as CheckpointSavingConfig
-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
-from nshtrainer.trainer._config import SanityCheckingConfig as SanityCheckingConfig
-from nshtrainer.trainer._config import TrainerConfig as TrainerConfig
-from nshtrainer.util._environment_info import EnvironmentClassInformationConfig as EnvironmentClassInformationConfig
-from nshtrainer.util._environment_info import EnvironmentConfig as EnvironmentConfig
-from nshtrainer.util._environment_info import EnvironmentLinuxEnvironmentConfig as EnvironmentLinuxEnvironmentConfig
-from nshtrainer.util._environment_info import EnvironmentSLURMInformationConfig as EnvironmentSLURMInformationConfig

nshtrainer-0.32.0/src/nshtrainer/util/config/__init__.py

@@ -1,4 +0,0 @@
-from . import duration as duration
-from .duration import Duration as Duration
-from .duration import Epochs as Epochs
-from .duration import Steps as Steps