nshtrainer 0.11.4__py3-none-any.whl → 0.11.6__py3-none-any.whl

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

{nshtrainer-0.11.4.dist-info → nshtrainer-0.11.6.dist-info}/METADATA

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

@@ -11,7 +11,7 @@ nshtrainer/callbacks/_throughput_monitor_callback.py,sha256=aJo_11rc4lo0IYOd-kHm
 nshtrainer/callbacks/actsave.py,sha256=qbnaKts4_dvjPeAaPtv7Ds12_vEWzaHUfg_--49NB9I,4041
 nshtrainer/callbacks/base.py,sha256=UnlYZAqSb8UwBJR-N5-XunxFx2yZjZ4lyGqUfhbCRlI,3555
 nshtrainer/callbacks/checkpoint/__init__.py,sha256=zrEVCGFikfkt0iOMceOFzXsZG2-6QrqY79RKBCS7bu4,738
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=SfM0VGCSDxYutJ21tY8m283eElN8jyKh2hhj6b7-1-s,6121
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=Ygblsf9NdLHxQPJUM47W0nGxlabj-ZnEBIMpvk-QMS8,7124
 nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py,sha256=NES-acaslPBiZQIMAdk_YwtnBrkm_y_BJQ8Ian0UKP0,4294
 nshtrainer/callbacks/checkpoint/model_checkpoint.py,sha256=mLFMbNzeMiBer3BCb7o3ucswKpOCQlYyN3wdB92N-LY,6884
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=s8tOHrnb_uVqLVeV2K38ZszXrXPTEGdDVfXuXgo_KDQ,3277
@@ -82,6 +82,6 @@ nshtrainer/util/seed.py,sha256=Or2wMPsnQxfnZ2xfBiyMcHFIUt3tGTNeMMyOEanCkqs,280
 nshtrainer/util/slurm.py,sha256=rofIU26z3SdL79SF45tNez6juou1cyDLz07oXEZb9Hg,1566
 nshtrainer/util/typed.py,sha256=NGuDkDzFlc1fAoaXjOFZVbmj0mRFjsQi1E_hPa7Bn5U,128
 nshtrainer/util/typing_utils.py,sha256=8ptjSSLZxlmy4FY6lzzkoGoF5fGNClo8-B_c0XHQaNU,385
-nshtrainer-0.11.4.dist-info/METADATA,sha256=6s-tc2dutPpIQaxv7j8cdZWiLU3l0Y_aS5f-k5JvAfM,860
-nshtrainer-0.11.4.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.4.dist-info/RECORD,,
+nshtrainer-0.11.6.dist-info/METADATA,sha256=tHGQ69o-paHEvlLUgo46bWeMlvuTlb8Q-upA00NxoKE,860
+nshtrainer-0.11.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.11.6.dist-info/RECORD,,