MS-Attn-Simulation/hparams.py at main · FFishy-git/MS-Attn-Simulation · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
hparams = {
    "model": {
        "dx": 50,
        "dy": 1,
        "L": 100,
        "num_heads": 8,
        "use_bias": False,
        "dropout_rate": 0.0,
        "attention_type": "softmax",
        "q_k_v_o_proj_enabled": [True, True, True, True],
    },
    "data": {
        "number_of_samples": 10000,
        "noise_std": 0.0,
    },
    "train": {
        "optimizer": "SGD",
        "weight_decay": 0.01,
        "gradient_clip": 1.0,
        "momentum": 0.0,
        "batch_size": 64,
        "max_steps": 100_000,
        "warmup_steps": 2_000,
        "lr_decay_steps": 100_000,
        "min_lr": 5e-5,
        "use_stored_data": False,
        "use_logger": True,
        "use_mixed_precision": True,
        "num_epochs": 1000000,
        "print_every_n_epochs": 1,
        "visualize_every_n_epochs": 10,
        "initialize_weights": .01, # lambda
        "max_learning_rate": 5e-1,
        "min_learning_rate": 1e-4,
        "learning_rate": 1e-2,
        "grad_elbow": .1,
        "lr_schedule_method": "constant", # "elbow_decay"
    }
}