nshtrainer 0.11.8__tar.gz → 0.11.10__tar.gz

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/PKG-INFO

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

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/pyproject.toml

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

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/src/nshtrainer/callbacks/__init__.py

@@ -8,12 +8,6 @@ from .checkpoint import BestCheckpoint as BestCheckpoint
 from .checkpoint import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .checkpoint import LastCheckpoint as LastCheckpoint
 from .checkpoint import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
-from .checkpoint import LatestEpochCheckpoint as LatestEpochCheckpoint
-from .checkpoint import (
-    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
-)
-from .checkpoint import ModelCheckpoint as ModelCheckpoint
-from .checkpoint import ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig
 from .checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
@@ -49,8 +43,6 @@ CallbackConfig = Annotated[
     | EMAConfig
     | BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
-    | ModelCheckpointCallbackConfig
-    | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig
     | WandbWatchConfig,
     C.Field(discriminator="name"),

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/src/nshtrainer/callbacks/checkpoint/__init__.py

@@ -6,14 +6,6 @@ from .last_checkpoint import LastCheckpoint as LastCheckpoint
 from .last_checkpoint import (
     LastCheckpointCallbackConfig as LastCheckpointCallbackConfig,
 )
-from .latest_epoch_checkpoint import LatestEpochCheckpoint as LatestEpochCheckpoint
-from .latest_epoch_checkpoint import (
-    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
-)
-from .model_checkpoint import ModelCheckpoint as ModelCheckpoint
-from .model_checkpoint import (
-    ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig,
-)
 from .on_exception_checkpoint import OnExceptionCheckpoint as OnExceptionCheckpoint
 from .on_exception_checkpoint import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/src/nshtrainer/model/__init__.py

@@ -5,17 +5,15 @@ from .base import LightningModuleBase as LightningModuleBase
 from .config import BaseConfig as BaseConfig
 from .config import BaseLoggerConfig as BaseLoggerConfig
 from .config import BaseProfilerConfig as BaseProfilerConfig
+from .config import BestCheckpointCallbackConfig as BestCheckpointCallbackConfig
 from .config import CheckpointLoadingConfig as CheckpointLoadingConfig
 from .config import CheckpointSavingConfig as CheckpointSavingConfig
 from .config import DirectoryConfig as DirectoryConfig
 from .config import EarlyStoppingConfig as EarlyStoppingConfig
 from .config import GradientClippingConfig as GradientClippingConfig
-from .config import (
-    LatestEpochCheckpointCallbackConfig as LatestEpochCheckpointCallbackConfig,
-)
+from .config import LastCheckpointCallbackConfig as LastCheckpointCallbackConfig
 from .config import LoggingConfig as LoggingConfig
 from .config import MetricConfig as MetricConfig
-from .config import ModelCheckpointCallbackConfig as ModelCheckpointCallbackConfig
 from .config import (
     OnExceptionCheckpointCallbackConfig as OnExceptionCheckpointCallbackConfig,
 )

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/src/nshtrainer/model/config.py

@@ -40,8 +40,6 @@ from ..callbacks import (
     BestCheckpointCallbackConfig,
     CallbackConfig,
     LastCheckpointCallbackConfig,
-    LatestEpochCheckpointCallbackConfig,
-    ModelCheckpointCallbackConfig,
     OnExceptionCheckpointCallbackConfig,
     WandbWatchConfig,
 )
@@ -772,10 +770,8 @@ class ReproducibilityConfig(C.Config):
 
 
 CheckpointCallbackConfig: TypeAlias = Annotated[
-    ModelCheckpointCallbackConfig
-    | BestCheckpointCallbackConfig
+    BestCheckpointCallbackConfig
     | LastCheckpointCallbackConfig
-    | LatestEpochCheckpointCallbackConfig
     | OnExceptionCheckpointCallbackConfig,
     C.Field(discriminator="name"),
 ]
@@ -786,7 +782,6 @@ class CheckpointSavingConfig(CallbackConfigBase):
     """Enable checkpoint saving."""
 
     checkpoint_callbacks: Sequence[CheckpointCallbackConfig] = [
-        # ModelCheckpointCallbackConfig(),
         BestCheckpointCallbackConfig(),
         LastCheckpointCallbackConfig(),
         OnExceptionCheckpointCallbackConfig(),
@@ -806,36 +801,6 @@ class CheckpointSavingConfig(CallbackConfigBase):
 
         return True
 
-    @property
-    def model_checkpoint(self) -> ModelCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, ModelCheckpointCallbackConfig)
-            ),
-        )
-
-    @property
-    def latest_epoch_checkpoint(self) -> LatestEpochCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, LatestEpochCheckpointCallbackConfig)
-            ),
-        )
-
-    @property
-    def on_exception_checkpoint(self) -> OnExceptionCheckpointCallbackConfig | None:
-        return next(
-            (
-                callback
-                for callback in self.checkpoint_callbacks
-                if isinstance(callback, OnExceptionCheckpointCallbackConfig)
-            ),
-        )
-
     @override
     def create_callbacks(self, root_config: "BaseConfig"):
         if not self.should_save_checkpoints(root_config):

