nshtrainer 0.19.3__tar.gz → 0.20.0__tar.gz

{nshtrainer-0.19.3 → nshtrainer-0.20.0}/PKG-INFO

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

{nshtrainer-0.19.3 → nshtrainer-0.20.0}/pyproject.toml

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

{nshtrainer-0.19.3 → nshtrainer-0.20.0}/src/nshtrainer/_checkpoint/loader.py

@@ -76,7 +76,11 @@ class CheckpointLoadingConfig(C.Config):
     """Whether to include checkpoints from HPC pre-emption."""
 
     @classmethod
-    def _auto_train(cls, ckpt: Literal["best", "last"] | str | Path | None):
+    def none(cls, include_hpc: bool = False):
+        return cls(strategies=[], include_hpc=include_hpc)
+
+    @classmethod
+    def _auto_train(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
         if ckpt is None:
             ckpt = "last"
         match ckpt:
@@ -90,6 +94,8 @@ class CheckpointLoadingConfig(C.Config):
                     strategies=[LastCheckpointStrategyConfig()],
                     include_hpc=True,
                 )
+            case "none":
+                return cls.none()
             case Path() | str():
                 ckpt = Path(ckpt)
                 return cls(
@@ -103,7 +109,7 @@ class CheckpointLoadingConfig(C.Config):
                 assert_never(ckpt)
 
     @classmethod
-    def _auto_eval(cls, ckpt: Literal["best", "last"] | str | Path | None):
+    def _auto_eval(cls, ckpt: Literal["best", "last", "none"] | str | Path | None):
         if ckpt is None:
             log.warn("No checkpoint specified for evaluation. Defaulting to `last`.")
             ckpt = "last"
@@ -119,6 +125,8 @@ class CheckpointLoadingConfig(C.Config):
                     strategies=[LastCheckpointStrategyConfig()],
                     include_hpc=False,
                 )
+            case "none":
+                return cls.none(include_hpc=False)
             case Path() | str():
                 ckpt = Path(ckpt)
                 return cls(
@@ -131,7 +139,7 @@ class CheckpointLoadingConfig(C.Config):
     @classmethod
     def auto(
         cls,
-        ckpt: Literal["best", "last"] | str | Path | None,
+        ckpt: Literal["best", "last", "none"] | str | Path | None,
         trainer_mode: TrainerFn,
     ):
         """
@@ -142,7 +150,7 @@ class CheckpointLoadingConfig(C.Config):
 
         Parameters:
         -----------
-        ckpt : Literal["best", "last"] | str | Path | None
+        ckpt : Literal["best", "last", "none"] | str | Path | None
             Specifies the checkpoint loading preference:
             - "best": Use the best checkpoint based on the primary metric.
             - "last": Use the most recent checkpoint.

{nshtrainer-0.19.3 → nshtrainer-0.20.0}/src/nshtrainer/model/config.py

@@ -811,11 +811,14 @@ class SanityCheckingConfig(C.Config):
 
 
 class TrainerConfig(C.Config):
-    ckpt_path: str | Path | None = None
-    """Path to a checkpoint to load and resume training from."""
+    ckpt_path: Literal["none"] | str | Path | None = None
+    """Path to a checkpoint to load and resume training from. If ``"none"``, will not load a checkpoint."""
 
-    checkpoint_loading: CheckpointLoadingConfig | Literal["auto"] = "auto"
-    """Checkpoint loading configuration options."""
+    checkpoint_loading: CheckpointLoadingConfig | Literal["auto", "none"] = "auto"
+    """Checkpoint loading configuration options.
+    `"auto"` will automatically determine the best checkpoint loading strategy based on the provided.
+    `"none"` will disable checkpoint loading.
+    """
 
     checkpoint_saving: CheckpointSavingConfig = CheckpointSavingConfig()
     """Checkpoint saving configuration options."""

{nshtrainer-0.19.3 → nshtrainer-0.20.0}/src/nshtrainer/trainer/checkpoint_connector.py

@@ -31,8 +31,14 @@ class _CheckpointConnector(_LightningCheckpointConnector):
 
         # Now, resolve the checkpoint loader config.
         root_config = cast("BaseConfig", trainer._base_module.config)
-        if (ckpt_loader_config := root_config.trainer.checkpoint_loading) == "auto":
-            ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
+        ckpt_loader_config = root_config.trainer.checkpoint_loading
+        match ckpt_loader_config:
+            case "auto":
+                ckpt_loader_config = CheckpointLoadingConfig.auto(ckpt_path, state_fn)
+            case "none":
+                ckpt_loader_config = CheckpointLoadingConfig.none()
+            case _:
+                pass
         log.debug(f"Checkpoint loader config: {ckpt_loader_config}")
 
         # Use the config to resolve the checkpoint.