nshtrainer 1.0.0b36__py3-none-any.whl → 1.0.0b39__py3-none-any.whl

nshtrainer/lr_scheduler/__init__.py

@@ -5,8 +5,8 @@ from typing import Annotated
 import nshconfig as C
 from typing_extensions import TypeAliasType
 
-from ._base import LRSchedulerConfigBase as LRSchedulerConfigBase
-from ._base import LRSchedulerMetadata as LRSchedulerMetadata
+from .base import LRSchedulerConfigBase as LRSchedulerConfigBase
+from .base import LRSchedulerMetadata as LRSchedulerMetadata
 from .linear_warmup_cosine import (
     LinearWarmupCosineAnnealingLR as LinearWarmupCosineAnnealingLR,
 )

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/_config.py

@@ -48,8 +48,8 @@ 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
@@ -770,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.0b36.dist-info → nshtrainer-1.0.0b39.dist-info}/METADATA

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

{nshtrainer-1.0.0b36.dist-info → nshtrainer-1.0.0b39.dist-info}/RECORD

@@ -3,7 +3,7 @@ nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
 nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
 nshtrainer/_checkpoint/metadata.py,sha256=PHy-54Cg-o3OtCffAqrVv6ZVMU7zhRo_-sZiSEEno1Y,5019
 nshtrainer/_checkpoint/saver.py,sha256=LOP8jjKF0Dw9x9H-BKrLMWlEp1XTan2DUK0zQUCWw5U,1360
-nshtrainer/_directory.py,sha256=p2uk1FnISFEpMqlDevKhoWhQsCEtvHUPg459K-86QA8,3053
+nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
 nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
 nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
 nshtrainer/callbacks/__init__.py,sha256=4giOYT8A709UOLRtQEt16QbOAFUHCjJ_aLB7ITTwXJI,3577
@@ -12,7 +12,7 @@ nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
 nshtrainer/callbacks/checkpoint/_base.py,sha256=ZVEUVl5kjCSSe69Q0rMUbKBNNUog0pxBwWkeyuxG2w0,6304
 nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=2CQuhPJ3Fi7lDw7z-J8kXXXuDU8-4HcU48oZxR49apk,2667
-nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=MJcNB0biOebx2si2IBFaSUiVOSLSCZTzxB-RcEgO2gY,3482
+nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
 nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
 nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
@@ -30,7 +30,7 @@ nshtrainer/callbacks/shared_parameters.py,sha256=s94jJTAIbDGukYJu6l247QonVOCudGC
 nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU,4731
 nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
-nshtrainer/configs/__init__.py,sha256=OevZEZxb4H8imadSQXK9huqdYUF4SrJPfNU_2fpMBvI,14084
+nshtrainer/configs/__init__.py,sha256=MZfcSKhnjtVObBvVv9lu8L2cFTLINP5zcTQvWnz8jdk,14505
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
@@ -58,29 +58,29 @@ nshtrainer/configs/callbacks/shared_parameters/__init__.py,sha256=AU7_bSnSRSlj16
 nshtrainer/configs/callbacks/timer/__init__.py,sha256=cOUtbsl0_OhCO0fIcBfLuIF6FEGBHQu7AvQFzwVznWQ,413
 nshtrainer/configs/callbacks/wandb_upload_code/__init__.py,sha256=CJeCc9OCu5F39lWiY5aIc4WxQlgBvB-8cga6cQtw0GQ,482
 nshtrainer/configs/callbacks/wandb_watch/__init__.py,sha256=dzz1oavL1BwELE33xus45_avBEAZDeB6xtcb6CsOEos,431