nshtrainer-0.11.8/src/nshtrainer/callbacks/checkpoint/latest_epoch_checkpoint.py

@@ -1,131 +0,0 @@
-import logging
-from pathlib import Path
-from typing import Literal
-
-from lightning.pytorch import LightningModule, Trainer
-from lightning.pytorch.callbacks import Checkpoint
-from typing_extensions import override
-
-from ..._checkpoint.metadata import _sort_ckpts_by_metadata
-from ..._checkpoint.saver import _link_checkpoint, _remove_checkpoint
-from ..base import CallbackConfigBase
-
-log = logging.getLogger(__name__)
-
-
-class LatestEpochCheckpointCallbackConfig(CallbackConfigBase):
-    name: Literal["latest_epoch_checkpoint"] = "latest_epoch_checkpoint"
-
-    dirpath: str | Path | None = None
-    """Directory path to save the checkpoint file."""
-
-    filename: str = "epoch{epoch:03d}_step{step:07d}"
-    """Checkpoint filename. This must not include the extension."""
-
-    save_weights_only: bool = False
-    """Whether to save only the model's weights or the entire model object."""
-
-    latest_symlink_filename: str | None = "latest"
-    """Filename for the latest symlink. If None, no symlink will be created."""
-
-    latest_k: int | Literal["all"] = 1
-    """Number of latest checkpoints to keep. If "all", all checkpoints are kept."""
-
-    @override
-    def create_callbacks(self, root_config):
-        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
-        )
-        dirpath = Path(dirpath)
-
-        yield LatestEpochCheckpoint(self, dirpath)
-
-
-class LatestEpochCheckpoint(Checkpoint):
-    PREFIX = "latest_"
-    EXTENSION = ".ckpt"
-
-    def __init__(self, config: LatestEpochCheckpointCallbackConfig, dirpath: Path):
-        super().__init__()
-
-        self.config = config
-        self.dirpath = dirpath
-
-        self._last_global_step_saved = 0
-
-    @override
-    def on_train_epoch_end(self, trainer: Trainer, pl_module: LightningModule):
-        self._save_new_checkpoint(trainer)
-
-    def _latest_symlink_filename(self):
-        if (filename := self.config.latest_symlink_filename) is None:
-            return None
-        return f"{filename}{self.EXTENSION}"
-
-    def _ckpt_path(self, trainer: Trainer):
-        filename = self.config.filename.format(
-            epoch=trainer.current_epoch, step=trainer.global_step
-        )
-        filename = f"{self.PREFIX}{filename}{self.EXTENSION}"
-        return self.dirpath / filename
-
-    def _remove_old_checkpoints(self, trainer: Trainer):
-        if (latest_k := self.config.latest_k) == "all":
-            return
-
-        # Get all configs, ignoring the latest symlink
-        ckpts = list(self.dirpath.glob(f"{self.PREFIX}*{self.EXTENSION}"))
-        # Ignore the latest symlink
-        if (latest_symlink_filename := self._latest_symlink_filename()) is not None:
-            ckpts = [p for p in ckpts if p.name != latest_symlink_filename]
-
-        # Sort by epoch, then step, then last modified
-        ckpts = _sort_ckpts_by_metadata(
-            ckpts,
-            key=lambda meta, p: (meta.epoch, meta.global_step, p.stat().st_mtime),
-            reverse=True,
-        )
-
-        # Remove all but the latest k checkpoints
-        # NOTE: We add 1 to the latest_k here because
-        # we're about to save a new checkpoint.
-        for _, ckpt_path in ckpts[latest_k:]:
-            _remove_checkpoint(trainer, ckpt_path, metadata=True)
-
-    def _save_new_checkpoint(self, trainer: Trainer):
-        if self._should_skip_saving_checkpoint(trainer):
-            return
-
-        # Save the new checkpoint
-        filepath = self._ckpt_path(trainer)
-        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._latest_symlink_filename()) is not None:
-                symlink_path = self.dirpath / symlink_filename
-                _link_checkpoint(filepath, symlink_path, metadata=True)
-                log.debug(f"Created latest symlink: {symlink_path}")
-
-        # Set 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.8/src/nshtrainer/callbacks/checkpoint/model_checkpoint.py

