nshtrainer 0.42.0__py3-none-any.whl → 0.44.0__py3-none-any.whl

nshtrainer/config/util/_environment_info/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING
@@ -43,30 +45,30 @@ else:
 
         if name in globals():
             return globals()[name]
-        if name == "EnvironmentPackageConfig":
-            return importlib.import_module(
-                "nshtrainer.util._environment_info"
-            ).EnvironmentPackageConfig
-        if name == "EnvironmentSnapshotConfig":
+        if name == "EnvironmentLinuxEnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSnapshotConfig
+            ).EnvironmentLinuxEnvironmentConfig
         if name == "EnvironmentLSFInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).EnvironmentLSFInformationConfig
-        if name == "EnvironmentLinuxEnvironmentConfig":
+        if name == "EnvironmentGPUConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentLinuxEnvironmentConfig
-        if name == "EnvironmentSLURMInformationConfig":
+            ).EnvironmentGPUConfig
+        if name == "EnvironmentPackageConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentSLURMInformationConfig
-        if name == "EnvironmentConfig":
+            ).EnvironmentPackageConfig
+        if name == "EnvironmentHardwareConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentConfig
+            ).EnvironmentHardwareConfig
+        if name == "EnvironmentSnapshotConfig":
+            return importlib.import_module(
+                "nshtrainer.util._environment_info"
+            ).EnvironmentSnapshotConfig
         if name == "EnvironmentClassInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
@@ -75,18 +77,18 @@ else:
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
             ).GitRepositoryConfig
-        if name == "EnvironmentCUDAConfig":
+        if name == "EnvironmentConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentCUDAConfig
-        if name == "EnvironmentGPUConfig":
+            ).EnvironmentConfig
+        if name == "EnvironmentCUDAConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentGPUConfig
-        if name == "EnvironmentHardwareConfig":
+            ).EnvironmentCUDAConfig
+        if name == "EnvironmentSLURMInformationConfig":
             return importlib.import_module(
                 "nshtrainer.util._environment_info"
-            ).EnvironmentHardwareConfig
+            ).EnvironmentSLURMInformationConfig
         raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
 
 # Submodule exports

nshtrainer/config/util/config/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/config/util/config/dtype/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/config/util/config/duration/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 __codegen__ = True
 
 from typing import TYPE_CHECKING

nshtrainer/data/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from . import transform as dataset_transform
 from .balanced_batch_sampler import BalancedBatchSampler as BalancedBatchSampler
 from .datamodule import LightningDataModuleBase as LightningDataModuleBase

nshtrainer/data/balanced_batch_sampler.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import heapq
 import logging
 from typing import Any, Protocol, runtime_checkable

nshtrainer/data/datamodule.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from lightning.pytorch import LightningDataModule
 
 

nshtrainer/data/transform.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 from collections.abc import Callable
 from typing import Any, cast

nshtrainer/ll/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TypeAlias
 
 from . import _experimental as _experimental

nshtrainer/ll/_experimental.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer._experimental import *  # noqa: F403

nshtrainer/ll/actsave.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshutils.actsave import *  # type: ignore # noqa: F403
 
 from nshtrainer.callbacks.actsave import ActSaveCallback as ActSaveCallback

nshtrainer/ll/callbacks.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.callbacks import *  # noqa: F403

nshtrainer/ll/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshconfig import *  # type: ignore  # noqa: F403
 from nshconfig import Config as TypedConfig  # type: ignore  # noqa: F401
 

nshtrainer/ll/data.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.data import *  # noqa: F403

nshtrainer/ll/log.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshutils import init_python_logging as init_python_logging
 from nshutils import lovely as lovely
 from nshutils import pretty as pretty

nshtrainer/ll/lr_scheduler.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.lr_scheduler import *  # noqa: F403

nshtrainer/ll/model.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshtrainer.model import *  # noqa: F403
 
 from ..trainer._config import CheckpointLoadingConfig as CheckpointLoadingConfig

nshtrainer/ll/nn.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.nn import *  # noqa: F403

nshtrainer/ll/optimizer.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.optimizer import *  # noqa: F403

nshtrainer/ll/runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from nshrunner import SnapshotConfig as SnapshotConfig
 
 from nshtrainer.runner import *  # type: ignore # noqa: F403

nshtrainer/ll/snapshot.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshsnap import *  # pyright: ignore[reportWildcardImportFromLibrary] # noqa: F403

nshtrainer/ll/snoop.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshutils import snoop as snoop

nshtrainer/ll/trainer.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.trainer import *  # noqa: F403

nshtrainer/ll/typecheck.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshutils.typecheck import *  # type: ignore # noqa: F403

nshtrainer/ll/util.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from nshtrainer.util import *  # noqa: F403

nshtrainer/loggers/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/loggers/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING
 

nshtrainer/loggers/csv.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Literal
 
 from typing_extensions import override

nshtrainer/loggers/tensorboard.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/loggers/wandb.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import importlib.metadata
 import logging
 from typing import TYPE_CHECKING, Literal
