deepmd-kit/examples/water/dpa4/input.json at c41e68af5968f3e38145ad5cda5b9e77f12d9ef5 · deepmodeling/deepmd-kit · 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
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
{
  "_comment": "DPA4-Mini energy-training example for the water dataset.",
  "model": {
    "type": "dpa4",
    "type_map": [
      "O",
      "H"
    ],
    "descriptor": {
      "rcut": 6.0,
      "channels": 32,
      "n_radial": 16,
      "edge_norm": false,
      "use_env_seed": true,
      "lmax": 2,
      "mmax": 1,
      "n_blocks": 2,
      "mixing_layers": 3,
      "radial_so2_mode": "degree_channel",
      "radial_so2_rank": 1,
      "n_focus": 1,
      "focus_dim": 0,
      "n_atten_head": 1,
      "message_node_so3": true,
      "ffn_neurons": 0,
      "ffn_so3_grid": true,
      "grid_mlp": false,
      "grid_branch": [
        0,
        0,
        1
      ],
      "ffn_blocks": 1,
      "so3_readout": "mlp",
      "use_amp": true,
      "precision": "float32",
      "seed": 42
    },
    "fitting_net": {
      "neuron": [
        0
      ],
      "precision": "float32",
      "seed": 42
    },
    "use_compile": false,
    "enable_tf32": true
  },
  "learning_rate": {
    "type": "wsd",
    "start_lr": 4.5e-4,
    "stop_lr": 1e-6,
    "warmup_ratio": 0.003,
    "warmup_start_factor": 0.2,
    "decay_phase_ratio": 0.65,
    "decay_type": "cosine"
  },
  "loss": {
    "type": "ener",
    "loss_func": "mae",
    "f_use_norm": true,
    "start_pref_e": 20,
    "limit_pref_e": 20,
    "start_pref_f": 20,
    "limit_pref_f": 20,
    "start_pref_v": 5,
    "limit_pref_v": 5
  },
  "optimizer": {
    "type": "HybridMuon",
    "weight_decay": 0.001
  },
  "training": {
    "stat_file": "./dpa4.hdf5",
    "stat_file_mode": "update",
    "training_data": {
      "systems": [
        "../data/data_0",
        "../data/data_1",
        "../data/data_2"
      ],
      "batch_size": 1
    },
    "validation_data": {
      "systems": [
        "../data/data_3"
      ],
      "batch_size": 1,
      "numb_batch": 1
    },
    "numb_steps": 2000000,
    "gradient_max_norm": 5.0,
    "save_freq": 2000,
    "save_dir": "ckpt",
    "max_ckpt_keep": 3,
    "enable_ema": true,
    "ema_decay": 0.999,
    "ema_ckpt_keep": 3,
    "disp_file": "lcurve.out",
    "disp_freq": 1000,
    "disp_avg": true,
    "disp_training": true,
    "time_training": true,
    "tensorboard": false,
    "enable_profiler": false,
    "tensorboard_freq": 1000,
    "tensorboard_log_dir": "tb_log",
    "profiling": false,
    "profiling_file": "timeline.json",
    "zero_stage": 1,
    "seed": 42
  },
  "validating": {
    "compiled_infer": false,
    "tf32_infer": false,
    "amp_infer": false,
    "save_best_dir": "ckpt_best"
  }
}