nshtrainer 0.11.5__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"
@@ -168,5 +175,22 @@ class BestCheckpoint(Checkpoint):
             _, 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.5.dist-info → nshtrainer-0.11.6.dist-info}/METADATA

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

{nshtrainer-0.11.5.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=w99O5GWRcV89XBe4j__v2TvNEHys0x_r3tSTr-6Lhec,6154
+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.5.dist-info/METADATA,sha256=KHgvYOhQXbc37awWeLbpbdVQbSEU4J7KoC7Lr5286KE,860
-nshtrainer-0.11.5.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.5.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,,