nshtrainer 1.0.0b37__py3-none-any.whl → 1.0.0b40__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
nshtrainer/_directory.py CHANGED
@@ -81,7 +81,7 @@ class DirectoryConfig(C.Config):
81
81
 
82
82
  # Save to nshtrainer/{id}/log/{logger name}
83
83
  log_dir = self.resolve_subdirectory(run_id, "log")
84
- log_dir = log_dir / getattr(logger, "name")
84
+ log_dir = log_dir / logger.resolve_logger_dirname()
85
85
  # ^ NOTE: Logger must have a `name` attribute, as this is
86
86
  # the discriminator for the logger registry
87
87
  log_dir.mkdir(exist_ok=True)
@@ -21,11 +21,8 @@ log = logging.getLogger(__name__)
21
21
  class LastCheckpointCallbackConfig(BaseCheckpointCallbackConfig):
22
22
  name: Literal["last_checkpoint"] = "last_checkpoint"
23
23
 
24
- save_on_time_interval: bool = True
25
- """Whether to save checkpoints based on time interval."""
26
-
27
- interval: timedelta = timedelta(hours=12)
28
- """Time interval between checkpoints when save_on_time_interval is True."""
24
+ save_on_time_interval: timedelta | None = None
25
+ """Save a checkpoint every `save_on_time_interval` seconds. If `None`, this feature is disabled."""
29
26
 
30
27
  @override
31
28
  def create_checkpoint(self, trainer_config, dirpath):
@@ -38,8 +35,6 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
38
35
  super().__init__(config, dirpath)
39
36
  self.start_time = time.time()
40
37
  self.last_checkpoint_time = self.start_time
41
- self.interval_seconds = config.interval.total_seconds()
42
- self.save_on_time_interval = config.save_on_time_interval
43
38
 
44
39
  @override
45
40
  def name(self):
@@ -57,12 +52,18 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
57
52
  def topk_sort_reverse(self):
58
53
  return True
59
54
 
60
- def _should_checkpoint(self) -> bool:
61
- if not self.save_on_time_interval:
55
+ def _local_should_checkpoint(self) -> bool:
56
+ if (interval := self.config.save_on_time_interval) is None:
62
57
  return False
58
+
63
59
  current_time = time.time()
64
60
  elapsed_time = current_time - self.last_checkpoint_time
65
- return elapsed_time >= self.interval_seconds
61
+ return elapsed_time >= interval.total_seconds()
62
+
63
+ def _should_checkpoint(self, trainer: Trainer):
64
+ if self.config.save_on_time_interval is None:
65
+ return False
66
+ return trainer.strategy.broadcast(self._local_should_checkpoint(), src=0)
66
67
 
67
68
  def _format_duration(self, seconds: float) -> str:
68
69
  """Format duration in seconds to a human-readable string."""
@@ -98,7 +99,7 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
98
99
  *args,
99
100
  **kwargs,
100
101
  ):
101
- if not self._should_checkpoint():
102
+ if not self._should_checkpoint(trainer):
102
103
  return
103
104
  self.save_checkpoints(trainer)
104
105
 
@@ -110,5 +111,5 @@ class LastCheckpointCallback(CheckpointBase[LastCheckpointCallbackConfig]):
110
111
  def save_checkpoints(self, trainer):
111
112
  super().save_checkpoints(trainer)
112
113
 
113
- if self.save_on_time_interval:
114
+ if self.config.save_on_time_interval is not None:
114
115
  self.last_checkpoint_time = time.time()
@@ -0,0 +1 @@
1
+ * linguist-generated=true
@@ -30,5 +30,14 @@ class LoggerConfigBase(C.Config, ABC):
30
30
  def __bool__(self):
31
31
  return self.enabled
32
32
 
33
+ def resolve_logger_dirname(self) -> str:
34
+ if not (name := getattr(self, "name", None)):
35
+ raise ValueError(
36
+ "Logger must have a name attribute to resolve the directory name.\n"
37
+ "Otherwise, you must override `resolve_logger_dirname`."
38
+ )
39
+
40
+ return name
41
+
33
42
 