-nshtrainer/configs/loggers/__init__.py,sha256=5wTekL79mQxit8f1K3AMllvb0mKertTzOKfC3gpE2Zk,1251
-nshtrainer/configs/loggers/_base/__init__.py,sha256=HxPPPePsEjlNuhnjsMgYIl0rwj_iqNKKOBTEk_zIOsM,169
-nshtrainer/configs/loggers/actsave/__init__.py,sha256=2lZQ4bpbjwd4MuUE_Z_PGbmQjjGtWCZUCtXqKO4dTSc,280
-nshtrainer/configs/loggers/csv/__init__.py,sha256=M3QGF5GKiRGENy3re6LJKpa4A4RThy1FlmaFuR4cPyo,260
-nshtrainer/configs/loggers/tensorboard/__init__.py,sha256=FbkYXnSohIX6JN5XyI-9y91IJv_T3VB3IwmpagXAnM4,309
-nshtrainer/configs/loggers/wandb/__init__.py,sha256=76qb0HhWojf0Ub1x9OkMjtzeXxE67KysBGa-MBbJyC4,651
-nshtrainer/configs/lr_scheduler/__init__.py,sha256=8ORO-QC12SjZ2F_reMoDgr8-O8nxZxX0IKU4fl-cC3A,1023
-nshtrainer/configs/lr_scheduler/_base/__init__.py,sha256=fvGjkUJ1K2RVXjXror22QOtEa-xWFJz2Cz3HrBC5XfA,189
-nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=i8LeZh0c4wqtZ1ehZb2LCq7kwOL0OyswMMOnwyI6R04,533
-nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=lpXEFZY4cM3znZqYG9IZ1xNNtzttt8VVspSuOz0fb-k,467
+nshtrainer/configs/loggers/__init__.py,sha256=GT7PO7UM3Mo87N616mGucc2ZRyGP8nQWBd_VJ_8RGXo,1337
+nshtrainer/configs/loggers/actsave/__init__.py,sha256=J7SnbD-zxUynWSskJezooFyBZdnhgTWyybRvwn9gzy4,377
+nshtrainer/configs/loggers/base/__init__.py,sha256=HLUfEDbjaAXqzsFmQbjdciIWzR1st1gRLKTCFvUFEX0,262
+nshtrainer/configs/loggers/csv/__init__.py,sha256=gawaDX92JObGSmBqYpfNHWMHBwVOofS694W-1Y2GWDU,353
+nshtrainer/configs/loggers/tensorboard/__init__.py,sha256=phzm-TnBkdkibTgoOxIIcAliqL3zU8gSNK61Mwxs1CM,410
+nshtrainer/configs/loggers/wandb/__init__.py,sha256=TDcD5WZSKenc2mgIXhwz2l96l8P_Ur3N5CzEol5AKGw,746
+nshtrainer/configs/lr_scheduler/__init__.py,sha256=xtiUx0isxA82-uXMn4-KmPnDCfbUkpAnd2_pFupAAKQ,1137
+nshtrainer/configs/lr_scheduler/base/__init__.py,sha256=6Cx8r4rdxeSYxc_z0o7drKCblGJU_zzqrOoYlWYR5qY,305
+nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=5ZMLDO9VL6SNU6pF-62lDnpmqix3_Ol9DdEwiuOPYlA,675
+nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=w-vq8UbRGPX8DZVWCMC5eIrbvVc_guxjj7Du9AaeKCw,609
 nshtrainer/configs/metrics/__init__.py,sha256=mK_xgXJDAyGY4K_x_Zo4aj36kjT45d850keuUe3U1rY,200
 nshtrainer/configs/metrics/_config/__init__.py,sha256=XDDvDPWULd_vd3lrgF2KGAVR2LVDuhdQvy-fF2ImarI,159
-nshtrainer/configs/nn/__init__.py,sha256=3hVc81Gs9AJYVkrwJkQ_ye7tLU2HOLdBj-mMkXx2c_I,1957
-nshtrainer/configs/nn/mlp/__init__.py,sha256=eMECrgz-My9mFS7lpWVI3dj1ApB-E7xwfmNc37hUsPI,347
-nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=Gjr2HCx8jJTcfu7sLgn54o2ucGKaBea4encm4AWpKNY,2040
-nshtrainer/configs/optimizer/__init__.py,sha256=IMEsEbiVFXSkj6WmDjNjmKQuRspphs5xZnYZ2gYE39Y,344
+nshtrainer/configs/nn/__init__.py,sha256=tkFG2Hb0oL_AmWP3_0WkDN2zI5PkVfrgwXhaAII7CZw,2072
+nshtrainer/configs/nn/mlp/__init__.py,sha256=O6kQ6utZNJPG9Fax5pRdZcHa3J-XFKKdXcc_PQg0jk0,347
+nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=LCTbTyelCMABVw505CGQ4UpEGlAnIhflSLFqwAQXLQA,2155
+nshtrainer/configs/optimizer/__init__.py,sha256=itIDIHQvGm50eZ7JLyNElahnNUMPJ__4PMmTjc0RQ6o,444
 nshtrainer/configs/profiler/__init__.py,sha256=2ssaIpfVnvcbfNvZ-JeKp1Cx4NO1LknkVqTm1hu7Lvw,768
 nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRDm1CKqjwUOQNbQjD4,176
 nshtrainer/configs/profiler/advanced/__init__.py,sha256=-ThpUat16Ij_0avkMUVVA8wCWDG_q_tM7KQofnWQCtg,308
 nshtrainer/configs/profiler/pytorch/__init__.py,sha256=soAU1s2_Pa1na4gW8CK-iysJBO5M_7YeZC2_x40iEdg,294
 nshtrainer/configs/profiler/simple/__init__.py,sha256=3Wb11lPuFuyasq8xS1CZ4WLuBCLS_nVSQGVllvOOi0Y,289
