nshtrainer 1.5.0__tar.gz → 1.5.1__tar.gz

{nshtrainer-1.5.0 → nshtrainer-1.5.1}/PKG-INFO

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

{nshtrainer-1.5.0 → nshtrainer-1.5.1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "nshtrainer"
-version = "1.5.0"
+version = "1.5.1"
 description = ""
 authors = [{ name = "Nima Shoghi", email = "nimashoghi@gmail.com" }]
 requires-python = ">=3.10,<4.0"

{nshtrainer-1.5.0 → nshtrainer-1.5.1}/src/nshtrainer/callbacks/log_epoch.py

@@ -35,7 +35,7 @@ class LogEpochCallbackConfig(CallbackConfigBase):
         yield LogEpochCallback(self)
 
 
-def _worker_fn(
+def _log_on_step(
     trainer: Trainer,
     pl_module: LightningModule,
     num_batches_prop: str,
@@ -75,6 +75,19 @@ def _worker_fn(
     pl_module.log(metric_name, epoch, on_step=True, on_epoch=False)
 
 
+def _log_on_epoch(
+    trainer: Trainer,
+    pl_module: LightningModule,
+    *,
+    metric_name: str,
+):
+    if trainer.logger is None:
+        return
+
+    epoch = pl_module.current_epoch + 1
+    pl_module.log(metric_name, epoch, on_step=False, on_epoch=True)
+
+
 class LogEpochCallback(Callback):
     def __init__(self, config: LogEpochCallbackConfig):
         super().__init__()
@@ -85,16 +98,27 @@ class LogEpochCallback(Callback):
     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:
+        if not self.config.train:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_training_batches",
             metric_name=self.config.metric_name,
         )
 
+    @override
+    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
+        if not self.config.train:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )
+
     @override
     def on_validation_batch_start(
         self,
@@ -104,10 +128,10 @@ class LogEpochCallback(Callback):
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:
-        if trainer.logger is None or not self.config.val:
+        if not self.config.val:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_val_batches",
@@ -115,6 +139,19 @@ class LogEpochCallback(Callback):
             metric_name=self.config.metric_name,
         )
 
+    @override
+    def on_validation_epoch_end(
+        self, trainer: Trainer, pl_module: LightningModule
+    ) -> None:
+        if not self.config.val:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )
+
     @override
     def on_test_batch_start(
         self,
@@ -124,13 +161,24 @@ class LogEpochCallback(Callback):
         batch_idx: int,
         dataloader_idx: int = 0,
     ) -> None:
-        if trainer.logger is None or not self.config.test:
+        if not self.config.test:
             return
 
-        _worker_fn(
+        _log_on_step(
             trainer,
             pl_module,
             "num_test_batches",
             dataloader_idx=dataloader_idx,
             metric_name=self.config.metric_name,
         )
+
+    @override
+    def on_test_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
+        if not self.config.test:
+            return
+
+        _log_on_epoch(
+            trainer,
+            pl_module,
+            metric_name=self.config.metric_name,
+        )