Skip to content

Commit

Permalink
chore(pt): examples for new dpa2 model
Browse files Browse the repository at this point in the history
  • Loading branch information
@iProzd
iProzd committed Sep 18, 2024
1 parent 2c9be6f commit aff85e2
Show file tree
Hide file tree
Showing 4 changed files with 244 additions and 8 deletions.
24 changes: 17 additions & 7 deletions examples/water/dpa2/input_torch.json → examples/water/dpa2_large/input_torch.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,16 +9,20 @@
"type": "dpa2",
"repinit": {
"tebd_dim": 8,
"rcut": 9.0,
"rcut_smth": 8.0,
"rcut": 6.0,
"rcut_smth": 0.5,
"nsel": 120,
"neuron": [
25,
50,
100
],
"axis_neuron": 12,
"activation_function": "tanh"
"activation_function": "tanh",
"three_body_sel": 40,
"three_body_rcut": 4.0,
"three_body_rcut_smth": 3.5,
"use_three_body": true
},
"repformer": {
"rcut": 4.0,
Expand All @@ -36,10 +40,16 @@
"update_g1_has_conv": true,
"update_g1_has_grrg": true,
"update_g1_has_drrd": true,
"update_g1_has_attn": true,
"update_g2_has_g1g1": true,
"update_g1_has_attn": false,
"update_g2_has_g1g1": false,
"update_g2_has_attn": true,
"attn2_has_gate": true
"update_style": "res_residual",
"update_residual": 0.01,
"update_residual_init": "norm",
"attn2_has_gate": true,
"use_sqrt_nnei": true,
"g1_out_conv": true,
"g1_out_mlp": true
},
"add_tebd_to_repinit_out": false
},
Expand All @@ -58,7 +68,7 @@
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.0002,
"start_lr": 0.001,
"stop_lr": 3.51e-08,
"_comment": "that's all"
},
Expand Down
112 changes: 112 additions & 0 deletions examples/water/dpa2_medium/input_torch.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
{
"_comment": "that's all",
"model": {
"type_map": [
"O",
"H"
],
"descriptor": {
"type": "dpa2",
"repinit": {
"tebd_dim": 8,
"rcut": 6.0,
"rcut_smth": 0.5,
"nsel": 120,
"neuron": [
25,
50,
100
],
"axis_neuron": 12,
"activation_function": "tanh",
"three_body_sel": 40,
"three_body_rcut": 4.0,
"three_body_rcut_smth": 3.5,
"use_three_body": true
},
"repformer": {
"rcut": 4.0,
"rcut_smth": 3.5,
"nsel": 40,
"nlayers": 6,
"g1_dim": 128,
"g2_dim": 32,
"attn2_hidden": 32,
"attn2_nhead": 4,
"attn1_hidden": 128,
"attn1_nhead": 4,
"axis_neuron": 4,
"update_h2": false,
"update_g1_has_conv": true,
"update_g1_has_grrg": true,
"update_g1_has_drrd": true,
"update_g1_has_attn": false,
"update_g2_has_g1g1": false,
"update_g2_has_attn": true,
"update_style": "res_residual",
"update_residual": 0.01,
"update_residual_init": "norm",
"attn2_has_gate": true,
"use_sqrt_nnei": true,
"g1_out_conv": true,
"g1_out_mlp": true
},
"add_tebd_to_repinit_out": false
},
"fitting_net": {
"neuron": [
240,
240,
240
],
"resnet_dt": true,
"seed": 1,
"_comment": " that's all"
},
"_comment": " that's all"
},
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.001,
"stop_lr": 3.51e-08,
"_comment": "that's all"
},
"loss": {
"type": "ener",
"start_pref_e": 0.02,
"limit_pref_e": 1,
"start_pref_f": 1000,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0,
"_comment": " that's all"
},
"training": {
"stat_file": "./dpa2.hdf5",
"training_data": {
"systems": [
"../data/data_0",
"../data/data_1",
"../data/data_2"
],
"batch_size": 1,
"_comment": "that's all"
},
"validation_data": {
"systems": [
"../data/data_3"
],
"batch_size": 1,
"_comment": "that's all"
},
"numb_steps": 1000000,
"warmup_steps": 0,
"gradient_max_norm": 5.0,
"seed": 10,
"disp_file": "lcurve.out",
"disp_freq": 100,
"save_freq": 2000,
"_comment": "that's all"
}
}
112 changes: 112 additions & 0 deletions examples/water/dpa2_small/input_torch.json
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,112 @@
{
"_comment": "that's all",
"model": {
"type_map": [
"O",
"H"
],
"descriptor": {
"type": "dpa2",
"repinit": {
"tebd_dim": 8,
"rcut": 6.0,
"rcut_smth": 0.5,
"nsel": 120,
"neuron": [
25,
50,
100
],
"axis_neuron": 12,
"activation_function": "tanh",
"three_body_sel": 40,
"three_body_rcut": 4.0,
"three_body_rcut_smth": 3.5,
"use_three_body": true
},
"repformer": {
"rcut": 4.0,
"rcut_smth": 3.5,
"nsel": 40,
"nlayers": 3,
"g1_dim": 128,
"g2_dim": 32,
"attn2_hidden": 32,
"attn2_nhead": 4,
"attn1_hidden": 128,
"attn1_nhead": 4,
"axis_neuron": 4,
"update_h2": false,
"update_g1_has_conv": true,
"update_g1_has_grrg": true,
"update_g1_has_drrd": true,
"update_g1_has_attn": false,
"update_g2_has_g1g1": false,
"update_g2_has_attn": false,
"update_style": "res_residual",
"update_residual": 0.01,
"update_residual_init": "norm",
"attn2_has_gate": true,
"use_sqrt_nnei": true,
"g1_out_conv": true,
"g1_out_mlp": true
},
"add_tebd_to_repinit_out": false
},
"fitting_net": {
"neuron": [
240,
240,
240
],
"resnet_dt": true,
"seed": 1,
"_comment": " that's all"
},
"_comment": " that's all"
},
"learning_rate": {
"type": "exp",
"decay_steps": 5000,
"start_lr": 0.001,
"stop_lr": 3.51e-08,
"_comment": "that's all"
},
"loss": {
"type": "ener",
"start_pref_e": 0.02,
"limit_pref_e": 1,
"start_pref_f": 1000,
"limit_pref_f": 1,
"start_pref_v": 0,
"limit_pref_v": 0,
"_comment": " that's all"
},
"training": {
"stat_file": "./dpa2.hdf5",
"training_data": {
"systems": [
"../data/data_0",
"../data/data_1",
"../data/data_2"
],
"batch_size": 1,
"_comment": "that's all"
},
"validation_data": {
"systems": [
"../data/data_3"
],
"batch_size": 1,
"_comment": "that's all"
},
"numb_steps": 1000000,
"warmup_steps": 0,
"gradient_max_norm": 5.0,
"seed": 10,
"disp_file": "lcurve.out",
"disp_freq": 100,
"save_freq": 2000,
"_comment": "that's all"
}
}
4 changes: 3 additions & 1 deletion source/tests/common/test_examples.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,9 @@
p_examples / "dprc" / "generalized_force" / "input.json",
p_examples / "water" / "se_e2_a" / "input_torch.json",
p_examples / "water" / "se_atten" / "input_torch.json",
p_examples / "water" / "dpa2" / "input_torch.json",
p_examples / "water" / "dpa2_small" / "input_torch.json",
p_examples / "water" / "dpa2_medium" / "input_torch.json",
p_examples / "water" / "dpa2_large" / "input_torch.json",
p_examples / "property" / "train" / "input_torch.json",
p_examples / "water" / "se_e3_tebd" / "input_torch.json",
)
Expand Down

0 comments on commit aff85e2

Please sign in to comment.