diff --git a/configs/ocp_hydra_example.yml b/configs/ocp_hydra_example.yml new file mode 100755 index 000000000..dbcadeff3 --- /dev/null +++ b/configs/ocp_hydra_example.yml @@ -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