nshtrainer 1.4.0__tar.gz → 1.4.1__tar.gz

{nshtrainer-1.4.0 → nshtrainer-1.4.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.4.0
+Version: 1.4.1
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -14,7 +14,7 @@ Provides-Extra: extra
 Requires-Dist: GitPython ; extra == "extra"
 Requires-Dist: huggingface-hub ; extra == "extra"
 Requires-Dist: lightning
-Requires-Dist: nshconfig (>0.39)
+Requires-Dist: nshconfig (>=0.43)
 Requires-Dist: nshrunner ; extra == "extra"
 Requires-Dist: nshutils ; extra == "extra"
 Requires-Dist: numpy

{nshtrainer-1.4.0 → nshtrainer-1.4.1}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "nshtrainer"
-version = "1.4.0"
+version = "1.4.1"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"
 readme = "README.md"
 
 dependencies = [
-    "nshconfig>0.39",
+    "nshconfig>=0.43",
     "psutil",
     "numpy",
     "torch",

nshtrainer-1.4.1/src/nshtrainer/callbacks/log_epoch.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Literal
+
+from lightning.pytorch import LightningModule, Trainer
+from lightning.pytorch.callbacks import Callback
+from typing_extensions import final, override
+
+from .base import CallbackConfigBase, callback_registry
+
+log = logging.getLogger(__name__)
+
+
+@final
+@callback_registry.register
+class LogEpochCallbackConfig(CallbackConfigBase):
+    name: Literal["log_epoch"] = "log_epoch"
+
+    metric_name: str = "computed_epoch"
+    """The name of the metric to log the epoch as."""
+
+    train: bool = True
+    """Whether to log the epoch during training."""
+
+    val: bool = True
+    """Whether to log the epoch during validation."""
+
+    test: bool = True
+    """Whether to log the epoch during testing."""
+
+    @override
+    def create_callbacks(self, trainer_config):
+        yield LogEpochCallback(self)
+
+
+def _worker_fn(
+    trainer: Trainer,
+    pl_module: LightningModule,
+    num_batches_prop: str,
+    dataloader_idx: int | None = None,
+    *,
+    metric_name: str,
+):
+    if trainer.logger is None:
+        return
+
+    # If trainer.num_{training/val/test}_batches is not set or is nan/inf, we cannot calculate the epoch
+    if not (num_batches := getattr(trainer, num_batches_prop, None)):
+        log.warning(f"Trainer has no valid `{num_batches_prop}`. Cannot log epoch.")
+        return
+
+    # If the trainer has a dataloader_idx, num_batches is a list of num_batches for each dataloader.
+    if dataloader_idx is not None:
+        assert isinstance(num_batches, list), (
+            f"Expected num_batches to be a list, got {type(num_batches)}"
+        )
+        assert 0 <= dataloader_idx < len(num_batches), (
+            f"Expected dataloader_idx to be between 0 and {len(num_batches)}, got {dataloader_idx}"
+        )
+        num_batches = num_batches[dataloader_idx]
+
+    if (
+        not isinstance(num_batches, (int, float))
+        or math.isnan(num_batches)
+        or math.isinf(num_batches)
+    ):
+        log.warning(
+            f"Trainer has no valid `{num_batches_prop}` (got {num_batches=}). Cannot log epoch."
+        )
+        return
+
+    epoch = pl_module.global_step / num_batches
+    pl_module.log(metric_name, epoch, on_step=True, on_epoch=False)
+
+
+class LogEpochCallback(Callback):
+    def __init__(self, config: LogEpochCallbackConfig):
+        super().__init__()
+
+        self.config = config
+
+    @override
+    def on_train_batch_start(
+        self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
+    ):
+        if trainer.logger is None or not self.config.train:
+            return
+
+        _worker_fn(
+            trainer,
+            pl_module,
+            "num_training_batches",
+            metric_name=self.config.metric_name,
+        )
+
+    @override
+    def on_validation_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        if trainer.logger is None or not self.config.val:
+            return
+
+        _worker_fn(
+            trainer,
+            pl_module,
+            "num_val_batches",
+            dataloader_idx=dataloader_idx,
+            metric_name=self.config.metric_name,
+        )
+
+    @override
+    def on_test_batch_start(
+        self,
+        trainer: Trainer,
+        pl_module: LightningModule,
+        batch: Any,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        if trainer.logger is None or not self.config.test:
+            return
+
+        _worker_fn(
+            trainer,
+            pl_module,
+            "num_test_batches",
+            dataloader_idx=dataloader_idx,
+            metric_name=self.config.metric_name,
+        )

{nshtrainer-1.4.0 → nshtrainer-1.4.1}/src/nshtrainer/trainer/_config.py

@@ -419,7 +419,7 @@ class DirectoryConfig(C.Config):
 
 class TrainerConfig(C.Config):
     # region Active Run Configuration
-    id: Annotated[str, C.AllowMissing()] = C.MISSING
+    id: C.AllowMissing[str] = C.MISSING
     """ID of the run."""
     name: list[str] = []
     """Run name in parts. Full name is constructed by joining the parts with spaces."""

nshtrainer-1.4.0/src/nshtrainer/callbacks/log_epoch.py

@@ -1,49 +0,0 @@
-from __future__ import annotations
-
-import logging
-import math
-from typing import Any, Literal
-
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Callback
-from typing_extensions import final, override
-
-from .base import CallbackConfigBase, callback_registry
-
-log = logging.getLogger(__name__)
-
-
-@final
-@callback_registry.register
-class LogEpochCallbackConfig(CallbackConfigBase):
-    name: Literal["log_epoch"] = "log_epoch"
-
-    @override
-    def create_callbacks(self, trainer_config):
-        yield LogEpochCallback()
-
-
-class LogEpochCallback(Callback):
-    def __init__(self, metric_name: str = "computed_epoch"):
-        super().__init__()
-
-        self.metric_name = metric_name
-
-    @override
-    def on_train_batch_start(
-        self, trainer: Trainer, pl_module: LightningModule, batch: Any, batch_idx: int
-    ):
-        if trainer.logger is None:
-            return
-
-        # If trainer.num_training_batches is not set or is nan/inf, we cannot calculate the epoch
-        if (
-            not trainer.num_training_batches
-            or math.isnan(trainer.num_training_batches)
-            or math.isinf(trainer.num_training_batches)
-        ):
-            log.warning("Trainer has no valid num_training_batches. Cannot log epoch.")
-            return
-
-        epoch = pl_module.global_step / trainer.num_training_batches
-        pl_module.log(self.metric_name, epoch, on_step=True, on_epoch=False)