nshtrainer 1.0.0b33__py3-none-any.whl → 1.0.0b37__py3-none-any.whl

nshtrainer/lr_scheduler/{_base.py → base.py}

@@ -94,3 +94,6 @@ class LRSchedulerConfigBase(C.Config, ABC):
         # ^ This is a hack to trigger the computation of the estimated stepping batches
         #   and make sure that the `trainer.num_training_batches` attribute is set.
         return math.ceil(trainer.num_training_batches / trainer.accumulate_grad_batches)
+
+
+lr_scheduler_registry = C.Registry(LRSchedulerConfigBase, discriminator="name")

nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -6,10 +6,64 @@ from typing import Literal
 
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
-from typing_extensions import override
+from typing_extensions import final, override
 
 from ..util.config import DurationConfig
-from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
+from .base import LRSchedulerConfigBase, LRSchedulerMetadata, lr_scheduler_registry
+
+
+@final
+@lr_scheduler_registry.register
+class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
+    name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
+
+    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: DurationConfig
+    r"""The total duration.
+    The learning rate is decayed to `min_lr` over this duration."""
+
+    warmup_start_lr_factor: float = 0.0
+    r"""The initial learning rate for the linear warmup phase, as a factor of the initial learning rate.
+    The learning rate is linearly increased from this value to the initial learning rate over `warmup_epochs` epochs."""
+
+    min_lr_factor: float = 0.0
+    r"""The minimum learning rate, as a factor of the initial learning rate.
+    The learning rate is decayed to this value over `max_epochs` epochs."""
+
+    annealing: bool = False
+    r"""Whether to restart the learning rate schedule after `max_epochs` epochs.
+    If `False`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be set to `min_lr` for all subsequent epochs.
+    If `True`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be increased back to the initial learning rate over `max_epochs` epochs, and so on (this is called a cosine annealing schedule)."""
+
+    @override
+    def metadata(self) -> LRSchedulerMetadata:
+        return {
+            "interval": "step",
+        }
+
+    @override
+    def create_scheduler_impl(self, optimizer, lightning_module):
+        num_steps_per_epoch = self.compute_num_steps_per_epoch(lightning_module)
+        warmup_steps = self.warmup_duration.to_steps(num_steps_per_epoch).value
+        max_steps = self.max_duration.to_steps(num_steps_per_epoch).value
+
+        # Warmup and max steps should be at least 1.
+        warmup_steps = max(warmup_steps, 1)
+        max_steps = max(max_steps, 1)
+
+        # Create the scheduler
+        scheduler = LinearWarmupCosineAnnealingLR(
+            optimizer=optimizer,
+            warmup_epochs=warmup_steps,
+            max_epochs=max_steps,
+            warmup_start_lr_factor=self.warmup_start_lr_factor,
+            eta_min_factor=self.min_lr_factor,
+            should_restart=self.annealing,
+        )
+        return scheduler
 
 
 class LinearWarmupCosineAnnealingLR(LRScheduler):
@@ -89,55 +143,3 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
             + self.eta_min_factor * base_lr
             for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
         ]
-
-
-class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
-    name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
-
-    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: DurationConfig
-    r"""The total duration.
-    The learning rate is decayed to `min_lr` over this duration."""
-
-    warmup_start_lr_factor: float = 0.0
-    r"""The initial learning rate for the linear warmup phase, as a factor of the initial learning rate.
-    The learning rate is linearly increased from this value to the initial learning rate over `warmup_epochs` epochs."""
-
-    min_lr_factor: float = 0.0
-    r"""The minimum learning rate, as a factor of the initial learning rate.
-    The learning rate is decayed to this value over `max_epochs` epochs."""
-
-    annealing: bool = False
-    r"""Whether to restart the learning rate schedule after `max_epochs` epochs.
-    If `False`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be set to `min_lr` for all subsequent epochs.
-    If `True`, the learning rate will be decayed to `min_lr` over `max_epochs` epochs, and then the learning rate will be increased back to the initial learning rate over `max_epochs` epochs, and so on (this is called a cosine annealing schedule)."""
-
-    @override
-    def metadata(self) -> LRSchedulerMetadata:
-        return {
-            "interval": "step",
-        }
-
-    @override
-    def create_scheduler_impl(self, optimizer, lightning_module):
-        num_steps_per_epoch = self.compute_num_steps_per_epoch(lightning_module)
-        warmup_steps = self.warmup_duration.to_steps(num_steps_per_epoch).value
-        max_steps = self.max_duration.to_steps(num_steps_per_epoch).value
-
-        # Warmup and max steps should be at least 1.
-        warmup_steps = max(warmup_steps, 1)
-        max_steps = max(max_steps, 1)
-
-        # Create the scheduler
-        scheduler = LinearWarmupCosineAnnealingLR(
-            optimizer=optimizer,
-            warmup_epochs=warmup_steps,
-            max_epochs=max_steps,
-            warmup_start_lr_factor=self.warmup_start_lr_factor,
-            eta_min_factor=self.min_lr_factor,
-            should_restart=self.annealing,
-        )
-        return scheduler

nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -4,12 +4,14 @@ from typing import Literal
 
 from lightning.pytorch.utilities.types import LRSchedulerConfigType
 from torch.optim.lr_scheduler import ReduceLROnPlateau
-from typing_extensions import override
+from typing_extensions import final, override
 
 from ..metrics._config import MetricConfig
-from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
+from .base import LRSchedulerConfigBase, LRSchedulerMetadata, lr_scheduler_registry
 
 
+@final
+@lr_scheduler_registry.register
 class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
     """Reduce learning rate when a metric has stopped improving."""
 

nshtrainer/nn/__init__.py

@@ -6,12 +6,12 @@ from .mlp import MLPConfigDict as MLPConfigDict
 from .mlp import ResidualSequential as ResidualSequential
 from .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList
-from .nonlinearity import BaseNonlinearityConfig as BaseNonlinearityConfig
 from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
 from .nonlinearity import GELUNonlinearityConfig as GELUNonlinearityConfig
 from .nonlinearity import LeakyReLUNonlinearityConfig as LeakyReLUNonlinearityConfig
 from .nonlinearity import MishNonlinearityConfig as MishNonlinearityConfig
 from .nonlinearity import NonlinearityConfig as NonlinearityConfig
+from .nonlinearity import NonlinearityConfigBase as NonlinearityConfigBase
 from .nonlinearity import PReLUConfig as PReLUConfig
 from .nonlinearity import ReLUNonlinearityConfig as ReLUNonlinearityConfig
 from .nonlinearity import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig

nshtrainer/nn/mlp.py

@@ -9,7 +9,7 @@ import torch
 import torch.nn as nn
 from typing_extensions import TypedDict, override
 
-from .nonlinearity import BaseNonlinearityConfig, NonlinearityConfig
+from .nonlinearity import NonlinearityConfig, NonlinearityConfigBase
 
 
 @runtime_checkable
