nshtrainer 0.29.1__tar.gz → 0.30.0__tar.gz

{nshtrainer-0.29.1 → nshtrainer-0.30.0}/PKG-INFO

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

{nshtrainer-0.29.1 → nshtrainer-0.30.0}/pyproject.toml

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

{nshtrainer-0.29.1 → nshtrainer-0.30.0}/src/nshtrainer/__init__.py

@@ -1,5 +1,6 @@
 from . import _experimental as _experimental
 from . import callbacks as callbacks
+from . import config as config
 from . import data as data
 from . import lr_scheduler as lr_scheduler
 from . import metrics as metrics

nshtrainer-0.30.0/src/nshtrainer/config/__init__.py

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

nshtrainer-0.30.0/src/nshtrainer/config/duration.py

@@ -0,0 +1,31 @@
+import math
+from typing import Annotated, Literal
+
+import nshconfig as C
+
+
+class Steps(C.Config):
+    kind: Literal["steps"] = "steps"
+
+    value: Annotated[int, C.Field(ge=0)]
+    """Number of steps."""
+
+    def to_steps(self, steps_per_epoch: int):
+        return self
+
+
+class Epochs(C.Config):
+    kind: Literal["epochs"] = "epochs"
+
+    value: Annotated[int | float, C.Field(ge=0)]
+    """Number of epochs."""
+
+    def to_steps(self, steps_per_epoch: int):
+        value = self.value * steps_per_epoch
+        if not isinstance(value, int):
+            value = int(math.ceil(value))
+
+        return Steps(value=value)
+
+
+Duration = Annotated[Steps | Epochs, C.Field(discriminator="kind")]

{nshtrainer-0.29.1 → nshtrainer-0.30.0}/src/nshtrainer/lr_scheduler/linear_warmup_cosine.py

@@ -2,11 +2,11 @@ import math
 import warnings
 from typing import Literal
 
-import nshconfig as C
 from torch.optim import Optimizer
 from torch.optim.lr_scheduler import LRScheduler
 from typing_extensions import override
 
+from ..config import Duration
 from ._base import LRSchedulerConfigBase, LRSchedulerMetadata
 
 
@@ -91,13 +91,13 @@ class LinearWarmupCosineAnnealingLR(LRScheduler):
 class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
     name: Literal["linear_warmup_cosine_decay"] = "linear_warmup_cosine_decay"
 
-    warmup_epochs: int = C.Field(ge=0)
-    r"""The number of epochs for the linear warmup phase.
-    The learning rate is linearly increased from `warmup_start_lr` to the initial learning rate over this number of epochs."""
+    warmup_duration: Duration
+    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_epochs: int = C.Field(gt=0)
-    r"""The total number of epochs.
-    The learning rate is decayed to `min_lr` over this number of epochs."""
+    max_duration: Duration
+    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.
@@ -121,11 +121,20 @@ class LinearWarmupCosineDecayLRSchedulerConfig(LRSchedulerConfigBase):
     @override
     def create_scheduler_impl(self, optimizer, lightning_module, lr):
         num_steps_per_epoch = self.compute_num_steps_per_epoch(lightning_module)
-        warmup_steps = self.warmup_epochs * num_steps_per_epoch
-        max_steps = self.max_epochs * num_steps_per_epoch
+        warmup_steps = (
+            self.warmup_duration.to_steps(num_steps_per_epoch).value
+            * num_steps_per_epoch
+        )
+        max_steps = (
+            self.max_duration.to_steps(num_steps_per_epoch).value * num_steps_per_epoch
+        )
         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)
+        max_steps = max(max_steps, 1)
+
         # Create the scheduler
         scheduler = LinearWarmupCosineAnnealingLR(
             optimizer=optimizer,