@@ -1,207 +0,0 @@
-import logging
-import re
-from datetime import timedelta
-from pathlib import Path
-from typing import TYPE_CHECKING, Literal
-
-from lightning.pytorch import Trainer
-from lightning.pytorch.callbacks.model_checkpoint import (
-    ModelCheckpoint as _ModelCheckpoint,
-)
-from typing_extensions import override
-
-from ..._checkpoint.saver import _link_checkpoint
-from ..._checkpoint.saver import _remove_checkpoint as _ckpt_saver_remove_checkpoint
-from ...metrics import MetricConfig
-from ..base import CallbackConfigBase
-
-if TYPE_CHECKING:
-    from ...model.config import BaseConfig
-
-log = logging.getLogger(__name__)
-
-
-def _convert_string(input_string: str):
-    # Find all variables enclosed in curly braces
-    variables = re.findall(r"\{(.*?)\}", input_string)
-
-    # Replace each variable with its corresponding key-value pair
-    output_string = input_string
-    for variable in variables:
-        # If the name is something like {variable:format}, we shouldn't process the format.
-        key_name = variable
-        if ":" in variable:
-            key_name, _ = variable.split(":", 1)
-            continue
-
-        # Replace '/' with '_' in the key name
-        key_name = key_name.replace("/", "_")
-        output_string = output_string.replace(
-            f"{{{variable}}}", f"{key_name}={{{variable}}}"
-        )
-
-    return output_string
-
-
-class ModelCheckpointCallbackConfig(CallbackConfigBase):
-    """Arguments for the ModelCheckpoint callback."""
-
-    name: Literal["model_checkpoint"] = "model_checkpoint"
-
-    dirpath: str | Path | None = None
-    """
-    Directory path to save the model file. If `None`, we save to the checkpoint directory set in `config.directory`.
-    """
-
-    filename: str | None = None
-    """
-    Checkpoint filename.
-        If None, a default template is used (see :attr:`ModelCheckpoint.CHECKPOINT_JOIN_CHAR`).
-    """
-
-    metric: MetricConfig | None = None
-    """
-    Metric to monitor for saving checkpoints.
-        If None, the primary metric of the runner will be used, if available.
-    """
-
-    verbose: bool = False
-    """Verbosity mode. If True, print additional information about checkpoints."""
-
-    save_last: Literal[True, False, "link"] | None = "link"
-    """
-    Whether to save the last checkpoint.
-        If True, saves a copy of the last checkpoint separately.
-        If "link", creates a symbolic link to the last checkpoint.
-    """
-
-    save_top_k: int | Literal["all"] = 1
-    """
-    Number of best models to save.
-        If "all" or -1, all models are saved.
-        If 0, no models are saved.
-    """
-
-    save_weights_only: bool = False
-    """Whether to save only the model's weights or the entire model object."""
-
-    auto_insert_metric_name: bool = True
-    """Whether to automatically insert the metric name in the checkpoint filename."""
-
-    every_n_train_steps: int | None = None
-    """
-    Number of training steps between checkpoints.
-        If None or 0, no checkpoints are saved during training.
-    """
-
-    train_time_interval: timedelta | None = None
-    """
-    Time interval between checkpoints during training.
-        If None, no checkpoints are saved during training based on time.
-    """
-
-    every_n_epochs: int | None = None
-    """
-    Number of epochs between checkpoints.
-        If None or 0, no checkpoints are saved at the end of epochs.
-    """
-
-    save_on_train_epoch_end: bool | None = None
-    """
-    Whether to run checkpointing at the end of the training epoch.
-        If False, checkpointing runs at the end of the validation.
-    """
-
-    enable_version_counter: bool = True
-    """Whether to append a version to the existing file name."""
-
-    auto_append_metric: bool = True
-    """If enabled, this will automatically add "-{monitor}" to the filename."""
-
-    def metric_or_default(self, root_config: "BaseConfig"):
-        if self.metric is not None:
-            return self.metric
-        if root_config.primary_metric is not None:
-            return root_config.primary_metric
-        raise ValueError("Primary metric must be provided if metric is not specified.")
-
-    def resolve_filename(self, root_config: "BaseConfig"):
-        metric = self.metric_or_default(root_config)
-
-        filename = self.filename
-        if not filename:
-            filename = "{epoch}-{step}"
-        if self.auto_append_metric:
-            filename = f"{filename}-{{{metric.validation_monitor}}}"
-
-        if self.auto_insert_metric_name and filename:
-            new_filename = _convert_string(filename)
-            log.critical(
-                f"Updated ModelCheckpoint filename: {filename} -> {new_filename}"
-            )
-            filename = new_filename
-
-        return filename
-
-    @override
-    def create_callbacks(self, root_config):
-        dirpath = self.dirpath or root_config.directory.resolve_subdirectory(
-            root_config.id, "checkpoint"
-        )
-
-        metric = self.metric_or_default(root_config)
-        filename = self.resolve_filename(root_config)
-
-        yield ModelCheckpoint(
-            self,
-            dirpath=Path(dirpath),
-            filename=filename,
-            metric=metric,
-        )
-
-    def _save_top_k_model_ckpt_input(self):
-        if self.save_top_k == "all":
-            return -1
-        return self.save_top_k
-
-
-class ModelCheckpoint(_ModelCheckpoint):
-    CHECKPOINT_NAME_LAST = "best"
-
-    @override
-    def __init__(
-        self,
-        config: ModelCheckpointCallbackConfig,
-        dirpath: Path,
-        filename: str,
-        metric: MetricConfig,
-    ):
-        self.config = config
-        del config
-
-        super().__init__(
-            dirpath=dirpath,
-            filename=filename,
-            monitor=metric.validation_monitor,
-            mode=metric.mode,
-            verbose=self.config.verbose,
-            save_last=self.config.save_last,
-            save_top_k=self.config._save_top_k_model_ckpt_input(),
-            save_weights_only=self.config.save_weights_only,
-            auto_insert_metric_name=False,
-            every_n_train_steps=self.config.every_n_train_steps,
-            train_time_interval=self.config.train_time_interval,
-            every_n_epochs=self.config.every_n_epochs,
-            save_on_train_epoch_end=self.config.save_on_train_epoch_end,
-            enable_version_counter=self.config.enable_version_counter,
-        )
-
-    @override
-    def _link_checkpoint(self, trainer: Trainer, filepath: str, linkpath: str):  # pyright: ignore[reportIncompatibleMethodOverride]
-        if trainer.is_global_zero:
-            _link_checkpoint(filepath, linkpath, metadata=True)
-        trainer.strategy.barrier()
-
-    @override
-    def _remove_checkpoint(self, trainer: Trainer, filepath: str):
-        _ckpt_saver_remove_checkpoint(trainer, filepath, metadata=True)

{nshtrainer-0.11.8 → nshtrainer-0.11.10}/src/nshtrainer/callbacks/checkpoint/best_checkpoint.py

@@ -46,8 +46,8 @@ class BestCheckpoint(CheckpointBase[BestCheckpointCallbackConfig]):
         dirpath: Path,
         metric: MetricConfig,
     ):
-        super().__init__(config, dirpath)
         self.metric = metric
+        super().__init__(config, dirpath)
 
     @override
     def name(self):