Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Dev hierarchical multibranch output heads #323

Open
wants to merge 10 commits into
base: Predictive_GFM_2025
Choose a base branch
from
Open

Dev hierarchical multibranch output heads #323

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
f59d4d9
add multiple branches to Base model
pzhanggit Jan 21, 2025
9ed64de
adjust based on Max's new datasets with dataset_name
pzhanggit Jan 27, 2025
c896903
add json for compute grad energy physics-informed force prediction
Jan 29, 2025
7837e07
revert time change
Jan 29, 2025
e50c103
reverse prior change
Jan 30, 2025
5679361
Add check to ensure correct output dim for energy in physics-informed…
Jan 30, 2025
3293994
Implement the check after setting var_config
Jan 30, 2025
81bb0b7
revert change
Jan 30, 2025
a38e9db
Shorten comment
Jan 30, 2025
86c3751
Merge pull request #7 from RylieWeaver/dev-multifidelitytraining-rylie
pzhanggit Jan 31, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
107 changes: 107 additions & 0 deletions examples/multibranch/gfm_multibranch.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
{
"Verbosity": {
"level": 2
},
"NeuralNetwork": {
"Architecture": {
"mpnn_type": "EGNN",
"equivariance": true,
"radius": 5.0,
"max_neighbours": 100000,
"num_gaussians": 50,
"envelope_exponent": 5,
"int_emb_size": 64,
"basis_emb_size": 8,
"out_emb_size": 128,
"num_after_skip": 2,
"num_before_skip": 1,
"num_radial": 6,
"num_spherical": 7,
"num_filters": 126,
"pe_dim": 1,
"global_attn_heads": 8,
"global_attn_engine": "",
"global_attn_type": "",
"hidden_dim": 50,
"num_conv_layers": 3,
"output_heads": {
"graph": [
{
"type": "branch-0",
"architecture": {
"num_sharedlayers": 2,
"dim_sharedlayers": 50,
"num_headlayers": 2,
"dim_headlayers": [
50,
25
]
}
},
{
"type": "branch-1",
"architecture": {
"num_sharedlayers": 1,
"dim_sharedlayers": 20,
"num_headlayers": 1,
"dim_headlayers": [
50
]
}
}
],
"node": [
{
"type": "branch-0",
"architecture": {
"num_headlayers": 2,
"dim_headlayers": [
200,
200
],
"type": "mlp"
}
},
{
"type": "branch-1",
"architecture": {
"num_headlayers": 1,
"dim_headlayers": [
100
],
"type": "mlp"
}
}
]
},
"task_weights": [
1.0,
1.0
]
},
"Variables_of_interest": {
"input_node_features": [0, 1, 2, 3],
"output_names": ["energy", "force"],
"output_index": [0, 2],
"output_dim": [1, 3],
"type": ["graph", "node"]
},
"Training": {
"num_epoch": 250,
"EarlyStopping": true,
"perc_train": 0.9,
"loss_function_type": "mae",
"batch_size": 32,
"continue": 0,
"Optimizer": {
"type": "AdamW",
"learning_rate": 1e-3
}
}
},
"Visualization": {
"plot_init_solution": true,
"plot_hist_solution": false,
"create_plots": true
}
}
82 changes: 82 additions & 0 deletions examples/multibranch/gfm_multibranch_physics-informed.json
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
{
"Verbosity": {
"level": 2
},
"NeuralNetwork": {
"Architecture": {
"mpnn_type": "PNAEq",
"equivariance": true,
"radius": 5.0,
"max_neighbours": 100,
"num_gaussians": 50,
"envelope_exponent": 5,
"int_emb_size": 64,
"basis_emb_size": 8,
"out_emb_size": 128,
"num_after_skip": 2,
"num_before_skip": 1,
"num_radial": 6,
"num_spherical": 7,
"num_filters": 126,
"pe_dim": 1,
"global_attn_heads": 8,
"global_attn_engine": "",
"global_attn_type": "",
"hidden_dim": 50,
"num_conv_layers": 3,
"output_heads": {
"node": [
{
"type": "branch-0",
"architecture": {
"num_headlayers": 2,
"dim_headlayers": [
200,
200
],
"type": "mlp"
}
},
{
"type": "branch-1",
"architecture": {
"num_headlayers": 1,
"dim_headlayers": [
100
],
"type": "mlp"
}
}
]
},
"task_weights": [
1.0
]
},
"Variables_of_interest": {
"input_node_features": [0, 1, 2, 3],
"output_names": ["energy"],
"output_index": [0],
"output_dim": [1],
"type": ["node"]
},
"Training": {
"num_epoch": 250,
"EarlyStopping": true,
"perc_train": 0.9,
"loss_function_type": "mae",
"batch_size": 32,
"continue": 0,
"Optimizer": {
"type": "AdamW",
"learning_rate": 1e-3
},
"compute_grad_energy": true
}
},
"Visualization": {
"plot_init_solution": true,
"plot_hist_solution": false,
"create_plots": true
}
}
Loading