nshtrainer 1.0.0b53__py3-none-any.whl → 1.0.0b55__py3-none-any.whl

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/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/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/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/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/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/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/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

{nshtrainer-1.0.0b53.dist-info → nshtrainer-1.0.0b55.dist-info}/METADATA

@@ -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.dist-info → nshtrainer-1.0.0b55.dist-info}/RECORD

@@ -32,7 +32,7 @@ nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU
 nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
 nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
 nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
-nshtrainer/configs/__init__.py,sha256=0BzCgE1iEJ0Ywmy__mqJZipLQtwZVdz6XK-gHbkA7GY,14650
+nshtrainer/configs/__init__.py,sha256=-rGk9pnRnuz4yKvACGOpY3nkrWnHholqZGk7UP2Vkrc,14716
 nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
 nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
 nshtrainer/configs/_directory/__init__.py,sha256=_oO7vM9DhzHSxtZcv86sTi7hZIptnK1gr-AP9mqQ370,386
@@ -67,15 +67,16 @@ nshtrainer/configs/loggers/base/__init__.py,sha256=HLUfEDbjaAXqzsFmQbjdciIWzR1st
 nshtrainer/configs/loggers/csv/__init__.py,sha256=gawaDX92JObGSmBqYpfNHWMHBwVOofS694W-1Y2GWDU,353
 nshtrainer/configs/loggers/tensorboard/__init__.py,sha256=phzm-TnBkdkibTgoOxIIcAliqL3zU8gSNK61Mwxs1CM,410
 nshtrainer/configs/loggers/wandb/__init__.py,sha256=TDcD5WZSKenc2mgIXhwz2l96l8P_Ur3N5CzEol5AKGw,746
-nshtrainer/configs/lr_scheduler/__init__.py,sha256=xtiUx0isxA82-uXMn4-KmPnDCfbUkpAnd2_pFupAAKQ,1137
+nshtrainer/configs/lr_scheduler/__init__.py,sha256=PvH2d8QEC3TsC3_svcUbxeQEMMzIK_In0_Bp9xntSms,1243
 nshtrainer/configs/lr_scheduler/base/__init__.py,sha256=6Cx8r4rdxeSYxc_z0o7drKCblGJU_zzqrOoYlWYR5qY,305
 nshtrainer/configs/lr_scheduler/linear_warmup_cosine/__init__.py,sha256=5ZMLDO9VL6SNU6pF-62lDnpmqix3_Ol9DdEwiuOPYlA,675
-nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=w-vq8UbRGPX8DZVWCMC5eIrbvVc_guxjj7Du9AaeKCw,609
+nshtrainer/configs/lr_scheduler/reduce_lr_on_plateau/__init__.py,sha256=DBwZV590I5qwyOS5M43YhUzgYy1-AjzkM5aEnTA6XdI,715
 nshtrainer/configs/metrics/__init__.py,sha256=mK_xgXJDAyGY4K_x_Zo4aj36kjT45d850keuUe3U1rY,200
 nshtrainer/configs/metrics/_config/__init__.py,sha256=XDDvDPWULd_vd3lrgF2KGAVR2LVDuhdQvy-fF2ImarI,159
-nshtrainer/configs/nn/__init__.py,sha256=tkFG2Hb0oL_AmWP3_0WkDN2zI5PkVfrgwXhaAII7CZw,2072
+nshtrainer/configs/nn/__init__.py,sha256=Ms2gIqbRxNVm6GHKCddCJTTqMwUPifjjHD_fCfJpcMQ,2174
 nshtrainer/configs/nn/mlp/__init__.py,sha256=O6kQ6utZNJPG9Fax5pRdZcHa3J-XFKKdXcc_PQg0jk0,347
 nshtrainer/configs/nn/nonlinearity/__init__.py,sha256=LCTbTyelCMABVw505CGQ4UpEGlAnIhflSLFqwAQXLQA,2155
