Skip to content

Commit 8272b0c

Browse files
EltonCNEltonCN
committed
2 parents b9fbb79 + 55e5a60 commit 8272b0c

3 files changed

Lines changed: 96 additions & 70 deletions

File tree

‎dev/toy1d_experiments/configutation_test_stats.ipynb‎

Lines changed: 18 additions & 5 deletions
Large diffs are not rendered by default.

‎experiments/src/world_machine_experiments/toy1d/experiment1_configuration_test.py‎

Lines changed: 74 additions & 63 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030
else:
3131
devices.append('cpu')
3232

33-
max_jobs_per_device = 10
33+
max_jobs_per_device = 12
3434
n_worker = n_device * max_jobs_per_device
3535

3636
print("DEVICES", devices)
@@ -44,18 +44,12 @@
4444
"context_size": 200,
4545
"batch_size": 256,
4646
"n_epoch": n_epoch,
47-
"learning_rate": 5e-4,
48-
"weight_decay": 5e-5,
4947
"accumulation_steps": 1,
5048
"state_dimensions": [0],
51-
"optimizer_class": AdamW,
5249
"device": devices,
5350
"state_control": "periodic",
5451
"discover_state": True,
5552
"sensorial_train_losses": [Dimensions.MEASUREMENT],
56-
"state_size": 128,
57-
"positional_encoder_type": "alibi",
58-
"n_attention_head": 4,
5953
}
6054

6155
n_variation = 0
@@ -78,65 +72,82 @@
7872
recall_n_past_choices = [5, 1, 0]
7973
recall_stride_choices = [3, 1]
8074

81-
for recall_stride in recall_stride_choices:
82-
for recall_n_past in recall_n_past_choices:
83-
84-
recall_n_future_choices = [0]
85-
if len(short_time_recall) > 0 and recall_n_past > 0:
86-
recall_n_future_choices = [5, 1, 0]
87-
elif len(short_time_recall) > 0:
88-
recall_n_future_choices = [5, 1]
89-
90-
for recall_n_future in recall_n_future_choices:
91-
for positional_encoder_type in ["alibi"]:
92-
for block_configuration in [[Dimensions.MEASUREMENT, Dimensions.MEASUREMENT],
93-
[Dimensions.MEASUREMENT, Dimensions.STATE_INPUT]]:
94-
for state_activation in ["tanh", None]:
95-
state_regularizer = None
96-
if state_activation is None:
97-
state_regularizer = "mse"
98-
99-
for train_mse in [True]:
100-
for train_stdw in [False]:
101-
for noise_config in [None,
102-
{"state": {
103-
"mean": 0.0, "std": 0.1}},
104-
{"measurement": {
105-
"mean": 0.0, "std": 0.1}},
106-
{"state": {"mean": 0.0, "std": 0.1}, "measurement": {"mean": 0.0, "std": 0.1}}]:
107-
108-
n_variation += 1
109-
110-
config = {"n_segment": n_segment,
111-
"fast_forward": fast_forward,
112-
"stable_state_epochs": stable_state_epochs,
113-
"check_input_masks": check_input_masks,
114-
"state_save_method": state_save_method,
115-
"mask_sensorial_data": mask_sensorial_data,
116-
"short_time_recall": short_time_recall,
117-
"recall_stride_past": recall_stride,
118-
"recall_stride_future": recall_stride,
119-
"recall_n_past": recall_n_past,
120-
"recall_n_future": recall_n_future,
121-
"positional_encoder_type": positional_encoder_type,
122-
"block_configuration": block_configuration,
123-
"state_activation": state_activation,
124-
"state_regularizer": state_regularizer,
125-
"train_mse": train_mse,
126-
"train_sdtw": train_stdw,
127-
"noise_config": noise_config}
128-
129-
model = make_model(
130-
config, "ParametersModel").model_validate(config)
131-
model_json = model.model_dump_json()
132-
variation_hash = int(
133-
hashlib.md5(model_json.encode('utf-8')).hexdigest(), 16)
134-
135-
configurations["variation"+str(
136-
variation_hash)] = config
75+
for recall_stride_past in recall_stride_choices:
76+
for recall_stride_future in recall_stride_choices:
77+
for recall_n_past in recall_n_past_choices:
78+
79+
recall_n_future_choices = [0]
80+
if len(short_time_recall) > 0 and recall_n_past > 0:
81+
recall_n_future_choices = [
82+
5, 1, 0]
83+
elif len(short_time_recall) > 0:
84+
recall_n_future_choices = [
85+
5, 1]
86+
87+
for recall_n_future in recall_n_future_choices:
88+
for positional_encoder_type in ["alibi"]:
89+
for block_configuration in [[Dimensions.MEASUREMENT, Dimensions.MEASUREMENT],
90+
[Dimensions.MEASUREMENT, Dimensions.STATE_INPUT]]:
91+
for state_activation in ["tanh", None]:
92+
state_regularizer = None
93+
if state_activation is None:
94+
state_regularizer = "mse"
95+
96+
for train_mse in [True]:
97+
for train_stdw in [False]:
98+
for noise_config in [None,
99+
{"state": {
100+
"mean": 0.0, "std": 0.1}},
101+
{"measurement": {
102+
"mean": 0.0, "std": 0.1}},
103+
{"state": {"mean": 0.0, "std": 0.1}, "measurement": {"mean": 0.0, "std": 0.1}}]:
104+
for local_chance in [None, 0.25]:
105+
n_variation += 1
106+
107+
config = {"n_segment": n_segment,
108+
"fast_forward": fast_forward,
109+
"stable_state_epochs": stable_state_epochs,
110+
"check_input_masks": check_input_masks,
111+
"state_save_method": state_save_method,
112+
"mask_sensorial_data": mask_sensorial_data,
113+
"short_time_recall": short_time_recall,
114+
"recall_stride_past": recall_stride_past,
115+
"recall_stride_future": recall_stride_future,
116+
"recall_n_past": recall_n_past,
117+
"recall_n_future": recall_n_future,
118+
"positional_encoder_type": positional_encoder_type,
119+
"block_configuration": block_configuration,
120+
"state_activation": state_activation,
121+
"state_regularizer": state_regularizer,
122+
"train_mse": train_mse,
123+
"train_sdtw": train_stdw,
124+
"noise_config": noise_config,
125+
"local_chance": local_chance,
126+
127+
# Fixed
128+
"batch_size": 256,
129+
"learning_rate": 1e-3,
130+
"cosine_annealing": True,
131+
"cosine_annealing_T_mult": 1,
132+
"cosine_annealing_T0": 25,
133+
"weight_decay": 5e-5,
134+
"optimizer_class": AdamW,
135+
"state_size": 128,
136+
"n_attention_head": 4}
137+
138+
model = make_model(
139+
config, "ParametersModel").model_validate(config)
140+
model_json = model.model_dump_json()
141+
variation_hash = int(
142+
hashlib.md5(model_json.encode('utf-8')).hexdigest(), 16)
143+
144+
configurations["variation"+str(
145+
variation_hash)] = config
137146