-nshtrainer/configs/trainer/__init__.py,sha256=8Z4E1IeJHtDW8fpDxJkiC9CgDqKrTBIR5VMK1q4DYy4,7729
-nshtrainer/configs/trainer/_config/__init__.py,sha256=t72kmUn60UtjpD6H38XzKbEs50gU2dS1IH0u-RnHZ04,3666
+nshtrainer/configs/trainer/__init__.py,sha256=jYCp4Q9uvutA6NYqfthbREMg09-obD3gHtzEI2Ta-hU,7729
+nshtrainer/configs/trainer/_config/__init__.py,sha256=uof_oJfhwjB1pft7KsRdk_RvNj-tE8wcDBEM7X5qtNc,3666
 nshtrainer/configs/trainer/accelerator/__init__.py,sha256=3H6R3wlwbKL1TzDqGCChZk78-BcE2czLouo7Djiq3nA,898
 nshtrainer/configs/trainer/plugin/__init__.py,sha256=NkHQxMPkrtTtdIAO4dQUE9SWEcHRDB0yUXLkTjnl4dA,3332
 nshtrainer/configs/trainer/plugin/base/__init__.py,sha256=GuubcKrbXt4VjJRT8VpNUQqBtyuutre_CfkS0EWZ5_E,368
@@ -99,16 +99,16 @@ nshtrainer/data/__init__.py,sha256=K4i3Tw4g9EOK2zlMMbidi99y0SyI4w8P7_XUf1n42Ts,2
 nshtrainer/data/balanced_batch_sampler.py,sha256=r1cBKRXKHD8E1Ax6tj-FUbE-z1qpbO58mQ9VrK9uLnc,5481
 nshtrainer/data/datamodule.py,sha256=lSOgH32nysJWa6Y7ba1QyOdUV0DVVdO98qokP8wigjk,4138
 nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
-nshtrainer/loggers/__init__.py,sha256=-y8B-9TF6vJdZUQewJNDcZ2aOv04FEUFtKwaiDobIO0,670
-nshtrainer/loggers/_base.py,sha256=nw4AZzJP3Z-fljgQlgq7FkuMkPmYKTsXj7OfJJSmtXI,811
-nshtrainer/loggers/actsave.py,sha256=23Kre-mq-Y9Iw1SRyGmHnHK1bc_0gTWFXJViv9bkVz0,1324
-nshtrainer/loggers/csv.py,sha256=Deh5gm3oROJbQzigV4SHni5JRwSrBdm-4YD3yrcGnHo,1104
-nshtrainer/loggers/tensorboard.py,sha256=jP9V4nlq_MXUaoD6xv1Cws2ioft83Lm8yUJhhGhuUrQ,2268
-nshtrainer/loggers/wandb.py,sha256=EjKQQznLSUSCWO7uIviz9g0dVW4ZLxb_8UVhY4vR7r0,6800
-nshtrainer/lr_scheduler/__init__.py,sha256=BGnO-okUTZOtF15-UmQ05U4oEatSF5VNs3YeidNEWn4,853
-nshtrainer/lr_scheduler/_base.py,sha256=EhA2f_WiZ79RcXL2nJbwCwNK620c8ugEVUmJ8CcVpsE,3723
-nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=gvUuv031lvWdXboDeH7iAF3ZgNPQK40bQwfmqb11TNk,5492
-nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=vXH5S26ESHO_LPPqW8aDC3S5NGoZYkXeFjAOgttaUX8,2870
+nshtrainer/loggers/__init__.py,sha256=Ddd3JJXVzew_ZpwHA9kGnGmvq4OwhItwghDL5PzNhDc,614
+nshtrainer/loggers/actsave.py,sha256=wgNrpBB6wQM7qff8iLDb_sQnbiAcYHRmH56pcEJPB3o,1409
+nshtrainer/loggers/base.py,sha256=ON92XbwTSgadQOSyw5PiRRFzyH6uJ-xLtE0nB3cbgPc,1205
+nshtrainer/loggers/csv.py,sha256=xJ8mSmw4vJwinIfqhF6t2HWmh_1dXEYyLfGuXwL7WHo,1160
+nshtrainer/loggers/tensorboard.py,sha256=E7iO_fDt9bfH02hBL430bXPLljOo5iGgq2QyPqmx2gQ,2324
+nshtrainer/loggers/wandb.py,sha256=KZXAUWrrmdX_L8rqej77oUHaM0JxZRM8y9z6JP9PISw,6856
+nshtrainer/lr_scheduler/__init__.py,sha256=daMMK3erUcNXGGd_nZB8AWu3ZTYqfS1RSWeK4FV2udw,851
+nshtrainer/lr_scheduler/base.py,sha256=062fGcH5sYeEKwoY55RydCTvfPwTnyZHCi049a3nMbM,3805
+nshtrainer/lr_scheduler/linear_warmup_cosine.py,sha256=MsoXgCcWTKsrkNZiGnKS6yC-slRuleuwFxeM_lmG_pQ,5560
+nshtrainer/lr_scheduler/reduce_lr_on_plateau.py,sha256=v9T0GpvOoHV30atFB0MwExHgHcTpMCYxbMRoPjPBjt8,2938
 nshtrainer/metrics/__init__.py,sha256=Nqkn_jsDf3n5WtfMcnaaEftYjIIT2b-S7rmsB1MOMkU,86
 nshtrainer/metrics/_config.py,sha256=XIRokFM8PHrhBa3w2R6BM6a4es3ncsoBqE_LqXQFsFE,1223
 nshtrainer/model/__init__.py,sha256=3G-bwPPSRStWdsdwG9-rn0bXcRpEiP1BiQpF_qavtls,97