34
43
  logger_registry = C.Registry(LoggerConfigBase, discriminator="name")
nshtrainer/nn/mlp.py CHANGED
@@ -1,5 +1,6 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import contextlib
3
4
  import copy
4
5
  from collections.abc import Callable, Sequence
5
6
  from typing import Literal, Protocol, runtime_checkable
@@ -44,6 +45,9 @@ class MLPConfigDict(TypedDict):
44
45
  residual: bool
45
46
  """Whether to use residual connections between layers."""
46
47
 
48
+ seed: int | None
49
+ """Random seed to use for initialization. If None, the default Torch behavior is used."""
50
+
47
51
 
48
52
  class MLPConfig(C.Config):
49
53
  bias: bool = True
@@ -64,15 +68,20 @@ class MLPConfig(C.Config):
64
68
  residual: bool = False
65
69
  """Whether to use residual connections between layers."""
66
70
 
71
+ seed: int | None = None
72
+ """Random seed to use for initialization. If None, the default Torch behavior is used."""
73
+
67
74
  def to_kwargs(self) -> MLPConfigDict:
68
- return {
75
+ kwargs: MLPConfigDict = {
69
76
  "bias": self.bias,
70
77
  "no_bias_scalar": self.no_bias_scalar,
71
78
  "nonlinearity": self.nonlinearity,
72
79
  "ln": self.ln,
73
80
  "dropout": self.dropout,
74
81
  "residual": self.residual,
82
+ "seed": self.seed,
75
83
  }
84
+ return kwargs
76
85
 
77
86
  def create_module(
78
87
  self,
@@ -108,6 +117,7 @@ def MLP(
108
117
  pre_layers: Sequence[nn.Module] = [],
109
118
  post_layers: Sequence[nn.Module] = [],
110
119
  linear_cls: LinearModuleConstructor = nn.Linear,
120
+ seed: int | None = None,
111
121
  ):
112
122
  """
113
123
  Constructs a multi-layer perceptron (MLP) with the given dimensions and activation function.
@@ -123,52 +133,61 @@ def MLP(
123
133
  residual (bool, optional): Whether to use residual connections between layers. Defaults to False.
124
134
  pre_layers (Sequence[nn.Module], optional): List of layers to insert before the linear layers. Defaults to [].
125
135
  post_layers (Sequence[nn.Module], optional): List of layers to insert after the linear layers. Defaults to [].
136
+ linear_cls (LinearModuleConstructor, optional): Linear module constructor to use. Defaults to nn.Linear.
137
+ seed (int | None, optional): Random seed to use for initialization. If None, the default Torch behavior is used. Defaults to None.
126
138
 
127
139
  Returns:
128
140
  nn.Sequential: The constructed MLP.
129
141
  """
130
142
 
131
- if activation is None:
132
- activation = nonlinearity
133
-
134
- if len(dims) < 2:
135
- raise ValueError("mlp requires at least 2 dimensions")
136
- if ln is True:
137
- ln = "pre"
138
- elif isinstance(ln, str) and ln not in ("pre", "post"):
139
- raise ValueError("ln must be a boolean or 'pre' or 'post'")
140
-
141
- layers: list[nn.Module] = []
142
- if ln == "pre":
143
- layers.append(nn.LayerNorm(dims[0]))
144
-
145
- layers.extend(pre_layers)
146
-
147
- for i in range(len(dims) - 1):
148
- in_features = dims[i]
149
- out_features = dims[i + 1]
150
- bias_ = bias and not (no_bias_scalar and out_features == 1)
151
- layers.append(linear_cls(in_features, out_features, bias=bias_))
152
- if dropout is not None:
153
- layers.append(nn.Dropout(dropout))
154
- if i < len(dims) - 2:
155
- match activation:
156
- case NonlinearityConfigBase():
157
- layers.append(activation.create_module())
158
- case nn.Module():
159
- # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
160
- layers.append(copy.deepcopy(activation))
161
- case Callable():
162
- layers.append(activation())
163
- case _:
164
- raise ValueError(
165
- "Either `nonlinearity` or `activation` must be provided"
166
- )
167
-
168
- layers.extend(post_layers)
169
-
170
- if ln == "post":
171
- layers.append(nn.LayerNorm(dims[-1]))
172
-
173
- cls = ResidualSequential if residual else nn.Sequential
174
- return cls(*layers)
143
+ with contextlib.ExitStack() as stack:
144
+ if seed is not None:
145
+ stack.enter_context(
146
+ torch.random.fork_rng(devices=range(torch.cuda.device_count()))
147
+ )
148
+ torch.manual_seed(seed)
149
+
150
+ if activation is None:
151
+ activation = nonlinearity
152
+
153
+ if len(dims) < 2:
154
+ raise ValueError("mlp requires at least 2 dimensions")
155
+ if ln is True:
156
+ ln = "pre"
157
+ elif isinstance(ln, str) and ln not in ("pre", "post"):
158
+ raise ValueError("ln must be a boolean or 'pre' or 'post'")
159
+
160
+ layers: list[nn.Module] = []
161
+ if ln == "pre":
162
+ layers.append(nn.LayerNorm(dims[0]))
163
+
164
+ layers.extend(pre_layers)
165
+
166
+ for i in range(len(dims) - 1):
167
+ in_features = dims[i]
168
+ out_features = dims[i + 1]
169
+ bias_ = bias and not (no_bias_scalar and out_features == 1)
170
+ layers.append(linear_cls(in_features, out_features, bias=bias_))
171
+ if dropout is not None:
172
+ layers.append(nn.Dropout(dropout))
173
+ if i < len(dims) - 2:
174
+ match activation:
175
+ case NonlinearityConfigBase():
176
+ layers.append(activation.create_module())
177
+ case nn.Module():
178
+ # In this case, we create a deep copy of the module to avoid sharing parameters (if any).
179
+ layers.append(copy.deepcopy(activation))
180
+ case Callable():
181
+ layers.append(activation())
182
+ case _:
183
+ raise ValueError(
184
+ "Either `nonlinearity` or `activation` must be provided"
185
+ )
186
+
187
+ layers.extend(post_layers)
188
+
189
+ if ln == "post":
190
+ layers.append(nn.LayerNorm(dims[-1]))
191
+
192
+ cls = ResidualSequential if residual else nn.Sequential
193
+ return cls(*layers)
@@ -0,0 +1,55 @@
1
+ from __future__ import annotations
2
+
3
+ from typing import cast
4
+
5
+ import pytest
6
+ import torch
7
+
8
+ from nshtrainer.nn.mlp import MLP
9
+
10
+
11
+ def test_mlp_seed_reproducibility():
12
+ """Test that the seed parameter in MLP ensures reproducible weights."""
13
+
14
+ # Test dimensions
15
+ dims = [10, 20, 5]
16
+
17
+ # Create two MLPs with the same seed
18
+ seed1 = 42
19
+ mlp1 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
20
+ mlp2 = MLP(dims, activation=torch.nn.ReLU(), seed=seed1)
21
+
22
+ # Create an MLP with a different seed
23
+ seed2 = 123
24
+ mlp3 = MLP(dims, activation=torch.nn.ReLU(), seed=seed2)
25
+
26
+ # Check first layer weights
27
+ layer1_weights1 = cast(torch.Tensor, mlp1[0].weight)
28
+ layer1_weights2 = cast(torch.Tensor, mlp2[0].weight)
29
+ layer1_weights3 = cast(torch.Tensor, mlp3[0].weight)
30
+
31
+ # Same seed should produce identical weights
32
+ assert torch.allclose(layer1_weights1, layer1_weights2)
33
+
34
+ # Different seeds should produce different weights
35
+ assert not torch.allclose(layer1_weights1, layer1_weights3)
36
+
37
+ # Check second layer weights
38
+ layer2_weights1 = cast(torch.Tensor, mlp1[2].weight)
39
+ layer2_weights2 = cast(torch.Tensor, mlp2[2].weight)
40
+ layer2_weights3 = cast(torch.Tensor, mlp3[2].weight)
41
+
42
+ # Same seed should produce identical weights for all layers
43
+ assert torch.allclose(layer2_weights1, layer2_weights2)
44
+
45
+ # Different seeds should produce different weights for all layers
46
+ assert not torch.allclose(layer2_weights1, layer2_weights3)
47
+
48
+ # Test that not providing a seed gives different results each time
49
+ mlp4 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
50
+ mlp5 = MLP(dims, activation=torch.nn.ReLU(), seed=None)
51
+
52
+ # Without seeds, weights should be different
53
+ assert not torch.allclose(
54
+ cast(torch.Tensor, mlp4[0].weight), cast(torch.Tensor, mlp5[0].weight)
55
+ )
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.1
1
+ Metadata-Version: 2.3
2
2
  Name: nshtrainer
3
- Version: 1.0.0b37
3
+ Version: 1.0.0b40
4
4
  Summary:
5
5
  Author: Nima Shoghi
6
6
  Author-email: nimashoghi@gmail.com
@@ -3,7 +3,7 @@ nshtrainer/__init__.py,sha256=g_moPnfQxSxFZX5NB9ILQQOJrt4RTRuiFt9N0STIpxM,874
3
3
  nshtrainer/_callback.py,sha256=tXQCDzS6CvMTuTY5lQSH5qZs1pXUi-gt9bQdpXMVdEs,12715
4
4
  nshtrainer/_checkpoint/metadata.py,sha256=PHy-54Cg-o3OtCffAqrVv6ZVMU7zhRo_-sZiSEEno1Y,5019
5
5
  nshtrainer/_checkpoint/saver.py,sha256=LOP8jjKF0Dw9x9H-BKrLMWlEp1XTan2DUK0zQUCWw5U,1360
6
- nshtrainer/_directory.py,sha256=xY8Z9POZJw0Uh56yqffZbnNZvdA_tnWCucT31dhwFCM,3183
6
+ nshtrainer/_directory.py,sha256=TJR9ccyuzRlAVfVjGyeQ3E2AFAcz-XbBCxWfiXo2SlY,3191
7
7
  nshtrainer/_experimental/__init__.py,sha256=U4S_2y3zgLZVfMenHRaJFBW8yqh2mUBuI291LGQVOJ8,35
8
8
  nshtrainer/_hf_hub.py,sha256=4OsCbIITnZk_YLyoMrVyZ0SIN04FBxlC0ig2Et8UAdo,14287
9
9
  nshtrainer/callbacks/__init__.py,sha256=4giOYT8A709UOLRtQEt16QbOAFUHCjJ_aLB7ITTwXJI,3577
@@ -12,7 +12,7 @@ nshtrainer/callbacks/base.py,sha256=Alaou1IHAIlMEM7g58d_02ozY2xWlshBN7fsw5Ee21s,
12
12
  nshtrainer/callbacks/checkpoint/__init__.py,sha256=l8tkHc83_mLiU0-wT09SWdRzwpm2ulbkLzcuCmuTwzE,620
13
13
  nshtrainer/callbacks/checkpoint/_base.py,sha256=ZVEUVl5kjCSSe69Q0rMUbKBNNUog0pxBwWkeyuxG2w0,6304
14
14
  nshtrainer/callbacks/checkpoint/best_checkpoint.py,sha256=2CQuhPJ3Fi7lDw7z-J8kXXXuDU8-4HcU48oZxR49apk,2667
15
- nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=MJcNB0biOebx2si2IBFaSUiVOSLSCZTzxB-RcEgO2gY,3482
15
+ nshtrainer/callbacks/checkpoint/last_checkpoint.py,sha256=vn-as3ex7kaTRcKsIurVtM6kUSHYNwHJeYG82j2dMcc,3554
16
16
  nshtrainer/callbacks/checkpoint/on_exception_checkpoint.py,sha256=nljzETqkHwA-4g8mxaeFK5HxA8My0dlIPzIUscSMWyk,3525
17
17
  nshtrainer/callbacks/debug_flag.py,sha256=96fuP0C7C6dSs1GiMeUYzzs0X3Q4Pjt9JVWg3b75fU4,1748
18
18
  nshtrainer/callbacks/directory_setup.py,sha256=wPas_Ren8ANejogmIdKhqqgj4ulxz9AS_8xVIAfRXa0,2565
@@ -30,6 +30,7 @@ nshtrainer/callbacks/shared_parameters.py,sha256=s94jJTAIbDGukYJu6l247QonVOCudGC
30
30
  nshtrainer/callbacks/timer.py,sha256=gDcw_K_ikf0bkVgxQ0cDhvvNvz6GLZVLcatuKfh0ORU,4731
31
31
  nshtrainer/callbacks/wandb_upload_code.py,sha256=shV7UtnXgY2bUlXdVrXiaDs0PNLlIt7TzNJkJPkzvzI,2414
32
32
  nshtrainer/callbacks/wandb_watch.py,sha256=VB14Dy5ZRXQ3di0fPv0K_DFJurLhroLPytnuwQBiJFg,3037
33
+ nshtrainer/configs/.gitattributes,sha256=VeZmarvNEqiRBOHGcllpKm90nL6C8u4tBu7SEm7fj-E,26
33
34
  nshtrainer/configs/__init__.py,sha256=MZfcSKhnjtVObBvVv9lu8L2cFTLINP5zcTQvWnz8jdk,14505
34
35
  nshtrainer/configs/_checkpoint/__init__.py,sha256=6s7Y68StboqscY2G4P_QG443jz5aiym5SjOogIljWLg,342
35
36
  nshtrainer/configs/_checkpoint/metadata/__init__.py,sha256=oOPfYkXTjKgm6pluGsG6V1TPyCEGjsQpHVL-LffSUFQ,290
@@ -101,7 +102,7 @@ nshtrainer/data/datamodule.py,sha256=lSOgH32nysJWa6Y7ba1QyOdUV0DVVdO98qokP8wigjk
101
102
  nshtrainer/data/transform.py,sha256=qd0lIocO59Fk_m90xyOHgFezbymd1mRwly8nbYIfHGc,2263
102
103
  nshtrainer/loggers/__init__.py,sha256=Ddd3JJXVzew_ZpwHA9kGnGmvq4OwhItwghDL5PzNhDc,614
103
104
  nshtrainer/loggers/actsave.py,sha256=wgNrpBB6wQM7qff8iLDb_sQnbiAcYHRmH56pcEJPB3o,1409
104
- nshtrainer/loggers/base.py,sha256=1-HoPmOiyXevQvMLXboiKe-4GOE1V5SvjURohOHakVc,882
105
+ nshtrainer/loggers/base.py,sha256=ON92XbwTSgadQOSyw5PiRRFzyH6uJ-xLtE0nB3cbgPc,1205
105
106
  nshtrainer/loggers/csv.py,sha256=xJ8mSmw4vJwinIfqhF6t2HWmh_1dXEYyLfGuXwL7WHo,1160
106
107
  nshtrainer/loggers/tensorboard.py,sha256=E7iO_fDt9bfH02hBL430bXPLljOo5iGgq2QyPqmx2gQ,2324
107
108
  nshtrainer/loggers/wandb.py,sha256=KZXAUWrrmdX_L8rqej77oUHaM0JxZRM8y9z6JP9PISw,6856
@@ -117,10 +118,11 @@ nshtrainer/model/mixins/callback.py,sha256=0LPgve4VszHbLipid4mpI1qnnmdGS2spivs0d
117
118
  nshtrainer/model/mixins/debug.py,sha256=1LX9KzeFX9JDPs_a6YCdYDZXLhEk_5rBO2aCqlfBy7w,2087
118
119
  nshtrainer/model/mixins/logger.py,sha256=27H99FuLaxc6_dDLG2pid4E_5E0-eLGnc2Ifpt0HYIM,6066
119
120
  nshtrainer/nn/__init__.py,sha256=7KCs-GDOynCXAIdwkgAQacc0p3FHLEION50UtrvgAOc,1463
120
- nshtrainer/nn/mlp.py,sha256=ZbkLyOc08stgIugvu1G5_h66DYtxAFDnboikBaJvvZ8,5988
121
+ nshtrainer/nn/mlp.py,sha256=OatI_pzAqa2awKfvMY4CnznO-LltLT8NHQPJKSNn8IM,6979
121
122
  nshtrainer/nn/module_dict.py,sha256=9plb8aQUx5TUEPhX5jI9u8LrpTeKe7jZAHi8iIqcN8w,2365
122
123
  nshtrainer/nn/module_list.py,sha256=UB43pcwD_3nUke_DyLQt-iXKhWdKM6Zjm84lRC1hPYA,1755
123
124
  nshtrainer/nn/nonlinearity.py,sha256=xmaL4QCRvCxqmaGIOwetJeKK-6IK4m2OV7D3SjxSwJQ,6322
125
+ nshtrainer/nn/tests/test_mlp.py,sha256=xBPiHlBvOCn67EbpzzKL-2FU7ikGxHT3i6CMSp1wk7M,1840
124
126
  nshtrainer/optimizer.py,sha256=u968GRNPUNn3f_9BEY2RBNuJq5O3wJWams3NG0dkrOA,1738
125
127
  nshtrainer/profiler/__init__.py,sha256=RjaNBoVcTFu8lF0dNlFp-2LaPYdonoIbDy2_KhgF0Ek,594
126
128
  nshtrainer/profiler/_base.py,sha256=kFcSVn9gJuMwgDxbfyHh46CmEAIPZjxw3yjPbKgzvwA,950
@@ -151,6 +153,6 @@ nshtrainer/util/seed.py,sha256=diMV8iwBKN7Xxt5pELmui-gyqyT80_CZzomrWhNss0k,316
151
153
  nshtrainer/util/slurm.py,sha256=HflkP5iI_r4UHMyPjw9R4dD5AHsJUpcfJw5PLvGYBRM,1603
152
154
  nshtrainer/util/typed.py,sha256=Xt5fUU6zwLKSTLUdenovnKK0N8qUq89Kddz2_XeykVQ,164
153
155
  nshtrainer/util/typing_utils.py,sha256=MjY-CUX9R5Tzat-BlFnQjwl1PQ_W2yZQoXhkYHlJ_VA,442
154
- nshtrainer-1.0.0b37.dist-info/METADATA,sha256=ObMgpZ_qJLmBAkeRDN7ufTuRSTltiB_LYPFTphNvWks,988
155
- nshtrainer-1.0.0b37.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
156
- nshtrainer-1.0.0b37.dist-info/RECORD,,
156
+ nshtrainer-1.0.0b40.dist-info/METADATA,sha256=m_MxoWKJnesp95-EClArQl90vTm____cgqMdNJ6n-Ng,988
157
+ nshtrainer-1.0.0b40.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
158
+ nshtrainer-1.0.0b40.dist-info/RECORD,,
@@ -1,4 +1,4 @@
1
1
  Wheel-Version: 1.0
2
- Generator: poetry-core 1.9.1
2
+ Generator: poetry-core 2.1.1
3
3
  Root-Is-Purelib: true
4
4
  Tag: py3-none-any