138147
assert len(configurations) == n_variation
139148

149+
print(f"Running {n_variation} variations")
150+
140151
os.makedirs(output_dir, exist_ok=True)
141152
configurations_path = os.path.join(output_dir, "configurations.bin")
142153
with open(configurations_path, "wb") as file:

‎src/world_machine/train/stages/state_manager.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,17 +20,19 @@ class StateManager(TrainStage):
2020
def __init__(self,
2121
stable_state_epochs: int = 1,
2222
check_input_masks: bool = False,
23-
state_save_method: StateSaveMethod = StateSaveMethod.REPLACE):
23+
state_save_method: StateSaveMethod = StateSaveMethod.REPLACE,
24+
reset_in_epoch0: bool = False):
2425
super().__init__(1)
2526

2627
self._stable_state_epochs = stable_state_epochs
2728
self._check_input_masks = check_input_masks
2829
self._state_save_method = state_save_method
30+
self._reset_in_epoch0 = reset_in_epoch0
2931

3032
def pre_segment(self, itens: list[TensorDict], losses: dict, batch_size: int,
3133
seq_len: int, epoch_index: int, device: torch.device, state_size: int, mode: DatasetPassMode, model: WorldMachine) -> None:
3234

33-
if epoch_index == 0 and "state" not in itens[0]["inputs"]:
35+
if (epoch_index == 0 and self._reset_in_epoch0) or "state" not in itens[0]["inputs"]:
3436
for item in itens:
3537
seq_len = item["inputs"][next(
3638
iter(item["inputs"].keys()))].shape[1]

0 commit comments

Comments
 (0)