nshtrainer 1.0.0b11__tar.gz → 1.0.0b12__tar.gz

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b12}/PKG-INFO

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

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b12}/pyproject.toml

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

nshtrainer-1.0.0b12/src/nshtrainer/data/datamodule.py

@@ -0,0 +1,126 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Mapping
+from pathlib import Path
+from typing import Any, Generic, cast
+
+import nshconfig as C
+import torch
+from lightning.pytorch import LightningDataModule
+from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.rank_zero import rank_zero_warn
+from typing_extensions import Never, TypeVar, deprecated, override
+
+from ..model.mixins.callback import CallbackRegistrarModuleMixin
+from ..model.mixins.debug import _DebugModuleMixin
+
+THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
+
+
+class LightningDataModuleBase(
+    _DebugModuleMixin,
+    CallbackRegistrarModuleMixin,
+    LightningDataModule,
+    ABC,
+    Generic[THparams],
+):
+    @property
+    @override
+    def hparams(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
+        return cast(THparams, super().hparams)
+
+    @property
+    @override
+    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
+        hparams = cast(THparams, super().hparams_initial)
+        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
+
+    @property
+    @deprecated("Use `hparams` instead")
+    def config(self):
+        return cast(Never, self.hparams)
+
+    @classmethod
+    @abstractmethod
+    def hparams_cls(cls) -> type[THparams]: ...
+
+    @override
+    def __init__(self, hparams: THparams | Mapping[str, Any]):
+        super().__init__()
+
+        # Validate and save hyperparameters
+        hparams_cls = self.hparams_cls()
+        if isinstance(hparams, Mapping):
+            hparams = hparams_cls.model_validate(hparams)
+        elif not isinstance(hparams, hparams_cls):
+            raise TypeError(
+                f"Expected hparams to be either a Mapping or an instance of {hparams_cls}, got {type(hparams)}"
+            )
+        hparams = hparams.model_deep_validate()
+        self.save_hyperparameters(hparams)
+
+    @override
+    @classmethod
+    def load_from_checkpoint(cls, *args, **kwargs) -> Never:
+        raise ValueError("This method is not supported. Use `from_checkpoint` instead.")
+
+    @classmethod
+    def hparams_from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
+
+        if (hparams := ckpt.get(cls.CHECKPOINT_HYPER_PARAMS_KEY)) is None:
+            raise ValueError(
+                f"The checkpoint does not contain hyperparameters. It must contain the key '{cls.CHECKPOINT_HYPER_PARAMS_KEY}'."
+            )
+        if update_hparams_dict is not None:
+            hparams = update_hparams_dict(hparams)
+
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
+
+        return hparams
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        map_location: torch.serialization.MAP_LOCATION = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        # Load checkpoint
+        if isinstance(ckpt_or_path, Mapping):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location=map_location)
+
+        # Load hyperparameters from checkpoint
+        hparams = cls.hparams_from_checkpoint(
+            ckpt,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+
+        # Load datamodule from checkpoint
+        datamodule = cls(hparams)
+        if datamodule.__class__.__qualname__ in ckpt:
+            datamodule.load_state_dict(ckpt[datamodule.__class__.__qualname__])
+
+        return datamodule

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b12}/src/nshtrainer/model/base.py

@@ -2,7 +2,8 @@ from __future__ import annotations
 
 import logging
 from abc import ABC, abstractmethod
-from collections.abc import Mapping
+from collections.abc import Callable, Mapping
+from pathlib import Path
 from typing import Any, Generic, Literal, cast
 
 import nshconfig as C
@@ -10,11 +11,13 @@ import torch
 import torch.distributed
 from lightning.pytorch import LightningModule
 from lightning.pytorch.profilers import PassThroughProfiler, Profiler
+from lightning.pytorch.utilities.model_helpers import is_overridden
+from lightning.pytorch.utilities.rank_zero import rank_zero_warn
 from typing_extensions import Never, TypeVar, deprecated, override
 
 from ..callbacks.rlp_sanity_checks import _RLPSanityCheckModuleMixin
 from .mixins.callback import CallbackModuleMixin