@@ -92,11 +92,11 @@ class MLPConfig(C.Config):
 
 def MLP(
     dims: Sequence[int],
-    activation: BaseNonlinearityConfig
+    activation: NonlinearityConfigBase
     | nn.Module
     | Callable[[], nn.Module]
     | None = None,
-    nonlinearity: BaseNonlinearityConfig
+    nonlinearity: NonlinearityConfigBase
     | nn.Module
     | Callable[[], nn.Module]
     | None = None,
@@ -153,7 +153,7 @@ def MLP(
             layers.append(nn.Dropout(dropout))
         if i < len(dims) - 2:
             match activation:
-                case BaseNonlinearityConfig():
+                case NonlinearityConfigBase():
                     layers.append(activation.create_module())
                 case nn.Module():
                     # In this case, we create a deep copy of the module to avoid sharing parameters (if any).

nshtrainer/nn/nonlinearity.py

@@ -7,10 +7,10 @@ import nshconfig as C
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from typing_extensions import final, override
+from typing_extensions import TypeAliasType, final, override
 
 
-class BaseNonlinearityConfig(C.Config, ABC):
+class NonlinearityConfigBase(C.Config, ABC):
     @abstractmethod
     def create_module(self) -> nn.Module: ...
 
@@ -18,8 +18,12 @@ class BaseNonlinearityConfig(C.Config, ABC):
     def __call__(self, x: torch.Tensor) -> torch.Tensor: ...
 
 
+nonlinearity_registry = C.Registry(NonlinearityConfigBase, discriminator="name")
+
+
 @final
-class ReLUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class ReLUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["relu"] = "relu"
 
     @override
@@ -31,7 +35,8 @@ class ReLUNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SigmoidNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SigmoidNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["sigmoid"] = "sigmoid"
 
     @override
@@ -43,7 +48,8 @@ class SigmoidNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class TanhNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class TanhNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["tanh"] = "tanh"
 
     @override
@@ -55,7 +61,8 @@ class TanhNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SoftmaxNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SoftmaxNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["softmax"] = "softmax"
 
     dim: int = -1
@@ -70,7 +77,8 @@ class SoftmaxNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SoftplusNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SoftplusNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["softplus"] = "softplus"
 
     beta: float = 1.0
@@ -88,7 +96,8 @@ class SoftplusNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SoftsignNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SoftsignNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["softsign"] = "softsign"
 
     @override
@@ -100,7 +109,8 @@ class SoftsignNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class ELUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class ELUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["elu"] = "elu"
 
     alpha: float = 1.0
@@ -115,7 +125,8 @@ class ELUNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class LeakyReLUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class LeakyReLUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["leaky_relu"] = "leaky_relu"
 
     negative_slope: float = 1.0e-2
@@ -130,7 +141,8 @@ class LeakyReLUNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class PReLUConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class PReLUConfig(NonlinearityConfigBase):
     name: Literal["prelu"] = "prelu"
 
     num_parameters: int = 1
@@ -152,7 +164,8 @@ class PReLUConfig(BaseNonlinearityConfig):
 
 
 @final
-class GELUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class GELUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["gelu"] = "gelu"
 
     approximate: Literal["tanh", "none"] = "none"
@@ -167,7 +180,8 @@ class GELUNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SwishNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SwishNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["swish"] = "swish"
 
     @override
@@ -179,7 +193,8 @@ class SwishNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class SiLUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SiLUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["silu"] = "silu"
 
     @override
@@ -191,7 +206,8 @@ class SiLUNonlinearityConfig(BaseNonlinearityConfig):
 
 
 @final
-class MishNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class MishNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["mish"] = "mish"
 
     @override
@@ -210,7 +226,8 @@ class SwiGLU(nn.SiLU):
 
 
 @final
-class SwiGLUNonlinearityConfig(BaseNonlinearityConfig):
+@nonlinearity_registry.register
+class SwiGLUNonlinearityConfig(NonlinearityConfigBase):
     name: Literal["swiglu"] = "swiglu"
 
     @override
@@ -222,20 +239,7 @@ class SwiGLUNonlinearityConfig(BaseNonlinearityConfig):
         return input * F.silu(gate)
 
 
-NonlinearityConfig = Annotated[
-    ReLUNonlinearityConfig
-    | SigmoidNonlinearityConfig
-    | TanhNonlinearityConfig
-    | SoftmaxNonlinearityConfig
-    | SoftplusNonlinearityConfig
-    | SoftsignNonlinearityConfig
-    | ELUNonlinearityConfig
-    | LeakyReLUNonlinearityConfig
-    | PReLUConfig
-    | GELUNonlinearityConfig
-    | SwishNonlinearityConfig
-    | SiLUNonlinearityConfig
-    | MishNonlinearityConfig
-    | SwiGLUNonlinearityConfig,
-    C.Field(discriminator="name"),
-]
+NonlinearityConfig = TypeAliasType(
+    "NonlinearityConfig",
+    Annotated[NonlinearityConfigBase, nonlinearity_registry.DynamicResolution()],
+)

nshtrainer/optimizer.py

@@ -7,7 +7,7 @@ from typing import Annotated, Any, Literal
 import nshconfig as C
 import torch.nn as nn
 from torch.optim import Optimizer
-from typing_extensions import TypeAliasType, override
+from typing_extensions import TypeAliasType, final, override
 
 
 class OptimizerConfigBase(C.Config, ABC):
@@ -18,6 +18,11 @@ class OptimizerConfigBase(C.Config, ABC):
     ) -> Optimizer: ...
 
 
+optimizer_registry = C.Registry(OptimizerConfigBase, discriminator="name")
+
+
+@final
+@optimizer_registry.register
 class AdamWConfig(OptimizerConfigBase):
     name: Literal["adamw"] = "adamw"
 
@@ -58,5 +63,6 @@ class AdamWConfig(OptimizerConfigBase):
 
 
 OptimizerConfig = TypeAliasType(
-    "OptimizerConfig", Annotated[AdamWConfig, C.Field(discriminator="name")]
+    "OptimizerConfig",
+    Annotated[OptimizerConfigBase, optimizer_registry.DynamicResolution()],
 )

nshtrainer/trainer/__init__.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
+from ..callbacks import callback_registry as callback_registry
 from ._config import TrainerConfig as TrainerConfig
-from ._config import accelerator_registry as accelerator_registry
-from ._config import plugin_registry as plugin_registry
+from .accelerator import accelerator_registry as accelerator_registry
+from .plugin import plugin_registry as plugin_registry
 from .trainer import Trainer as Trainer

nshtrainer/trainer/_config.py

@@ -37,7 +37,6 @@ from ..callbacks import (
     OnExceptionCheckpointCallbackConfig,
 )
 from ..callbacks.base import CallbackConfigBase
-from ..callbacks.checkpoint.time_checkpoint import TimeCheckpointCallbackConfig
 from ..callbacks.debug_flag import DebugFlagCallbackConfig
 from ..callbacks.log_epoch import LogEpochCallbackConfig
 from ..callbacks.lr_monitor import LearningRateMonitorConfig
@@ -49,14 +48,14 @@ from ..loggers import (
     TensorboardLoggerConfig,
     WandbLoggerConfig,
 )
-from ..loggers._base import BaseLoggerConfig
 from ..loggers.actsave import ActSaveLoggerConfig
+from ..loggers.base import LoggerConfigBase
 from ..metrics._config import MetricConfig
 from ..profiler import ProfilerConfig
 from ..util._environment_info import EnvironmentConfig
-from .accelerator import AcceleratorConfig, AcceleratorLiteral, accelerator_registry
-from .plugin import PluginConfig, plugin_registry
-from .strategy import StrategyConfig
+from .accelerator import AcceleratorConfig, AcceleratorLiteral
+from .plugin import PluginConfig
+from .strategy import StrategyConfig, StrategyLiteral
 
 log = logging.getLogger(__name__)
 
@@ -70,46 +69,12 @@ class GradientClippingConfig(C.Config):
     """Norm type to use for gradient clipping."""
 
 
-StrategyLiteral = TypeAliasType(
-    "StrategyLiteral",
-    Literal[
-        "auto",
-        "ddp",
-        "ddp_find_unused_parameters_false",
-        "ddp_find_unused_parameters_true",
-        "ddp_spawn",
-        "ddp_spawn_find_unused_parameters_false",
-        "ddp_spawn_find_unused_parameters_true",
-        "ddp_fork",
-        "ddp_fork_find_unused_parameters_false",
-        "ddp_fork_find_unused_parameters_true",
-        "ddp_notebook",
-        "dp",
-        "deepspeed",
-        "deepspeed_stage_1",
-        "deepspeed_stage_1_offload",
-        "deepspeed_stage_2",
-        "deepspeed_stage_2_offload",
-        "deepspeed_stage_3",
-        "deepspeed_stage_3_offload",
-        "deepspeed_stage_3_offload_nvme",
-        "fsdp",
-        "fsdp_cpu_offload",
-        "single_xla",
-        "xla_fsdp",
-        "xla",
-        "single_tpu",
-    ],
-)
-
-
 CheckpointCallbackConfig = TypeAliasType(
     "CheckpointCallbackConfig",
     Annotated[
         BestCheckpointCallbackConfig
         | LastCheckpointCallbackConfig
-        | OnExceptionCheckpointCallbackConfig
-        | TimeCheckpointCallbackConfig,
+        | OnExceptionCheckpointCallbackConfig,
         C.Field(discriminator="name"),
     ],
 )
@@ -123,7 +88,6 @@ class CheckpointSavingConfig(CallbackConfigBase):
         BestCheckpointCallbackConfig(throw_on_no_metric=False),
         LastCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
-        TimeCheckpointCallbackConfig(interval=timedelta(hours=12)),
     ]
     """Checkpoint callback configurations."""
 
@@ -397,8 +361,6 @@ class SanityCheckingConfig(C.Config):
     """
 
 
-@plugin_registry.rebuild_on_registers
-@accelerator_registry.rebuild_on_registers
 class TrainerConfig(C.Config):
     # region Active Run Configuration
     id: str = C.Field(default_factory=lambda: TrainerConfig.generate_id())
@@ -808,7 +770,7 @@ class TrainerConfig(C.Config):
         yield self.auto_set_debug_flag
         yield from self.callbacks
 
-    def _nshtrainer_all_logger_configs(self) -> Iterable[BaseLoggerConfig | None]:
+    def _nshtrainer_all_logger_configs(self) -> Iterable[LoggerConfigBase | None]:
         # Disable all loggers if barebones mode is enabled
         if self.barebones:
             return

{nshtrainer-1.0.0b33.dist-info → nshtrainer-1.0.0b37.dist-info}/METADATA

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