nshtrainer 0.11.12__py3-none-any.whl → 0.12.0__py3-none-any.whl

nshtrainer/callbacks/checkpoint/_base.py

@@ -79,6 +79,9 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
     @abstractmethod
     def topk_sort_key(self, metadata: CheckpointMetadata) -> Any: ...
 
+    @abstractmethod
+    def topk_sort_reverse(self) -> bool: ...
+
     def symlink_path(self):
         if not self.config.save_symlink:
             return None
@@ -102,7 +105,7 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         ]
 
         # Sort by the topk sort key
-        metas = sorted(metas, key=self.topk_sort_key)
+        metas = sorted(metas, key=self.topk_sort_key, reverse=self.topk_sort_reverse())
 
         # Now, the metas are sorted from the best to the worst,
         # so we can remove the worst checkpoints
@@ -145,15 +148,15 @@ class CheckpointBase(Checkpoint, ABC, Generic[TConfig]):
         trainer.save_checkpoint(filepath, self.config.save_weights_only)
 
         if trainer.is_global_zero:
-            # Remove old checkpoints
-            self.remove_old_checkpoints(trainer)
-
             # Create the latest symlink
             if (symlink_filename := self.symlink_path()) is not None:
                 symlink_path = self.dirpath / symlink_filename
                 _link_checkpoint(filepath, symlink_path, metadata=True)
                 log.debug(f"Created latest symlink: {symlink_path}")
 
+            # Remove old checkpoints
+            self.remove_old_checkpoints(trainer)
+
         # Barrier to ensure all processes have saved the checkpoint,
         # deleted the old checkpoints, and created the symlink before continuing
         trainer.strategy.barrier()

nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -55,7 +55,7 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
 
     @override
     def default_filename(self):
-        return f"epoch{{epoch:03d}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
+        return f"epoch{{epoch}}-step{{step}}-{self._metric_name_normalized}{{{self.metric.validation_monitor}}}"
 
     @override
     def topk_sort_key(self, metadata: CheckpointMetadata):
@@ -64,6 +64,10 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
             float("-inf" if self.metric.mode == "max" else "inf"),
         )
 
+    @override
+    def topk_sort_reverse(self):
+        return self.metric.mode == "max"
+
     # Events
     @override
     def on_validation_end(self, trainer: Trainer, pl_module: LightningModule):

nshtrainer/callbacks/checkpoint/last_checkpoint.py

@@ -28,12 +28,16 @@ class LastCheckpoint(CheckpointBase[LastCheckpointCallbackConfig]):
 
     @override
     def default_filename(self):
-        return "epoch{epoch:03d}-step{step:07d}"
+        return "epoch{epoch}-step{step}"
 
     @override
     def topk_sort_key(self, metadata: CheckpointMetadata):
         return metadata.checkpoint_timestamp
 
+    @override
+    def topk_sort_reverse(self):
+        return True
+
     @override
     def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
         self.save_checkpoints(trainer)

{nshtrainer-0.11.12.dist-info → nshtrainer-0.12.0.dist-info}/METADATA

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

{nshtrainer-0.11.12.dist-info → nshtrainer-0.12.0.dist-info}/RECORD

@@ -11,9 +11,9 @@ 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=g-3zIthupERKqWZQw-A_busQPaPRkto6iHBV-M7nK1Y,527
-nshtrainer/callbacks/checkpoint/_base.py,sha256=wb5ARqqSslZXB5FSwNd9a_7qtt4-F9dC7E1vNHKCK1o,5994
-nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=ySeyALxc-YJaS5IhmW0MkAhr41Mxsm_BllHuufXGy4Y,2067
-nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CM8f37dwaYHkjQFfJNTZTzSoF45zEjFRm-Fg1CzYmP4,1037
+nshtrainer/callbacks/checkpoint/_base.py,sha256=yBUZ63P4a3m5jviIUGxBd2WXq_TigpKt3w4KdzqrzLs,6094
+nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=DJiLo7NDzd-lp-O3v7Cv8WejyjXPV_6_RmfltKO9fvE,2165
+nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=CqB_8Xv32rtpLCaEEPi6DbRZm4ph5TWS-LfqIHXUIUA,1097
 nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=ctT88EGT22_t_6tr5r7Sfo43cuve6XeroBnBYRMPOus,3372
 nshtrainer/callbacks/early_stopping.py,sha256=LGn3rdbvkFfUo9kwMzK4eMGlPAqD9uFdowDx6VdfozQ,3761
 nshtrainer/callbacks/ema.py,sha256=8-WHmKFP3VfnzMviJaIFmVD9xHPqIPmq9NRF5xdu3c8,12131
@@ -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.12.dist-info/METADATA,sha256=Lthmwj2EpbkCNWlDg2n7z4zS197Pn5Kt1aUMayB70oo,861
-nshtrainer-0.11.12.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-nshtrainer-0.11.12.dist-info/RECORD,,
+nshtrainer-0.12.0.dist-info/METADATA,sha256=YKxr6_dhF7jHF1njUUJi0CZRR0BuOPDNrQiwiIO7iTU,860
+nshtrainer-0.12.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+nshtrainer-0.12.0.dist-info/RECORD,,