nshtrainer 1.4.0__py3-none-any.whl → 1.4.1__py3-none-any.whl

nshtrainer/callbacks/log_epoch.py

@@ -18,32 +18,119 @@ log = logging.getLogger(__name__)
 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()
+        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, metric_name: str = "computed_epoch"):
+    def __init__(self, config: LogEpochCallbackConfig):
         super().__init__()
 
-        self.metric_name = metric_name
+        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:
+        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
 
-        # 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.")
+        _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
 
-        epoch = pl_module.global_step / trainer.num_training_batches
-        pl_module.log(self.metric_name, epoch, on_step=True, on_epoch=False)
+        _worker_fn(
+            trainer,
+            pl_module,
+            "num_test_batches",
+            dataloader_idx=dataloader_idx,
+            metric_name=self.config.metric_name,
+        )

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.dist-info → nshtrainer-1.4.1.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-1.4.1.dist-info}/RECORD

@@ -21,7 +21,7 @@ nshtrainer/callbacks/ema.py,sha256=dBFiUXG0xmyCw8-ayuSzJMKqSbepl6Ii5VIbhFlT5ug,1
 nshtrainer/callbacks/finite_checks.py,sha256=3lZ3kEIjmYQfqTF0DcrgZ9_98ZLQhQj8usH7SgWst3o,2185
 nshtrainer/callbacks/gradient_skipping.py,sha256=8g7oC7PF0LTAEzwiNoaS5tWOnkjk_EB0QG3JdHkQ8ek,3523
 nshtrainer/callbacks/interval.py,sha256=UCzUzt3XCFVyQyCWL9lOrStkkxesvduNOYk8yMrGTTk,8116
-nshtrainer/callbacks/log_epoch.py,sha256=B5Dm8XVZwCzKUhUWfT_5PDdDac993191OsbcxxuSVJE,1457
+nshtrainer/callbacks/log_epoch.py,sha256=C2yUww8lAuCX-dy06tsw95yCBOfFd2mfGs0VhrEq1oU,3775
 nshtrainer/callbacks/lr_monitor.py,sha256=v45ehnwNO987087HfiOY5aIrVRbwdKMgPYRFHs1fyEE,1444
 nshtrainer/callbacks/metric_validation.py,sha256=4RDr1FuNKfro-6QEtmcFqT4iNf2twmJVNk9y-8nq9bg,2882
 nshtrainer/callbacks/norm_logging.py,sha256=nVIDWe-ASl5zN830-ODR8QMCqI1ma-QPCIwoy0Wb-Nk,6390
@@ -133,7 +133,7 @@ nshtrainer/profiler/advanced.py,sha256=XrM3FX0ThCv5UwUrrH0l4Ow4LGAtpiBww2N8QAU5N
 nshtrainer/profiler/pytorch.py,sha256=8K37XvPnCApUpIK8tA2zNMFIaIiTLSoxKQoiyCPBm1Q,2757
 nshtrainer/profiler/simple.py,sha256=PimjqcU-JuS-8C0ZGHAdwCxgNLij4x0FH6WXsjBQzZs,1005
 nshtrainer/trainer/__init__.py,sha256=jRaHdaFK8wxNrN1bleT9cf29iZahL_-XkWo5TWz2CmA,550
-nshtrainer/trainer/_config.py,sha256=FWEspBYt_bjLhUSkJApkC9pfYBTlFBHmIQRFNGpGjAc,45849
+nshtrainer/trainer/_config.py,sha256=GL8DtuH-6x2aHcRlEcmzyhEBMRRldiSazNAeNmPw7gM,45836
 nshtrainer/trainer/_distributed_prediction_result.py,sha256=bQw8Z6PT694UUf-zQPkech6CxyUSy8bAIexfSfPej0U,2507
 nshtrainer/trainer/_log_hparams.py,sha256=XH2lZ4U_3AZBhOt91ocsEhdL_NRz35oWvqLCUFDohUs,2389
 nshtrainer/trainer/_runtime_callback.py,sha256=6F2Gq27Q8OFfN3RtdNC6QRA8ac0LC1hh4DUE3V5WgbI,4217
@@ -159,6 +159,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.4.0.dist-info/METADATA,sha256=PIV_5Swp1HhgFU2ZBj_X1tCeOBfNhrhTXOFB1vgunno,979
-nshtrainer-1.4.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.4.0.dist-info/RECORD,,
+nshtrainer-1.4.1.dist-info/METADATA,sha256=QL69Trcmw3NF3UOovpqVJbzBTtHJtnDDxAzxyj9EX24,980
+nshtrainer-1.4.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.4.1.dist-info/RECORD,,