-from .mixins.debug import _DebugModuleMixin, _trainer
+from .mixins.debug import _DebugModuleMixin
 from .mixins.logger import LoggerLightningModuleMixin
 
 log = logging.getLogger(__name__)
@@ -241,3 +244,98 @@ class LightningModuleBase(
         loss = sum((0.0 * v).sum() for v in self.parameters() if v.requires_grad)
         loss = cast(torch.Tensor, loss)
         return loss
+
+    @override
+    @classmethod
+    def load_from_checkpoint(cls, *args, **kwargs) -> Never:
+        raise ValueError("This method is not supported. Use `from_checkpoint` instead.")
+
+    @classmethod
+    def hparams_from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
+
+        if (hparams := ckpt.get(cls.CHECKPOINT_HYPER_PARAMS_KEY)) is None:
+            raise ValueError(
+                f"The checkpoint does not contain hyperparameters. It must contain the key '{cls.CHECKPOINT_HYPER_PARAMS_KEY}'."
+            )
+        if update_hparams_dict is not None:
+            hparams = update_hparams_dict(hparams)
+
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
+
+        return hparams
+
+    @classmethod
+    def from_checkpoint(
+        cls,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        map_location: torch.serialization.MAP_LOCATION = None,
+        *,
+        update_hparams: Callable[[THparams], THparams] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        # Load checkpoint
+        if isinstance(ckpt_or_path, Mapping):
+            ckpt = ckpt_or_path
+        else:
+            ckpt = torch.load(ckpt_or_path, map_location=map_location)
+
+        # Load hyperparameters from checkpoint
+        hparams = cls.hparams_from_checkpoint(
+            ckpt,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+
+        # Load model from checkpoint
+        model = cls(hparams)
+
+        # Load model state from checkpoint
+        if (
+            model._strict_loading is not None
+            and strict is not None
+            and strict != model.strict_loading
+        ):
+            raise ValueError(
+                f"You set `.load_from_checkpoint(..., strict={strict!r})` which is in conflict with"
+                f" `{cls.__name__}.strict_loading={model.strict_loading!r}. Please set the same value for both of them."
+            )
+        strict = model.strict_loading if strict is None else strict
+
+        if is_overridden("configure_model", model):
+            model.configure_model()
+
+        # give model a chance to load something
+        model.on_load_checkpoint(ckpt)
+
+        # load the state_dict on the model automatically
+
+        keys = model.load_state_dict(ckpt["state_dict"], strict=strict)
+
+        if not strict:
+            if keys.missing_keys:
+                rank_zero_warn(
+                    f"Found keys that are in the model state dict but not in the checkpoint: {keys.missing_keys}"
+                )
+            if keys.unexpected_keys:
+                rank_zero_warn(
+                    f"Found keys that are not in the model state dict but in the checkpoint: {keys.unexpected_keys}"
+                )
+
+        return model

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b12}/src/nshtrainer/trainer/_config.py

@@ -9,7 +9,6 @@ from collections.abc import Iterable, Sequence
 from datetime import timedelta
 from pathlib import Path
 from typing import (
-    TYPE_CHECKING,
     Annotated,
     Any,
     ClassVar,

{nshtrainer-1.0.0b11 → nshtrainer-1.0.0b12}/src/nshtrainer/trainer/trainer.py

@@ -2,28 +2,19 @@ from __future__ import annotations
 
 import logging
 import os
-from collections.abc import Mapping, Sequence
+from collections.abc import Callable, Mapping, Sequence
 from pathlib import Path
-from typing import IO, TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, cast
 
 import torch
 from lightning.fabric.plugins.environments.lsf import LSFEnvironment
 from lightning.fabric.plugins.environments.slurm import SLURMEnvironment
 from lightning.fabric.plugins.precision.precision import _PRECISION_INPUT
-from lightning.fabric.utilities.cloud_io import _load as pl_load
-from lightning.fabric.utilities.types import _MAP_LOCATION_TYPE, _PATH
 from lightning.pytorch import LightningModule
 from lightning.pytorch import Trainer as LightningTrainer
 from lightning.pytorch.callbacks import Callback
-from lightning.pytorch.core.saving import (
-    _default_map_location,
-    load_hparams_from_tags_csv,
-    load_hparams_from_yaml,
-)
 from lightning.pytorch.profilers import Profiler
 from lightning.pytorch.trainer.states import TrainerFn
-from lightning.pytorch.utilities.migration import pl_legacy_patch
-from lightning.pytorch.utilities.migration.utils import _pl_migrate_checkpoint
 from lightning.pytorch.utilities.types import _EVALUATE_OUTPUT, _PREDICT_OUTPUT
 from typing_extensions import Never, Unpack, assert_never, deprecated, override
 
@@ -473,62 +464,46 @@ class Trainer(LightningTrainer):
         _callback._call_on_checkpoint_saved(self, filepath, metadata_path)
 
     @classmethod
-    def load_from_checkpoint(
+    def hparams_from_checkpoint(
         cls,
-        checkpoint_path: _PATH | IO,
-        map_location: _MAP_LOCATION_TYPE = None,
-        hparams_file: _PATH | None = None,
-        **kwargs: Any,
+        ckpt_or_path: dict[str, Any] | str | Path,
+        /,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[TrainerConfig], TrainerConfig] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
     ):
-        loaded = _load_from_checkpoint(
-            checkpoint_path,
-            map_location=map_location,
-            hparams_file=hparams_file,
-            **kwargs,
-        )
-        return loaded
-
-
-def _load_from_checkpoint(
-    checkpoint_path: _PATH | IO,
-    map_location: _MAP_LOCATION_TYPE = None,
-    hparams_file: _PATH | None = None,
-    **kwargs: Any,
-):
-    map_location = map_location or _default_map_location
-    with pl_legacy_patch():
-        checkpoint = pl_load(checkpoint_path, map_location=map_location)
-
-    # convert legacy checkpoints to the new format
-    checkpoint = _pl_migrate_checkpoint(
-        checkpoint,
-        checkpoint_path=(
-            checkpoint_path if isinstance(checkpoint_path, (str, Path)) else None
-        ),
-    )
-
-    if hparams_file is not None:
-        extension = str(hparams_file).split(".")[-1]
-        if extension.lower() == "csv":
-            hparams = load_hparams_from_tags_csv(hparams_file)
-        elif extension.lower() in ("yml", "yaml"):
-            hparams = load_hparams_from_yaml(hparams_file)
+        if isinstance(ckpt_or_path, dict):
+            ckpt = ckpt_or_path
         else:
-            raise ValueError(".csv, .yml or .yaml is required for `hparams_file`")
+            ckpt = torch.load(ckpt_or_path, map_location="cpu")
 
-        # overwrite hparams by the given file
-        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY] = hparams
+        if (hparams := ckpt.get(cls.CHECKPOINT_HYPER_PARAMS_KEY)) is None:
+            raise ValueError(
+                f"The checkpoint does not contain hyperparameters. It must contain the key '{cls.CHECKPOINT_HYPER_PARAMS_KEY}'."
+            )
+        if update_hparams_dict is not None:
+            hparams = update_hparams_dict(hparams)
 
-    # for past checkpoint need to add the new key
-    checkpoint.setdefault(Trainer.CHECKPOINT_HYPER_PARAMS_KEY, {})
-    # override the hparams with values that were passed in
-    checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY].update(kwargs)
+        hparams = cls.hparams_cls().model_validate(hparams, strict=strict)
+        if update_hparams is not None:
+            hparams = update_hparams(hparams)
 
-    # load the hparams
-    hparams = Trainer.hparams_cls().model_validate(
-        checkpoint[Trainer.CHECKPOINT_HYPER_PARAMS_KEY]
-    )
+        return hparams
 
-    # create the trainer
-    trainer = Trainer(hparams)
-    return trainer
+    @classmethod
+    def from_checkpoint(
+        cls,
+        path: str | Path,
+        strict: bool | None = None,
+        *,
+        update_hparams: Callable[[TrainerConfig], TrainerConfig] | None = None,
+        update_hparams_dict: Callable[[dict[str, Any]], dict[str, Any]] | None = None,
+    ):
+        hparams = cls.hparams_from_checkpoint(
+            path,
+            strict=strict,
+            update_hparams=update_hparams,
+            update_hparams_dict=update_hparams_dict,
+        )
+        return cls(hparams)

nshtrainer-1.0.0b11/src/nshtrainer/data/datamodule.py

@@ -1,57 +0,0 @@
-from __future__ import annotations
-
-from abc import ABC, abstractmethod
-from collections.abc import Mapping
-from typing import Any, Generic, cast
-
-import nshconfig as C
-from lightning.pytorch import LightningDataModule
-from typing_extensions import Never, TypeVar, deprecated, override
-
-from ..model.mixins.callback import CallbackRegistrarModuleMixin
-from ..model.mixins.debug import _DebugModuleMixin
-
-THparams = TypeVar("THparams", bound=C.Config, infer_variance=True)
-
-
-class LightningDataModuleBase(
-    _DebugModuleMixin,
-    CallbackRegistrarModuleMixin,
-    LightningDataModule,
-    ABC,
-    Generic[THparams],
-):
-    @property
-    @override
-    def hparams(self) -> THparams:  # pyright: ignore[reportIncompatibleMethodOverride]
-        return cast(THparams, super().hparams)
-
-    @property
-    @override
-    def hparams_initial(self):  # pyright: ignore[reportIncompatibleMethodOverride]
-        hparams = cast(THparams, super().hparams_initial)
-        return cast(Never, {"datamodule": hparams.model_dump(mode="json")})
-
-    @property
-    @deprecated("Use `hparams` instead")
-    def config(self):
-        return cast(Never, self.hparams)
-
-    @classmethod
-    @abstractmethod
-    def hparams_cls(cls) -> type[THparams]: ...
-
-    @override
-    def __init__(self, hparams: THparams | Mapping[str, Any]):
-        super().__init__()
-
-        # Validate and save hyperparameters
-        hparams_cls = self.hparams_cls()
-        if isinstance(hparams, Mapping):
-            hparams = hparams_cls.model_validate(hparams)
-        elif not isinstance(hparams, hparams_cls):
-            raise TypeError(
-                f"Expected hparams to be either a Mapping or an instance of {hparams_cls}, got {type(hparams)}"
-            )
-        hparams = hparams.model_deep_validate()
-        self.save_hyperparameters(hparams)

nshtrainer-1.0.0b11/src/nshtrainer/scripts/find_packages.py

@@ -1,52 +0,0 @@
-from __future__ import annotations
-
-import argparse
-import ast
-import glob
-import sys
-from pathlib import Path
-
-
-def get_imports(file_path: Path):
-    with open(file_path, "r") as file:
-        try:
-            tree = ast.parse(file.read())
-        except SyntaxError:
-            print(f"Syntax error in file: {file_path}", file=sys.stderr)
-            return set()
-
-    imports = set()
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Import):
-            for alias in node.names:
-                imports.add(alias.name.split(".")[0])
-        elif isinstance(node, ast.ImportFrom):
-            if node.level == 0 and node.module:  # Absolute import
-                imports.add(node.module.split(".")[0])
-    return imports
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="Find unique Python packages used in files."
-    )
-    parser.add_argument("glob_pattern", help="Glob pattern to match files")
-    parser.add_argument(
-        "--exclude-std", action="store_true", help="Exclude Python standard libraries"
-    )
-    args = parser.parse_args()
-
-    all_imports = set()
-    for file_path in glob.glob(args.glob_pattern, recursive=True):
-        all_imports.update(get_imports(Path(file_path)))
-
-    if args.exclude_std:
-        std_libs = set(sys.stdlib_module_names)
-        all_imports = all_imports - std_libs
-
-    for package in sorted(all_imports):
-        print(package)
-
-
-if __name__ == "__main__":
-    main()