@@ -116,19 +116,19 @@ nshtrainer/model/base.py,sha256=JL3AmH17GQjQIoMrZl3O0vUI7dj5ZsO5iEJgoLPyzHw,1035
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
 nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
-nshtrainer/nn/__init__.py,sha256=sANhrZpeN5syLKOsmXMwhaFl2SBFPWcLaEe1EH22TWQ,1463
-nshtrainer/nn/mlp.py,sha256=2W8bzE96DzCMzGm6WPiPhNFQfhqaoG3GXPn_oKBnlUM,5988
+nshtrainer/nn/__init__.py,sha256=7KCs-GDOynCXAIdwkgAQacc0p3FHLEION50UtrvgAOc,1463
+nshtrainer/nn/mlp.py,sha256=ZbkLyOc08stgIugvu1G5_h66DYtxAFDnboikBaJvvZ8,5988
 nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
-nshtrainer/nn/nonlinearity.py,sha256=mp5XvXRHURB6jwuZ0YyTj5ZoHJYNJNgO2aLtUY1D-2Y,6114
-nshtrainer/optimizer.py,sha256=wmSRpSoU59rstj2RBoifQ15ZwRInYpm0tDBQZ1gqOfE,1596
+nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
+nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
 nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5NOQ,1160
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=fQ7gQRlGWX-90TYT0rttkQyvXDCzo7DAvJgr-jX1zsY,316
-nshtrainer/trainer/_config.py,sha256=SPg3WXjF3ufhnr27sTHQLq23hdebnW6CTWa8AJkRG0A,32982
+nshtrainer/trainer/_config.py,sha256=QDy6sINVDGEqfHfPTWXSN-06EoEuMSVscHn8fCRTvr0,32981
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
 nshtrainer/trainer/accelerator.py,sha256=Bqq-ry7DeCY4zw9_zBvTZiijpA-uUHrDjtbLV652m4M,2415
 nshtrainer/trainer/plugin/__init__.py,sha256=UM8f70Ml3RGpsXeQr1Yh1yBcxxjFNGFGgJbniCn_rws,366
@@ -151,6 +151,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.0b36.dist-info/METADATA,sha256=R9O2SnflaNiDkxtoOPD_YFCXIgnEl8YjkhbEU5CbWHQ,988
-nshtrainer-1.0.0b36.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-nshtrainer-1.0.0b36.dist-info/RECORD,,
+nshtrainer-1.0.0b39.dist-info/METADATA,sha256=zzE6nHlj-clB3HJs5_-bBePCHSOrtTkZTi9z_NrSeRY,988
+nshtrainer-1.0.0b39.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+nshtrainer-1.0.0b39.dist-info/RECORD,,

nshtrainer/configs/loggers/_base/__init__.py

@@ -1,9 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from nshtrainer.loggers._base import BaseLoggerConfig as BaseLoggerConfig
-
-__all__ = [
-    "BaseLoggerConfig",
-]

nshtrainer/configs/lr_scheduler/_base/__init__.py

@@ -1,9 +0,0 @@
-from __future__ import annotations
-
-__codegen__ = True
-
-from nshtrainer.lr_scheduler._base import LRSchedulerConfigBase as LRSchedulerConfigBase
-
-__all__ = [
-    "LRSchedulerConfigBase",
-]