[BE] Update all configs to use ocp2.0 format (#653)

* switch all main repo configs to ocp2.0 format

* move energy and force coefficients to loss functions

* add force magnitude and remove extra energy metrics from s2ef configs

* address PR comments

---------

Co-authored-by: Muhammed Shuaibi <[email protected]>

configs/is2re/100k/base.yml

@@ -1,18 +1,34 @@
-trainer: energy
+trainer: ocp
 
 dataset:
-  - src: data/is2re/100k/train/data.lmdb
-    normalize_labels: True
-    target_mean: -1.525913953781128
-    target_std: 2.279365062713623
-  - src: data/is2re/all/val_id/data.lmdb
+  train:
+    format: lmdb
+    src: data/is2re/100k/train/data.lmdb
+    key_mapping:
+      y_relaxed: energy
+    transforms:
+      normalizer:
+        energy:
+          mean: -1.525913953781128
+          stdev: 2.279365062713623
+  val:
+    src: data/is2re/all/val_id/data.lmdb
 
 logger: wandb
 
-task:
-  dataset: single_point_lmdb
-  description: "Relaxed state energy prediction from initial structure."
-  type: regression
-  metric: mae
-  labels:
-    - relaxed energy
+outputs:
+  energy:
+    shape: 1
+    level: system
+
+loss_functions:
+  - energy:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
+  primary_metric: energy_mae

configs/is2re/10k/base.yml

@@ -1,18 +1,34 @@
-trainer: energy
+trainer: ocp
 
 dataset:
-  - src: data/is2re/10k/train/data.lmdb
-    normalize_labels: True
-    target_mean: -1.525913953781128
-    target_std: 2.279365062713623
-  - src: data/is2re/all/val_id/data.lmdb
+  train:
+    format: lmdb
+    src: data/is2re/10k/train/data.lmdb
+    key_mapping:
+      y_relaxed: energy
+    transforms:
+      normalizer:
+        energy:
+          mean: -1.525913953781128
+          stdev: 2.279365062713623
+  val:
+    src: data/is2re/all/val_id/data.lmdb
 
 logger: wandb
 
-task:
-  dataset: single_point_lmdb
-  description: "Relaxed state energy prediction from initial structure."
-  type: regression
-  metric: mae
-  labels:
-    - relaxed energy
+outputs:
+  energy:
+    shape: 1
+    level: system
+
+loss_functions:
+  - energy:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
+  primary_metric: energy_mae

configs/is2re/all/base.yml

@@ -1,18 +1,34 @@
-trainer: energy
+trainer: ocp
 
 dataset:
-  - src: data/is2re/all/train/data.lmdb
-    normalize_labels: True
-    target_mean: -1.525913953781128
-    target_std: 2.279365062713623
-  - src: data/is2re/all/val_id/data.lmdb
+  train:
+    format: lmdb
+    src: data/is2re/all/train/data.lmdb
+    key_mapping:
+      y_relaxed: energy
+    transforms:
+      normalizer:
+        energy:
+          mean: -1.525913953781128
+          stdev: 2.279365062713623
+  val:
+    src: data/is2re/all/val_id/data.lmdb
 
 logger: wandb
 
-task:
-  dataset: single_point_lmdb
-  description: "Relaxed state energy prediction from initial structure."
-  type: regression
-  metric: mae
-  labels:
-    - relaxed energy
+outputs:
+  energy:
+    shape: 1
+    level: system
+
+loss_functions:
+  - energy:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
+  primary_metric: energy_mae

configs/is2re/all/painn/painn_h1024_bs8x4.yml

@@ -29,7 +29,5 @@ optim:
   factor: 0.8
   patience: 3
   max_epochs: 80
-  energy_coefficient: 1
   ema_decay: 0.999
   clip_grad_norm: 10
-  loss_energy: mae

configs/oc22/is2re/base.yml

@@ -1,19 +1,29 @@
-trainer: energy
+trainer: ocp
 
 dataset:
   train:
+    format: oc22_lmdb
     src: data/oc22/is2re/train
-    normalize_labels: False
+    key_mapping:
+      y_relaxed: energy
   val:
     src: data/oc22/is2re/val_id
 
 logger: wandb
 
-task:
-  dataset: oc22_lmdb
-  description: "Relaxed state energy prediction from initial structure."
-  type: regression
-  metric: mae
+outputs:
+  energy:
+    shape: 1
+    level: system
+
+loss_functions:
+  - energy:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
   primary_metric: energy_mae
-  labels:
-    - relaxed energy

configs/oc22/is2re/base_joint.yml

@@ -1,22 +1,30 @@
-trainer: energy
-
 dataset:
   train:
+    format: oc22_lmdb
     src: data/oc20+oc22/is2re/train
-    normalize_labels: False
     train_on_oc20_total_energies: True
-    #download at https://dl.fbaipublicfiles.com/opencatalystproject/data/oc22/oc20_ref.pkl
+    # download at https://dl.fbaipublicfiles.com/opencatalystproject/data/oc22/oc20_ref.pkl
     oc20_ref: path/to/oc20_ref.pkl
+    key_mapping:
+      y_relaxed: energy
   val:
     src: data/oc22/is2re/val_id
 
 logger: wandb
 
-task:
-  dataset: oc22_lmdb
-  description: "Relaxed state energy prediction from initial structure."
-  type: regression
-  metric: mae
+outputs:
+  energy:
+    shape: 1
+    level: system
+
+loss_functions:
+  - energy:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
   primary_metric: energy_mae
-  labels:
-    - relaxed energy

configs/oc22/is2re/gemnet-dT/gemnet-dT.yml

@@ -45,7 +45,5 @@ optim:
   factor: 0.8
   patience: 3
   max_epochs: 100
-  energy_coefficient: 1
   ema_decay: 0.999
   clip_grad_norm: 10
-  loss_energy: mae

configs/oc22/is2re/painn/painn.yml

@@ -29,7 +29,5 @@ optim:
   factor: 0.8
   patience: 3
   max_epochs: 80
-  energy_coefficient: 1
   ema_decay: 0.999
   clip_grad_norm: 10
-  loss_energy: mae

configs/oc22/s2ef/base.yml

@@ -1,29 +1,45 @@
-trainer: forces
+trainer: ocp
 
 dataset:
   train:
+    format: oc22_lmdb
     src: data/oc22/s2ef/train
-    normalize_labels: False
+    key_mapping:
+      y: energy
+      force: forces
   val:
     src: data/oc22/s2ef/val_id
 
 logger: wandb
 
-task:
-  dataset: oc22_lmdb
-  description: "Regressing to energies and forces for DFT trajectories from OCP"
-  type: regression
-  metric: mae
-  primary_metric: forces_mae
-  labels:
-    - potential energy
-  grad_input: atomic forces
-  train_on_free_atoms: True
-  eval_on_free_atoms: True
-  prediction_dtype: float32
+outputs:
+  energy:
+    shape: 1
+    level: system
+    prediction_dtype: float32
+  forces:
+    irrep_dim: 1
+    level: atom
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+    prediction_dtype: float32
+
+loss_functions:
+  - energy:
+      fn: mae
+      coefficient: 1
+  - forces:
+      fn: l2mae
+      coefficient: 100
 
-optim:
-  loss_energy: mae
-  loss_force: atomwisel2
-  force_coefficient: 1
-  energy_coefficient: 1
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+    forces:
+      - mae
+      - cosine_similarity
+      - magnitude_error
+    misc:
+      - energy_forces_within_threshold
+  primary_metric: forces_mae

configs/oc22/s2ef/base_joint.yml

@@ -1,26 +1,48 @@
-trainer: forces
+trainer: ocp
 
 dataset:
   train:
+    format: lmdb
     src: data/oc20+oc22/s2ef/train
-    normalize_labels: False
     train_on_oc20_total_energies: True
-    #download at https://dl.fbaipublicfiles.com/opencatalystproject/data/oc22/oc20_ref.pkl
+    # download at https://dl.fbaipublicfiles.com/opencatalystproject/data/oc22/oc20_ref.pkl
     oc20_ref: path/to/oc20_ref.pkl
+    key_mapping:
+      y: energy
+      force: forces
   val:
     src: data/oc22/s2ef/val_id
 
 logger: wandb
 
-task:
-  dataset: oc22_lmdb
-  description: "Regressing to energies and forces for DFT trajectories from OCP"
-  type: regression
-  metric: mae
+outputs:
+  energy:
+    shape: 1
+    level: system
+    prediction_dtype: float32
+  forces:
+    irrep_dim: 1
+    level: atom
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+    prediction_dtype: float32
+
+loss_functions:
+  - energy:
+      fn: mae
+      coefficient: 1
+  - forces:
+      fn: atomwisel2
+      coefficient: 1
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+    forces:
+      - mae
+      - cosine_similarity
+      - magnitude_error
+    misc:
+      - energy_forces_within_threshold
   primary_metric: forces_mae
-  labels:
-    - potential energy
-  grad_input: atomic forces
-  train_on_free_atoms: True
-  eval_on_free_atoms: True
-  prediction_dtype: float32