nshtrainer 0.3.0__tar.gz → 0.4.0__tar.gz

{nshtrainer-0.3.0 → nshtrainer-0.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.3.0
+Version: 0.4.0
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com
@@ -14,7 +14,7 @@ Requires-Dist: lovely-numpy (>=0.2.13,<0.3.0)
 Requires-Dist: lovely-tensors (>=0.1.16,<0.2.0)
 Requires-Dist: nshconfig (>=0.2.0,<0.3.0)
 Requires-Dist: nshrunner (>=0.5.4,<0.6.0)
-Requires-Dist: nshutils (>=0.2.0,<0.3.0)
+Requires-Dist: nshutils (>=0.3.0,<0.4.0)
 Requires-Dist: numpy
 Requires-Dist: pytorch-lightning
 Requires-Dist: rich

{nshtrainer-0.3.0 → nshtrainer-0.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.3.0"
+version = "0.4.0"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"
@@ -19,7 +19,7 @@ lovely-numpy = "^0.2.13"
 lovely-tensors = "^0.1.16"
 wrapt = "*"
 rich = "*"
-nshutils = "^0.2.0"
+nshutils = "^0.3.0"
 
 
 [tool.poetry.group.dev.dependencies]

nshtrainer-0.4.0/src/nshtrainer/actsave/__init__.py

@@ -0,0 +1,3 @@
+from nshutils.actsave import *  # type: ignore # noqa: F403
+
+from ._callback import ActSaveCallback as ActSaveCallback

{nshtrainer-0.3.0 → nshtrainer-0.4.0}/src/nshtrainer/actsave/_callback.py

@@ -3,10 +3,9 @@ from typing import TYPE_CHECKING, Literal, cast
 
 from lightning.pytorch import LightningModule, Trainer
 from lightning.pytorch.callbacks.callback import Callback
+from nshutils.actsave import ActSave
 from typing_extensions import TypeAlias, override
 
-from ._saver import ActSave
-
 if TYPE_CHECKING:
     from ..model.config import BaseConfig
 

nshtrainer-0.3.0/src/nshtrainer/actsave/__init__.py

@@ -1,7 +0,0 @@
-from ._callback import ActSaveCallback as ActSaveCallback
-from ._loader import ActivationLoader as ActivationLoader
-from ._loader import ActLoad as ActLoad
-from ._saver import Activation as Activation
-from ._saver import ActivationSaver as ActivationSaver
-from ._saver import ActSave as ActSave
-from ._saver import Transform as Transform

nshtrainer-0.3.0/src/nshtrainer/actsave/_loader.py

@@ -1,144 +0,0 @@
-import pprint
-from dataclasses import dataclass, field
-from functools import cached_property
-from logging import getLogger
-from pathlib import Path
-from typing import cast, overload
-
-import numpy as np
-from typing_extensions import TypeVar, override
-
-log = getLogger(__name__)
-
-T = TypeVar("T", infer_variance=True)
-
-
-@dataclass
-class LoadedActivation:
-    base_dir: Path = field(repr=False)
-    name: str
-    num_activations: int = field(init=False)
-    activation_files: list[Path] = field(init=False, repr=False)
-
-    def __post_init__(self):
-        if not self.activation_dir.exists():
-            raise ValueError(f"Activation dir {self.activation_dir} does not exist")
-
-        # The number of activations = the * of .npy files in the activation dir
-        self.activation_files = list(self.activation_dir.glob("*.npy"))
-        # Sort the activation files by the numerical index in the filename
-        self.activation_files.sort(key=lambda p: int(p.stem))
-        self.num_activations = len(self.activation_files)
-
-    @property
-    def activation_dir(self) -> Path:
-        return self.base_dir / self.name
-
-    def _load_activation(self, item: int):
-        activation_path = self.activation_files[item]
-        if not activation_path.exists():
-            raise ValueError(f"Activation {activation_path} does not exist")
-        return cast(np.ndarray, np.load(activation_path, allow_pickle=True))
-
-    @overload
-    def __getitem__(self, item: int) -> np.ndarray: ...
-
-    @overload
-    def __getitem__(self, item: slice | list[int]) -> list[np.ndarray]: ...
-
-    def __getitem__(
-        self, item: int | slice | list[int]
-    ) -> np.ndarray | list[np.ndarray]:
-        if isinstance(item, int):
-            return self._load_activation(item)
-        elif isinstance(item, slice):
-            return [
-                self._load_activation(i)
-                for i in range(*item.indices(self.num_activations))
-            ]
-        elif isinstance(item, list):
-            return [self._load_activation(i) for i in item]
-        else:
-            raise TypeError(f"Invalid type {type(item)} for item {item}")
-
-    def __iter__(self):
-        return iter(self[i] for i in range(self.num_activations))
-
-    def __len__(self):
-        return self.num_activations
-
-    def all_activations(self):
-        return [self[i] for i in range(self.num_activations)]
-
-    @override
-    def __repr__(self):
-        return f"<LoadedActivation {self.name} ({self.num_activations} activations)>"
-
-
-class ActLoad:
-    @classmethod
-    def all_versions(cls, dir: str | Path):
-        dir = Path(dir)
-
-        # If the dir is not an activation base directory, we return None
-        if not (dir / ".activationbase").exists():
-            return None
-
-        # The contents of `dir` should be directories, each of which is a version.
-        return [
-            (subdir, int(subdir.name)) for subdir in dir.iterdir() if subdir.is_dir()
-        ]
-
-    @classmethod
-    def is_valid_activation_base(cls, dir: str | Path):
-        return cls.all_versions(dir) is not None
-
-    @classmethod
-    def from_latest_version(cls, dir: str | Path):
-        # The contents of `dir` should be directories, each of which is a version
-        # We need to find the latest version
-        if (all_versions := cls.all_versions(dir)) is None:
-            raise ValueError(f"{dir} is not an activation base directory")
-
-        path, _ = max(all_versions, key=lambda p: p[1])
-        return cls(path)
-
-    def __init__(self, dir: Path):
-        self._dir = dir
-
-    def activation(self, name: str):
-        return LoadedActivation(self._dir, name)
-
-    @cached_property
-    def activations(self):
-        dirs = list(self._dir.iterdir())
-        # Sort the dirs by the last modified time
-        dirs.sort(key=lambda p: p.stat().st_mtime)
-
-        return {p.name: LoadedActivation(self._dir, p.name) for p in dirs}
-
-    def __iter__(self):
-        return iter(self.activations.values())
-
-    def __getitem__(self, item: str):
-        return self.activations[item]
-
-    def __len__(self):
-        return len(self.activations)
-
-    @override
-    def __repr__(self):
-        acts_str = pprint.pformat(
-            {
-                name: f"<{activation.num_activations} activations>"
-                for name, activation in self.activations.items()
-            }
-        )
-        acts_str = acts_str.replace("'<", "<").replace(">'", ">")
-        return f"ActLoad({acts_str})"
-
-    def get(self, name: str, /, default: T) -> LoadedActivation | T:
-        return self.activations.get(name, default)
-
-
-ActivationLoader = ActLoad

nshtrainer-0.3.0/src/nshtrainer/actsave/_saver.py

@@ -1,337 +0,0 @@
-import contextlib
-import fnmatch
-import tempfile
-import uuid
-import weakref
-from collections.abc import Callable, Mapping
-from dataclasses import dataclass
-from functools import wraps
-from logging import getLogger
-from pathlib import Path
-from typing import Generic, TypeAlias, cast, overload
-
-import numpy as np
-import torch
-from lightning_utilities.core.apply_func import apply_to_collection
-from typing_extensions import ParamSpec, TypeVar, override
-
-log = getLogger(__name__)
-
-Value: TypeAlias = int | float | complex | bool | str | np.ndarray | torch.Tensor | None
-ValueOrLambda = Value | Callable[..., Value]
-
-
-def _to_numpy(activation: Value) -> np.ndarray:
-    # Make sure it's not `None`
-    if activation is None:
-        raise ValueError("Activation should not be `None`")
-
-    if isinstance(activation, np.ndarray):
-        return activation
-    if isinstance(activation, torch.Tensor):
-        activation = activation.detach()
-        if activation.is_floating_point():
-            # NOTE: We need to convert to float32 because [b]float16 is not supported by numpy
-            activation = activation.float()
-        return activation.cpu().numpy()
-    if isinstance(activation, (int, float, complex, str, bool)):
-        return np.array(activation)
-
-    return activation
-
-
-T = TypeVar("T", infer_variance=True)
-
-
-# A wrapper around weakref.ref that allows for primitive types
-# To get around errors like:
-# TypeError: cannot create weak reference to 'int' object
-class WeakRef(Generic[T]):
-    _ref: Callable[[], T] | None
-
-    def __init__(self, obj: T):
-        try:
-            self._ref = cast(Callable[[], T], weakref.ref(obj))
-        except TypeError as e:
-            if "cannot create weak reference" not in str(e):
-                raise
-            self._ref = lambda: obj
-
-    def __call__(self) -> T:
-        if self._ref is None:
-            raise RuntimeError("WeakRef is deleted")
-        return self._ref()
-
-    def delete(self):
-        del self._ref
-        self._ref = None
-
-
-@dataclass
-class Activation:
-    name: str
-    ref: WeakRef[ValueOrLambda] | None
-    transformed: np.ndarray | None = None
-
-    def __post_init__(self):
-        # Update the `name` to replace `/` with `.`
-        self.name = self.name.replace("/", ".")
-
-    def __call__(self) -> np.ndarray | None:
-        # If we have a transformed value, we return it
-        if self.transformed is not None:
-            return self.transformed
-
-        if self.ref is None:
-            raise RuntimeError("Activation is deleted")
-
-        # If we have a lambda, we need to call it
-        unrwapped_ref = self.ref()
-        activation = unrwapped_ref
-        if callable(unrwapped_ref):
-            activation = unrwapped_ref()
-
-        # If we have a `None`, we return early
-        if activation is None:
-            return None
-
-        activation = apply_to_collection(activation, torch.Tensor, _to_numpy)
-        activation = _to_numpy(activation)
-
-        # Set the transformed value
-        self.transformed = activation
-
-        # Delete the reference
-        self.ref.delete()
-        del self.ref
-        self.ref = None
-
-        return self.transformed
-
-    @classmethod
-    def from_value_or_lambda(cls, name: str, value_or_lambda: ValueOrLambda):
-        return cls(name, WeakRef(value_or_lambda))
-
-    @classmethod
-    def from_dict(cls, d: Mapping[str, ValueOrLambda]):
-        return [cls.from_value_or_lambda(k, v) for k, v in d.items()]
-
-
-Transform = Callable[[Activation], Mapping[str, ValueOrLambda]]
-
-
-def _ensure_supported():
-    try:
-        import torch.distributed as dist
-
-        if dist.is_initialized() and dist.get_world_size() > 1:
-            raise RuntimeError("Only single GPU is supported at the moment")
-    except ImportError:
-        pass
-
-
-P = ParamSpec("P")
-
-
-def _ignore_if_scripting(fn: Callable[P, None]) -> Callable[P, None]:
-    @wraps(fn)
-    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
-        if torch.jit.is_scripting():
-            return
-
-        _ensure_supported()
-        fn(*args, **kwargs)
-
-    return wrapper
-
-
-class _Saver:
-    def __init__(
-        self,
-        save_dir: Path,
-        prefixes_fn: Callable[[], list[str]],
-        *,
-        filters: list[str] | None = None,
-    ):
-        # Create a directory under `save_dir` by autoincrementing
-        # (i.e., every activation save context, we create a new directory)
-        # The id = the number of activation subdirectories
-        self._id = sum(1 for subdir in save_dir.glob("*") if subdir.is_dir())
-        save_dir.mkdir(parents=True, exist_ok=True)
-
-        # Add a .activationbase file to the save_dir to indicate that this is an activation base
-        (save_dir / ".activationbase").touch(exist_ok=True)
-
-        self._save_dir = save_dir / f"{self._id:04d}"
-        # Make sure `self._save_dir` does not exist and create it
-        self._save_dir.mkdir(exist_ok=False)
-
-        self._prefixes_fn = prefixes_fn
-        self._filters = filters
-
-    def _save_activation(self, activation: Activation):
-        # If the activation value is `None`, we skip it.
-        if (activation_value := activation()) is None:
-            return
-
-        # Save the activation to self._save_dir / name / {id}.npz, where id is an auto-incrementing integer
-        file_name = ".".join(self._prefixes_fn() + [activation.name])
-        path = self._save_dir / file_name
-        path.mkdir(exist_ok=True, parents=True)
-
-        # Get the next id and save the activation
-        id = len(list(path.glob("*.npy")))
-        np.save(path / f"{id:04d}.npy", activation_value)
-
-    @_ignore_if_scripting
-    def save(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        kwargs.update(acts or {})
-
-        # Build activations
-        activations = Activation.from_dict(kwargs)
-
-        for activation in activations:
-            # Make sure name matches at least one filter if filters are specified
-            if self._filters is not None and all(
-                not fnmatch.fnmatch(activation.name, f) for f in self._filters
-            ):
-                continue
-
-            # Save the current activation
-            self._save_activation(activation)
-
-        del activations
-
-
-class ActSaveProvider:
-    _saver: _Saver | None = None
-    _prefixes: list[str] = []
-
-    def initialize(self, save_dir: Path | None = None):
-        """
-        Initializes the saver with the given configuration and save directory.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        if self._saver is None:
-            if save_dir is None:
-                save_dir = Path(tempfile.gettempdir()) / f"actsave-{uuid.uuid4()}"
-                log.critical(f"No save_dir specified, using {save_dir=}")
-            self._saver = _Saver(
-                save_dir,
-                lambda: self._prefixes,
-            )
-
-    @contextlib.contextmanager
-    def enabled(self, save_dir: Path | None = None):
-        """
-        Context manager that enables the actsave functionality with the specified configuration.
-
-        Args:
-            save_dir (Path): The directory where the saved files will be stored.
-        """
-        prev = self._saver
-        self.initialize(save_dir)
-        try:
-            yield
-        finally:
-            self._saver = prev
-
-    @override
-    def __init__(self):
-        super().__init__()
-
-        self._saver = None
-        self._prefixes = []
-
-    @contextlib.contextmanager
-    def context(self, label: str):
-        """
-        A context manager that adds a label to the current context.
-
-        Args:
-            label (str): The label for the context.
-        """
-        if torch.jit.is_scripting():
-            yield
-            return
-
-        if self._saver is None:
-            yield
-            return
-
-        _ensure_supported()
-
-        log.debug(f"Entering ActSave context {label}")
-        self._prefixes.append(label)
-        try:
-            yield
-        finally:
-            _ = self._prefixes.pop()
-
-    prefix = context
-
-    @overload
-    def __call__(
-        self,
-        acts: dict[str, ValueOrLambda] | None = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (dict[str, ValueOrLambda] | None, optional): A dictionary of acts. Defaults to None.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    @overload
-    def __call__(self, acts: Callable[[], dict[str, ValueOrLambda]], /):
-        """
-        Saves the activations to disk.
-
-        Args:
-            acts (Callable[[], dict[str, ValueOrLambda]]): A callable that returns a dictionary of acts.
-            **kwargs (ValueOrLambda): Additional keyword arguments.
-
-        Returns:
-            None
-
-        """
-        ...
-
-    def __call__(
-        self,
-        acts: (
-            dict[str, ValueOrLambda] | Callable[[], dict[str, ValueOrLambda]] | None
-        ) = None,
-        /,
-        **kwargs: ValueOrLambda,
-    ):
-        if torch.jit.is_scripting():
-            return
-
-        if self._saver is None:
-            return
-
-        if acts is not None and callable(acts):
-            acts = acts()
-        self._saver.save(acts, **kwargs)
-
-    save = __call__
-
-
-ActSave = ActSaveProvider()
-ActivationSaver = ActSave