nshtrainer 1.0.0b53__tar.gz → 1.0.0b55__tar.gz

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: nshtrainer
-Version: 1.0.0b53
+Version: 1.0.0b55
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/pyproject.toml

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

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/configs/__init__.py

@@ -85,6 +85,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -306,6 +307,7 @@ __all__ = [
     "ProfilerConfig",
     "PyTorchProfilerConfig",
     "RLPSanityChecksCallbackConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "ReduceLROnPlateauConfig",
     "SLURMEnvironmentPlugin",

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/configs/lr_scheduler/__init__.py

@@ -12,6 +12,7 @@ from nshtrainer.lr_scheduler.base import lr_scheduler_registry as lr_scheduler_r
 from nshtrainer.lr_scheduler.linear_warmup_cosine import (
     DurationConfig as DurationConfig,
 )
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import MetricConfig as MetricConfig
 
 from . import base as base
@@ -20,6 +21,7 @@ from . import reduce_lr_on_plateau as reduce_lr_on_plateau
 
 __all__ = [
     "DurationConfig",
+    "EpochsConfig",
     "LRSchedulerConfig",
     "LRSchedulerConfigBase",
     "LinearWarmupCosineDecayLRSchedulerConfig",

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 __codegen__ = True
 
+from nshtrainer.lr_scheduler.reduce_lr_on_plateau import EpochsConfig as EpochsConfig
 from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
     LRSchedulerConfigBase as LRSchedulerConfigBase,
 )
@@ -14,6 +15,7 @@ from nshtrainer.lr_scheduler.reduce_lr_on_plateau import (
 )
 
 __all__ = [
+    "EpochsConfig",
     "LRSchedulerConfigBase",
     "MetricConfig",
     "ReduceLROnPlateauConfig",

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/configs/nn/__init__.py

@@ -11,6 +11,7 @@ from nshtrainer.nn import NonlinearityConfig as NonlinearityConfig
 from nshtrainer.nn import NonlinearityConfigBase as NonlinearityConfigBase
 from nshtrainer.nn import PReLUConfig as PReLUConfig
 from nshtrainer.nn import ReLUNonlinearityConfig as ReLUNonlinearityConfig
+from nshtrainer.nn import RNGConfig as RNGConfig
 from nshtrainer.nn import SigmoidNonlinearityConfig as SigmoidNonlinearityConfig
 from nshtrainer.nn import SiLUNonlinearityConfig as SiLUNonlinearityConfig
 from nshtrainer.nn import SoftmaxNonlinearityConfig as SoftmaxNonlinearityConfig
@@ -25,6 +26,7 @@ from nshtrainer.nn.nonlinearity import nonlinearity_registry as nonlinearity_reg
 
 from . import mlp as mlp
 from . import nonlinearity as nonlinearity
+from . import rng as rng
 
 __all__ = [
     "ELUNonlinearityConfig",
@@ -35,6 +37,7 @@ __all__ = [
     "NonlinearityConfig",
     "NonlinearityConfigBase",
     "PReLUConfig",
+    "RNGConfig",
     "ReLUNonlinearityConfig",
     "SiLUNonlinearityConfig",
     "SigmoidNonlinearityConfig",
@@ -47,4 +50,5 @@ __all__ = [
     "mlp",
     "nonlinearity",
     "nonlinearity_registry",
+    "rng",
 ]

nshtrainer-1.0.0b55/src/nshtrainer/configs/nn/rng/__init__.py

@@ -0,0 +1,9 @@
+from __future__ import annotations
+
+__codegen__ = True
+
+from nshtrainer.nn.rng import RNGConfig as RNGConfig
+
+__all__ = [
+    "RNGConfig",
+]

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/nn/__init__.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from .mlp import MLP as MLP
 from .mlp import MLPConfig as MLPConfig
 from .mlp import ResidualSequential as ResidualSequential
-from .mlp import custom_seed_context as custom_seed_context
 from .module_dict import TypedModuleDict as TypedModuleDict
 from .module_list import TypedModuleList as TypedModuleList
 from .nonlinearity import ELUNonlinearityConfig as ELUNonlinearityConfig
@@ -21,3 +20,5 @@ from .nonlinearity import SoftplusNonlinearityConfig as SoftplusNonlinearityConf
 from .nonlinearity import SoftsignNonlinearityConfig as SoftsignNonlinearityConfig
 from .nonlinearity import SwishNonlinearityConfig as SwishNonlinearityConfig
 from .nonlinearity import TanhNonlinearityConfig as TanhNonlinearityConfig
+from .rng import RNGConfig as RNGConfig
+from .rng import rng_context as rng_context

nshtrainer-1.0.0b55/src/nshtrainer/nn/rng.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import contextlib
+
+import nshconfig as C
+import torch
+
+
+@contextlib.contextmanager
+def rng_context(config: RNGConfig | None):
+    with contextlib.ExitStack() as stack:
+        if config is not None:
+            stack.enter_context(
+                torch.random.fork_rng(devices=range(torch.cuda.device_count()))
+            )
+            torch.manual_seed(config.seed)
+
+        yield
+
+
+class RNGConfig(C.Config):
+    seed: int
+    """Random seed to use for initialization."""

{nshtrainer-1.0.0b53 → nshtrainer-1.0.0b55}/src/nshtrainer/trainer/trainer.py

@@ -457,7 +457,21 @@ class Trainer(LightningTrainer):
     ):
         filepath = Path(filepath)
 
-        super().save_checkpoint(filepath, weights_only, storage_options)
+        if self.model is None:
+            raise AttributeError(
+                "Saving a checkpoint is only possible if a model is attached to the Trainer. Did you call"
+                " `Trainer.save_checkpoint()` before calling `Trainer.{fit,validate,test,predict}`?"
+            )
+        with self.profiler.profile("save_checkpoint"):  # type: ignore
+            checkpoint = self._checkpoint_connector.dump_checkpoint(weights_only)
+            # Update the checkpoint for the trainer hyperparameters
+            checkpoint[self.CHECKPOINT_HYPER_PARAMS_KEY] = self.hparams.model_dump(
+                mode="json"
+            )
+            self.strategy.save_checkpoint(
+                checkpoint, filepath, storage_options=storage_options
+            )
+            self.strategy.barrier("Trainer.save_checkpoint")
 
         # Save the checkpoint metadata
         metadata_path = None