@@ -8,8 +10,8 @@ from packaging import version
 from typing_extensions import assert_never, override
 
 from ..callbacks.base import CallbackConfigBase
-from ..callbacks.wandb_upload_code import WandbUploadCodeConfig
-from ..callbacks.wandb_watch import WandbWatchConfig
+from ..callbacks.wandb_upload_code import WandbUploadCodeCallbackConfig
+from ..callbacks.wandb_watch import WandbWatchCallbackConfig
 from ._base import BaseLoggerConfig
 
 if TYPE_CHECKING:
@@ -92,10 +94,10 @@ class WandbLoggerConfig(CallbackConfigBase, BaseLoggerConfig):
     - "none" or False: Do not log any checkpoints
     """
 
-    log_code: WandbUploadCodeConfig | None = WandbUploadCodeConfig()
+    log_code: WandbUploadCodeCallbackConfig | None = WandbUploadCodeCallbackConfig()
     """WandB code upload configuration. Used to upload code to WandB."""
 
-    watch: WandbWatchConfig | None = WandbWatchConfig()
+    watch: WandbWatchCallbackConfig | None = WandbWatchCallbackConfig()
     """WandB model watch configuration. Used to log model architecture, gradients, and parameters."""
 
     offline: bool = False

nshtrainer/lr_scheduler/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/lr_scheduler/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 from abc import ABC, abstractmethod
 from collections.abc import Mapping
@@ -9,7 +11,7 @@ from lightning.pytorch.utilities.types import (
     LRSchedulerTypeUnion,
 )
 from torch.optim import Optimizer
-from typing_extensions import NotRequired, TypedDict
+from typing_extensions import Never, NotRequired, TypedDict
 
 if TYPE_CHECKING:
     from ..model.base import LightningModuleBase
@@ -42,20 +44,17 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
     @abstractmethod
     def create_scheduler_impl(
-        self,
-        optimizer: Optimizer,
-        lightning_module: "LightningModuleBase",
-        lr: float,
+        self, optimizer: Optimizer, lightning_module: LightningModuleBase
     ) -> LRSchedulerTypeUnion | LRSchedulerConfigType: ...
 
     def create_scheduler(
         self,
         optimizer: Optimizer,
-        lightning_module: "LightningModuleBase",
-        lr: float,
+        lightning_module: LightningModuleBase,
+        lr: Never,  # Backward compatibility, should be removed in the future
     ) -> LRSchedulerConfigType:
         # Create the scheduler.
-        scheduler = self.create_scheduler_impl(optimizer, lightning_module, lr)
+        scheduler = self.create_scheduler_impl(optimizer, lightning_module)
 
         # If the scheduler is not a `LRSchedulerConfigType`, then make it one.
         if not isinstance(scheduler, Mapping):
@@ -87,9 +86,7 @@ class LRSchedulerConfigBase(C.Config, ABC):
 
         return scheduler
 
-    def compute_num_steps_per_epoch(
-        self, lightning_module: "LightningModuleBase"
-    ) -> int:
+    def compute_num_steps_per_epoch(self, lightning_module: LightningModuleBase) -> int:
         trainer = lightning_module.trainer
         # Use the Lightning trainer to convert the epoch-based values to step-based values
         _ = trainer.estimated_stepping_batches

nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import math
 import warnings
 from typing import Literal
@@ -18,21 +20,21 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
         optimizer: Optimizer,
         warmup_epochs: int,
         max_epochs: int,
-        warmup_start_lr: float = 0.0,
-        eta_min: float = 0.0,
+        warmup_start_lr_factor: float = 0.0,
+        eta_min_factor: float = 0.0,
         last_epoch: int = -1,
         should_restart: bool = True,
     ) -> None:
         self.warmup_epochs = warmup_epochs
         self.max_epochs = max_epochs
-        self.warmup_start_lr = warmup_start_lr
-        self.eta_min = eta_min
+        self.warmup_start_lr_factor = warmup_start_lr_factor
+        self.eta_min_factor = eta_min_factor
         self.should_restart = should_restart
 
         super().__init__(optimizer, last_epoch)
 
     @override
-    def get_lr(self) -> list[float]:  # pyright: ignore[reportIncompatibleMethodOverride]
+    def get_lr(self) -> list[float]:
         if not self._get_lr_called_within_step:
             warnings.warn(
                 "To get the last learning rate computed by the scheduler, "
@@ -41,25 +43,26 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
             )
 
         if self.last_epoch == 0:
-            return [self.warmup_start_lr] * len(self.base_lrs)
+            return [self.warmup_start_lr_factor * base_lr for base_lr in self.base_lrs]
         if self.last_epoch < self.warmup_epochs:
             return [
                 group["lr"]
-                + (base_lr - self.warmup_start_lr) / (self.warmup_epochs - 1)
+                + (base_lr - self.warmup_start_lr_factor * base_lr)
+                / (self.warmup_epochs - 1)
                 for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
             ]
         if self.last_epoch == self.warmup_epochs:
             return self.base_lrs
 
         if not self.should_restart and self.last_epoch >= self.max_epochs:
-            return [self.eta_min] * len(self.base_lrs)
+            return [self.eta_min_factor * base_lr for base_lr in self.base_lrs]
 
         if (self.last_epoch - 1 - self.max_epochs) % (
             2 * (self.max_epochs - self.warmup_epochs)
         ) == 0:
             return [
                 group["lr"]
-                + (base_lr - self.eta_min)
+                + (base_lr - self.eta_min_factor * base_lr)
                 * (1 - math.cos(math.pi / (self.max_epochs - self.warmup_epochs)))
                 / 2
                 for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
@@ -82,9 +85,9 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
                     / (self.max_epochs - self.warmup_epochs)
                 )
             )
-            * (group["lr"] - self.eta_min)
-            + self.eta_min
-            for group in self.optimizer.param_groups
+            * (group["lr"] - self.eta_min_factor * base_lr)
+            + self.eta_min_factor * base_lr
+            for base_lr, group in zip(self.base_lrs, self.optimizer.param_groups)
         ]
 
 
@@ -119,12 +122,10 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
         }
 
     @override
-    def create_scheduler_impl(self, optimizer, lightning_module, lr):
+    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_start_lr = self.warmup_start_lr_factor * lr
-        min_lr = self.min_lr_factor * lr
 
         # Warmup and max steps should be at least 1.
         warmup_steps = max(warmup_steps, 1)
@@ -135,8 +136,8 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
             optimizer=optimizer,
             warmup_epochs=warmup_steps,
             max_epochs=max_steps,
-            warmup_start_lr=warmup_start_lr,
-            eta_min=min_lr,
+            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

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import TYPE_CHECKING, Literal, cast
 
 from lightning.pytorch.utilities.types import LRSchedulerConfigType
@@ -20,21 +22,21 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
     """Metric to monitor.
     If not provided, the primary metric of the runner will be used."""
 
-    patience: int = 10
+    patience: int
     r"""Number of epochs with no improvement after which learning rate will be reduced."""
 
-    factor: float = 0.1
+    factor: float
     r"""Factor by which the learning rate will be reduced. new_lr = lr * factor."""
 
+    cooldown: int = 0
+    r"""Number of epochs to wait before resuming normal operation after lr has been reduced."""
+
     min_lr: float | list[float] = 0.0
     r"""A scalar or a list of scalars. A lower bound on the learning rate of all param groups or each group respectively."""
 
     eps: float = 1.0e-8
     r"""Minimal decay applied to lr. If the difference between new and old lr is smaller than eps, the update is ignored."""
 
-    cooldown: int = 0
-    r"""Number of epochs to wait before resuming normal operation after lr has been reduced."""
-
     threshold: float = 1.0e-4
     r"""Threshold for measuring the new optimum, to only focus on significant changes."""
 
@@ -43,7 +45,7 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
 
     @override
     def create_scheduler_impl(
-        self, optimizer, lightning_module, lr
+        self, optimizer, lightning_module
     ) -> LRSchedulerConfigType:
         if (metric := self.metric) is None:
             lm_config = cast("BaseConfig", lightning_module.config)

nshtrainer/metrics/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from ._config import MetricConfig as MetricConfig

nshtrainer/metrics/_config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import builtins
 from typing import Literal
 

nshtrainer/model/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
 from .config import DirectoryConfig as DirectoryConfig

nshtrainer/model/base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import inspect
 import logging
 from abc import ABC, abstractmethod

nshtrainer/model/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import logging
 import os

nshtrainer/model/mixins/callback.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from collections.abc import Callable, Iterable, Sequence
 from typing import Any, TypeAlias, cast, final

nshtrainer/model/mixins/logger.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections import deque
 from collections.abc import Callable, Generator
 from contextlib import contextmanager

nshtrainer/nn/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import MLPConfigDict as MLPConfigDict

nshtrainer/nn/mlp.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 from collections.abc import Callable, Sequence
 from typing import Literal, Protocol, runtime_checkable

nshtrainer/nn/module_dict.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Mapping
 from typing import Generic, cast
 

nshtrainer/nn/module_list.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections.abc import Iterable, Iterator
 from typing import Generic, TypeVar, overload
 

nshtrainer/nn/nonlinearity.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from typing import Annotated, Literal
 

nshtrainer/optimizer.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from abc import ABC, abstractmethod
 from collections.abc import Iterable
 from typing import Annotated, Any, Literal, TypeAlias

nshtrainer/profiler/__init__.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from typing import Annotated, TypeAlias
 
 import nshconfig as C

nshtrainer/profiler/_base.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path

nshtrainer/profiler/advanced.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/profiler/pytorch.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Any, Literal
 

nshtrainer/profiler/simple.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import logging
 from typing import Literal
 

nshtrainer/runner.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import logging
 from collections.abc import Callable, Iterable, Mapping, Sequence

nshtrainer/scripts/find_packages.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import argparse
 import ast
 import glob

nshtrainer/trainer/__init__.py

@@ -1 +1,3 @@
+from __future__ import annotations
+
 from .trainer import Trainer as Trainer