+nshtrainer/configs/nn/rng/__init__.py,sha256=4iC6vwxbfNeXyvpwZ1Z5Kcy-he4cu7mg3UpLD-RLrHc,141
 nshtrainer/configs/optimizer/__init__.py,sha256=itIDIHQvGm50eZ7JLyNElahnNUMPJ__4PMmTjc0RQ6o,444
 nshtrainer/configs/profiler/__init__.py,sha256=2ssaIpfVnvcbfNvZ-JeKp1Cx4NO1LknkVqTm1hu7Lvw,768
 nshtrainer/configs/profiler/_base/__init__.py,sha256=ekYfPg-VDhCAFM5nJka2TxUYdRDm1CKqjwUOQNbQjD4,176
@@ -119,11 +120,12 @@ nshtrainer/model/base.py,sha256=bZMNap0rkxRbAbu2BOHV_6YS2iZZnvy6wVSMOXGa_ZM,8680
 nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0dXLvqHw,3154
 nshtrainer/model/mixins/debug.py,sha256=ydLuAAaa7M5bX0gougZ5gWuZnvn4Ra9assal3IZ9hq8,2086
 nshtrainer/model/mixins/logger.py,sha256=7u9fQig-SVFA9RFIB4U0gqJAzruh49mgmXXvZ6VkDUk,11694
-nshtrainer/nn/__init__.py,sha256=5Gg3nieGSC5_dXaI9KUVUUbM13hHexH9831m4hcf6no,1475
+nshtrainer/nn/__init__.py,sha256=Vd246v2N9tBQ8XxmTquWzj5lAmeSnngrjpYOfp4LTXM,1499
 nshtrainer/nn/mlp.py,sha256=nYUgAISzuhC8sav6PloAdyz0PdEoikwppiXIuToEVdE,7550
 nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
 nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
 nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
+nshtrainer/nn/rng.py,sha256=IJGvX9v8qBkfgBrMlNU2aj-MbYTPoncFyJzvPkzCQpM,512
 nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
 nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
 nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
@@ -142,7 +144,7 @@ nshtrainer/trainer/plugin/layer_sync.py,sha256=h-ydZwXepnsw5-paLgiDatqPyQ_8C0QEv
 nshtrainer/trainer/plugin/precision.py,sha256=I0QsB1bVxmsFmBOkgrAfGONsuYae_lD9Bz0PfJEQvH4,5598
 nshtrainer/trainer/signal_connector.py,sha256=GhfGcSzfaTNhnj2QFkBDq5aT7FqbLMA7eC8SYQs8_8w,10828
 nshtrainer/trainer/strategy.py,sha256=VPTn5z3zvXTydY8IJchjhjcOfpvtoejnvUkq5E4WTus,1368
-nshtrainer/trainer/trainer.py,sha256=QEK-0bcw1y5Cconi99PYFXr0MElUGgGYMZ_SlcJUQ1k,20364
+nshtrainer/trainer/trainer.py,sha256=Lo3vUo3ooTAjaX2fUYPFSMv5FP7sWfVov0QbA-T5hZ8,21113
 nshtrainer/util/_environment_info.py,sha256=MT8mBe6ZolRfKiwU-les1P-lPNPqXpHQcfADrh_A3uY,24629
 nshtrainer/util/bf16.py,sha256=9QhHZCkYSfYpIcxwAMoXyuh2yTSHBzT-EdLQB297jEs,762
 nshtrainer/util/config/__init__.py,sha256=Z39JJufSb61Lhn2GfVcv3eFW_eorOrN9-9llDWlnZZM,272
@@ -154,6 +156,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
 nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
 nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
 nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
-nshtrainer-1.0.0b53.dist-info/METADATA,sha256=xzPw_zTWp8gLpzKhvcAHTv8KbZCvxwMU_MIzMiq8j78,988
-nshtrainer-1.0.0b53.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
-nshtrainer-1.0.0b53.dist-info/RECORD,,
+nshtrainer-1.0.0b55.dist-info/METADATA,sha256=3JHfGqw8kR8FKZT-E--W3H_rS5M-QQHH9U2EooKyO70,988
+nshtrainer-1.0.0b55.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+nshtrainer-1.0.0b55.dist-info/RECORD,,