Skip to content

Commit e7ad8dc

Browse files
authored
chore(pt): update multitask example (#4419)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Updated multi-task model configuration with a new descriptor for enhanced representation learning. - Introduced additional parameters for model initialization and attention mechanisms. - **Bug Fixes** - Replaced outdated descriptor references in model configurations to ensure compatibility with new settings. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent ad30709 commit e7ad8dc

File tree

1 file changed

+50
-19
lines changed

1 file changed

+50
-19
lines changed

examples/water_multi_task/pytorch_example/input_torch.json

Lines changed: 50 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -6,31 +6,62 @@
66
"O",
77
"H"
88
],
9-
"sea_descriptor_1": {
10-
"type": "se_e2_a",
11-
"sel": [
12-
46,
13-
92
14-
],
15-
"rcut_smth": 0.50,
16-
"rcut": 6.00,
17-
"neuron": [
18-
25,
19-
50,
20-
100
21-
],
22-
"resnet_dt": false,
23-
"axis_neuron": 16,
24-
"type_one_side": true,
25-
"seed": 1,
9+
"dpa2_descriptor": {
10+
"type": "dpa2",
11+
"repinit": {
12+
"tebd_dim": 8,
13+
"rcut": 6.0,
14+
"rcut_smth": 0.5,
15+
"nsel": 120,
16+
"neuron": [
17+
25,
18+
50,
19+
100
20+
],
21+
"axis_neuron": 12,
22+
"activation_function": "tanh",
23+
"three_body_sel": 48,
24+
"three_body_rcut": 4.0,
25+
"three_body_rcut_smth": 3.5,
26+
"use_three_body": true
27+
},
28+
"repformer": {
29+
"rcut": 4.0,
30+
"rcut_smth": 3.5,
31+
"nsel": 48,
32+
"nlayers": 6,
33+
"g1_dim": 128,
34+
"g2_dim": 32,
35+
"attn2_hidden": 32,
36+
"attn2_nhead": 4,
37+
"attn1_hidden": 128,
38+
"attn1_nhead": 4,
39+
"axis_neuron": 4,
40+
"update_h2": false,
41+
"update_g1_has_conv": true,
42+
"update_g1_has_grrg": true,
43+
"update_g1_has_drrd": true,
44+
"update_g1_has_attn": false,
45+
"update_g2_has_g1g1": false,
46+
"update_g2_has_attn": true,
47+
"update_style": "res_residual",
48+
"update_residual": 0.01,
49+
"update_residual_init": "norm",
50+
"attn2_has_gate": true,
51+
"use_sqrt_nnei": true,
52+
"g1_out_conv": true,
53+
"g1_out_mlp": true
54+
},
55+
"precision": "float64",
56+
"add_tebd_to_repinit_out": false,
2657
"_comment": " that's all"
2758
},
2859
"_comment": "that's all"
2960
},
3061
"model_dict": {
3162
"water_1": {
3263
"type_map": "type_map_all",
33-
"descriptor": "sea_descriptor_1",
64+
"descriptor": "dpa2_descriptor",
3465
"fitting_net": {
3566
"neuron": [
3667
240,
@@ -44,7 +75,7 @@
4475
},
4576
"water_2": {
4677
"type_map": "type_map_all",
47-
"descriptor": "sea_descriptor_1",
78+
"descriptor": "dpa2_descriptor",
4879
"fitting_net": {
4980
"neuron": [
5081
240,

0 commit comments

Comments
 (0)