nshtrainer 0.11.4__tar.gz → 0.11.6__tar.gz

{nshtrainer-0.11.4 → nshtrainer-0.11.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.11.4
+Version: 0.11.6
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.11.4 → nshtrainer-0.11.6}/pyproject.toml

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

{nshtrainer-0.11.4 → nshtrainer-0.11.6}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -72,6 +72,8 @@ class BestCheckpoint(Checkpoint):
         self.metric = metric
         self.dirpath = dirpath
 
+        self._last_global_step_saved = 0  # no need to save when no steps were taken
+
     @override
     def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):
         self._save_best_checkpoint(trainer)
@@ -127,6 +129,11 @@ class BestCheckpoint(Checkpoint):
         log.debug(f"Created best symlink: {symlink_path}")
 
     def _save_best_checkpoint(self, trainer: Trainer):
+        # Skip saving the checkpoint if we're not in the fitting state
+        if self._should_skip_saving_checkpoint(trainer):
+            return
+
+        # Get the current metric value
         if (current := self._get_metric_value(trainer.callback_metrics)) is None:
             log.warning(
                 f"Can't save best model, {self.metric.validation_monitor} not found in metrics"
@@ -150,6 +157,7 @@ class BestCheckpoint(Checkpoint):
         # Save the current model
         filepath = self._ckpt_path(trainer)
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
+        log.debug(f"Saved best checkpoint: {filepath}")
 
         # Remove worst checkpoint if we've reached save_top_k
         # NOTE: We add 1 to save_top_k here because we have just saved a new checkpoint
@@ -163,9 +171,26 @@ class BestCheckpoint(Checkpoint):
             )
 
         # Create symlink to best model
-        _, best_ckpt_path = sorted_ckpts[0]
-        self._create_symlink(trainer, best_ckpt_path)
-        log.debug(f"Saved best checkpoint: {filepath}")
+        if sorted_ckpts:
+            _, best_ckpt_path = sorted_ckpts[0]
+            self._create_symlink(trainer, best_ckpt_path)
+
+        # Update the last global step saved
+        self._last_global_step_saved = trainer.global_step
 
         # Barrier to ensure all processes have saved the checkpoint before continuing
         trainer.strategy.barrier()
+
+    def _should_skip_saving_checkpoint(self, trainer: Trainer) -> bool:
+        from lightning.pytorch.trainer.states import TrainerFn
+
+        return (
+            bool(
+                getattr(trainer, "fast_dev_run", False)
+            )  # disable checkpointing with fast_dev_run
+            or trainer.state.fn
+            != TrainerFn.FITTING  # don't save anything during non-fit
+            or trainer.sanity_checking  # don't save anything during sanity check
+            or self._last_global_step_saved
+            == trainer.global_step  # already saved at the last step
+        )