nshtrainer 1.5.0__py3-none-any.whl → 1.5.1__py3-none-any.whl

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,
+        )

{nshtrainer-1.5.0.dist-info → nshtrainer-1.5.1.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-1.5.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=C2yUww8lAuCX-dy06tsw95yCBOfFd2mfGs0VhrEq1oU,3775
+nshtrainer/callbacks/log_epoch.py,sha256=-uC5ss9p_ngXUCrSIUwViFcaaVX6ALUzIAKxoDgZrac,4823
 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
@@ -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.5.0.dist-info/METADATA,sha256=fbtia7kDnNxHx_8VE0I-zFtmlF-HMAxH5raSiPjtl7w,980
-nshtrainer-1.5.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-nshtrainer-1.5.0.dist-info/RECORD,,
+nshtrainer-1.5.1.dist-info/METADATA,sha256=ct7S8c2O-oHJ2yw3-pApipwmP_r07z8lmFu40FhQY-k,980
+nshtrainer-1.5.1.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+nshtrainer-1.5.1.dist-info/RECORD,,