nshtrainer 1.0.0b47__tar.gz → 1.0.0b48__tar.gz

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b47
+Version: 1.0.0b48
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/pyproject.toml

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

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -51,7 +51,7 @@ class BestCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
 class BestCheckpointCallback(CheckpointBase[BestCheckpointCallbackConfig]):
     @property
     def _metric_name_normalized(self):
-        return self.metric.name.replace("/", "_").replace(" ", "_").replace(".", "_")
+        return self.metric.monitor.replace("/", "_").replace(" ", "_").replace(".", "_")
 
     @override
     def __init__(
@@ -69,12 +69,12 @@ class BestCheckpointCallback(CheckpointBase[BestCheckpointCallbackConfig]):
 
     @override
     def default_filename(self):
-        return f"epoch{{epoch}}-step{{step}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
+        return f"epoch{{epoch}}-step{{step}}-{self._metric_name_normalized}{{{self.metric.monitor}}}"
 
     @override
     def topk_sort_key(self, metadata: CheckpointMetadata):
         return metadata.metrics.get(
-            self.metric.validation_monitor,
+            self.metric.monitor,
             float("-inf" if self.metric.mode == "max" else "inf"),
         )
 

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/callbacks/early_stopping.py

@@ -68,7 +68,7 @@ class EarlyStoppingCallback(_EarlyStopping):
         del config, metric
 
         super().__init__(
-            monitor=self.metric.validation_monitor,
+            monitor=self.metric.monitor,
             mode=self.metric.mode,
             patience=self.config.patience,
             min_delta=self.config.min_delta,

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/callbacks/metric_validation.py

@@ -55,14 +55,14 @@ class MetricValidationCallback(Callback):
         self.metrics = metrics
 
     def _check_metrics(self, trainer: Trainer):
-        metric_names = ", ".join(metric.validation_monitor for metric in self.metrics)
+        metric_names = ", ".join(metric.monitor for metric in self.metrics)
         log.info(f"Validating metrics: {metric_names}...")
         logged_metrics = set(trainer.logged_metrics.keys())
 
         invalid_metrics: list[str] = []
         for metric in self.metrics:
-            if metric.validation_monitor not in logged_metrics:
-                invalid_metrics.append(metric.validation_monitor)
+            if metric.monitor not in logged_metrics:
+                invalid_metrics.append(metric.monitor)
 
         if invalid_metrics:
             msg = (

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/callbacks/rlp_sanity_checks.py

@@ -171,7 +171,7 @@ class CustomRLPImplementation(Protocol):
     __reduce_lr_on_plateau__: bool
 
 
-class _RLPSanityCheckModuleMixin(LightningModule):
+class RLPSanityCheckModuleMixin(LightningModule):
     def reduce_lr_on_plateau_config(
         self,
         lr_scheduler: LRSchedulerTypeUnion | LRSchedulerConfigType,

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/data/datamodule.py

@@ -11,13 +11,13 @@ from lightning.pytorch import LightningDataModule
 from typing_extensions import Never, TypeVar, deprecated, override
 
 from ..model.mixins.callback import CallbackRegistrarModuleMixin
-from ..model.mixins.debug import _DebugModuleMixin
+from ..model.mixins.debug import DebugModuleMixin
 
 THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
 
 
 class LightningDataModuleBase(
-    _DebugModuleMixin,
+    DebugModuleMixin,
     CallbackRegistrarModuleMixin,
     LightningDataModule,
     ABC,

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/loggers/__init__.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 from typing import Annotated
 
-import nshconfig as C
 from typing_extensions import TypeAliasType
 
 from .actsave import ActSaveLoggerConfig as ActSaveLoggerConfig

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/lr_scheduler/reduce_lr_on_plateau.py

@@ -49,13 +49,13 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
         if (metric := self.metric) is None:
             from ..trainer import Trainer
 
-            assert isinstance(
-                trainer := lightning_module.trainer, Trainer
-            ), "The trainer must be a `nshtrainer.Trainer` instance."
+            assert isinstance(trainer := lightning_module.trainer, Trainer), (
+                "The trainer must be a `nshtrainer.Trainer` instance."
+            )
 
-            assert (
-                metric := trainer.hparams.primary_metric
-            ) is not None, "Primary metric must be provided if metric is not specified."
+            assert (metric := trainer.hparams.primary_metric) is not None, (
+                "Primary metric must be provided if metric is not specified."
+            )
 
         lr_scheduler = ReduceLROnPlateau(
             optimizer,
@@ -70,7 +70,7 @@ class ReduceLROnPlateauConfig(LRSchedulerConfigBase):
         )
         return {
             "scheduler": lr_scheduler,
-            "monitor": metric.validation_monitor,
+            "monitor": metric.monitor,
         }
 
     @override

nshtrainer-1.0.0b48/src/nshtrainer/metrics/_config.py

@@ -0,0 +1,25 @@
+from __future__ import annotations
+
+import builtins
+from typing import Any, Literal
+
+import nshconfig as C
+
+
+class MetricConfig(C.Config):
+    monitor: str
+    """The name of the metric to monitor."""
+
+    mode: Literal["min", "max"]
+    """
+    The mode of the primary metric:
+    - "min" for metrics that should be minimized (e.g., loss)
+    - "max" for metrics that should be maximized (e.g., accuracy)
+    """
+
+    @property
+    def best(self):
+        return builtins.min if self.mode == "min" else builtins.max
+
+    def is_better(self, a: Any, b: Any):
+        return self.best(a, b) == a

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/model/base.py

@@ -15,9 +15,9 @@ from lightning.pytorch.utilities.model_helpers import is_overridden
 from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 from typing_extensions import Never, TypeVar, deprecated, override
 
-from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
+from ..callbacks.rlp_sanity_checks import RLPSanityCheckModuleMixin
 from .mixins.callback import CallbackModuleMixin
-from .mixins.debug import _DebugModuleMixin
+from .mixins.debug import DebugModuleMixin
 from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
@@ -54,8 +54,8 @@ VALID_REDUCE_OPS = (
 
 
 class LightningModuleBase(
-    _DebugModuleMixin,
-    _RLPSanityCheckModuleMixin,
+    DebugModuleMixin,
+    RLPSanityCheckModuleMixin,
     LoggerLightningModuleMixin,
     CallbackModuleMixin,
     LightningModule,

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/model/mixins/debug.py

@@ -28,7 +28,7 @@ def _trainer(module: Any):
     return trainer
 
 
-class _DebugModuleMixin:
+class DebugModuleMixin:
     @property
     def nshtrainer_or_none(self):
         return _trainer(self)

{nshtrainer-1.0.0b47 → nshtrainer-1.0.0b48}/src/nshtrainer/model/mixins/logger.py

@@ -54,6 +54,12 @@ class _LogContextKwargs:
         d = dataclasses.asdict(self)
         for field in self.__ignore_fields__:
             d.pop(field, None)
+
+        # Pop all None values
+        for k in list(d.keys()):
+            if d[k] is None:
+                d.pop(k)
+
         return d
 
 
@@ -134,18 +140,18 @@ class LoggerLightningModuleMixin(mixin_base_type(LightningModule)):
         self,
         name: str,
         value: _METRIC,
-        prog_bar: bool = False,
+        prog_bar: bool | None = None,
         logger: bool | None = None,
         on_step: bool | None = None,
         on_epoch: bool | None = None,
-        reduce_fx: str | Callable = "mean",
-        enable_graph: bool = False,
-        sync_dist: bool = False,
+        reduce_fx: str | Callable | None = None,
+        enable_graph: bool | None = None,
+        sync_dist: bool | None = None,
         sync_dist_group: Any | None = None,
-        add_dataloader_idx: bool = True,
+        add_dataloader_idx: bool | None = None,
         batch_size: int | None = None,
         metric_attribute: str | None = None,
-        rank_zero_only: bool = False,
+        rank_zero_only: bool | None = None,
     ) -> None:
         # If logging is disabled, then do nothing.
         if not self.logging_enabled:

nshtrainer-1.0.0b47/src/nshtrainer/metrics/_config.py

@@ -1,42 +0,0 @@
-from __future__ import annotations
-
-import builtins
-from typing import Any, Literal
-
-import nshconfig as C
-
-
-class MetricConfig(C.Config):
-    name: str
-    """The name of the primary metric."""
-
-    mode: Literal["min", "max"]
-    """
-    The mode of the primary metric:
-    - "min" for metrics that should be minimized (e.g., loss)
-    - "max" for metrics that should be maximized (e.g., accuracy)
-    """
-
-    @property
-    def validation_monitor(self) -> str:
-        return f"val/{self.name}"
-
-    def __post_init__(self):
-        for split in ("train", "val", "test", "predict"):
-            if self.name.startswith(f"{split}/"):
-                raise ValueError(
-                    f"Primary metric name should not start with '{split}/'. "
-                    f"Just use '{self.name[len(split) + 1:]}' instead. "
-                    "The split name is automatically added depending on the context."
-                )
-
-    @classmethod
-    def loss(cls, mode: Literal["min", "max"] = "min"):
-        return cls(name="loss", mode=mode)
-
-    @property
-    def best(self):
-        return builtins.min if self.mode == "min" else builtins.max
-
-    def is_better(self, a: Any, b: Any):
-        return self.best(a, b) == a