Skip to content

Commit

Permalink
add example hydra config (#799)
Browse files Browse the repository at this point in the history
* add example hydra config

* update ocp_hydra config

* remove old hydra example

* update ocp_hydra config
  • Loading branch information
misko authored Aug 16, 2024
1 parent 7877671 commit 96cf75d
Showing 1 changed file with 106 additions and 0 deletions.
106 changes: 106 additions & 0 deletions configs/ocp_hydra_example.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
trainer: equiformerv2_forces

dataset:
train:
format: lmdb
src: data/s2ef/train/
key_mapping:
y: energy
force: forces
transforms:
normalizer:
energy:
mean: -0.7554450631141663
stdev: 2.887317180633545
forces:
mean: 0
stdev: 2.887317180633545
val:
src: data/s2ef/val_id/

logger: wandb

outputs:
energy:
shape: 1
level: system
forces:
irrep_dim: 1
level: atom
train_on_free_atoms: True
eval_on_free_atoms: True

loss_functions:
- energy:
fn: mae
coefficient: 4
- forces:
fn: l2mae
coefficient: 100

evaluation_metrics:
metrics:
energy:
- mae
forces:
- mae
- cosine_similarity
- magnitude_error
misc:
- energy_forces_within_threshold
primary_metric: forces_mae

hide_eval_progressbar: False


model:
#The model definition here is using "hydra"
# The hydra model is really a container for a backbone model
# and a variable number of head modules
name: hydra
# Use a lightweight (4 layer) eSCN backbone
backbone:
model: escn_backbone
num_layers: 4
max_neighbors: 20
cutoff: 12.0
sphere_channels: 128
hidden_channels: 256
lmax_list: [6]
mmax_list: [2]
num_sphere_samples: 128
distance_function: "gaussian"
regress_forces: True
use_pbc: True
basis_width_scalar: 2.0
otf_graph: True
# Use a energy and direct forces head attached onto the eSCN backbone
# This will output both energy and direct forces for each input system
heads:
energy:
module: escn_energy_head
forces:
module: escn_force_head

optim:
batch_size: 8 # 6
eval_batch_size: 12 # 6
load_balancing: atoms
num_workers: 8
lr_initial: 0.0004 # [0.0002, 0.0004], eSCN uses 0.0008 for batch size 96

optimizer: AdamW
optimizer_params:
weight_decay: 0.001
scheduler: LambdaLR
scheduler_params:
lambda_type: cosine
warmup_factor: 0.2
warmup_epochs: 0.01
lr_min_factor: 0.01 #

max_epochs: 3
clip_grad_norm: 100
ema_decay: 0.999

eval_every: 10000

0 comments on commit 96cf75d

Please sign in to comment.