Unified OCP Trainer

.pre-commit-config.yaml

@@ -3,7 +3,6 @@ repos:
     rev: 22.3.0
     hooks:
     - id: black
-      language_version: python3.8
       additional_dependencies: ['click==8.0.4']
 -   repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v2.3.0

configs/goc_oc20_debug.yml

@@ -0,0 +1,129 @@
+trainer: ocp
+
+dataset:
+  train:
+    format: lmdb
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/train/2M
+    key_mapping:
+      y: energy
+      force: forces
+    transforms:
+      normalizer:
+        energy:
+          mean: -0.7554450631141663
+          stdev: 2.887317180633545
+        forces:
+          mean: 0
+          stdev: 2.887317180633545
+  val:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+  test:
+    src: /datasets01/open_catalyst/oc20/082422/struct_to_energy_forces/val/id_30k
+
+logger: tensorboard
+
+loss_functions:
+  - energy:
+      fn: mae
+      coefficient: 1
+  - forces:
+      fn: l2mae
+      coefficient: 100
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
+    forces:
+      - mae
+      - cosine_similarity
+    misc:
+      - energy_forces_within_threshold
+  primary_metric: forces_mae
+
+outputs:
+  energy:
+    shape: 1
+    level: system
+  forces:
+    shape: 3
+    level: atom
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+
+model:
+  name: gemnet_oc
+  num_spherical: 7
+  num_radial: 128
+  num_blocks: 4
+  emb_size_atom: 256
+  emb_size_edge: 512
+  emb_size_trip_in: 64
+  emb_size_trip_out: 64
+  emb_size_quad_in: 32
+  emb_size_quad_out: 32
+  emb_size_aint_in: 64
+  emb_size_aint_out: 64
+  emb_size_rbf: 16
+  emb_size_cbf: 16
+  emb_size_sbf: 32
+  num_before_skip: 2
+  num_after_skip: 2
+  num_concat: 1
+  num_atom: 3
+  num_output_afteratom: 3
+  cutoff: 12.0
+  cutoff_qint: 12.0
+  cutoff_aeaint: 12.0
+  cutoff_aint: 12.0
+  max_neighbors: 30
+  max_neighbors_qint: 8
+  max_neighbors_aeaint: 20
+  max_neighbors_aint: 1000
+  rbf:
+    name: gaussian
+  envelope:
+    name: polynomial
+    exponent: 5
+  cbf:
+    name: spherical_harmonics
+  sbf:
+    name: legendre_outer
+  extensive: True
+  output_init: HeOrthogonal
+  activation: silu
+  scale_file: configs/s2ef/all/gemnet/scaling_factors/gemnet-oc.pt
+
+  regress_forces: True
+  direct_forces: True
+  forces_coupled: False
+
+  quad_interaction: True
+  atom_edge_interaction: True
+  edge_atom_interaction: True
+  atom_interaction: True
+
+  num_atom_emb_layers: 2
+  num_global_out_layers: 2
+  qint_tags: [1, 2]
+  otf_graph: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 5000
+  num_workers: 2
+  lr_initial: 5.e-4
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  patience: 3
+  max_epochs: 80
+  ema_decay: 0.999
+  clip_grad_norm: 10
+  weight_decay: 0

configs/goc_stress_debug.yml

@@ -0,0 +1,162 @@
+trainer: ocp
+
+dataset:
+  train:
+    format: lmdb
+    src: /checkpoint/saro00/mpf_datasets/s2efs/0/train.lmdb
+    key_mapping:
+      y: energy
+      force: forces
+      stress: stress
+    transforms:
+      decompose_tensor:
+          tensor: stress
+          rank: 2
+          decomposition:
+            isotropic_stress:
+              irrep_dim: 0
+            anisotropic_stress:
+              irrep_dim: 2
+      normalizer:
+        energy:
+          mean: -5.9749126
+          stdev: 1.866159
+        forces:
+          mean: 0
+          stdev: 1.866159
+        isotropic_stress:
+          mean: 43.27065
+          stdev: 674.1657344451734
+        anisotropic_stress:
+          stdev: 143.72764771869745
+  val:
+    src: /checkpoint/saro00/mpf_datasets/s2efs/0/val.lmdb
+  test:
+    src: /checkpoint/saro00/mpf_datasets/s2efs/0/val.lmdb
+
+logger: tensorboard
+
+loss_functions:
+  - energy:
+      fn: mae
+      coefficient: 1
+  - forces:
+      fn: l2mae
+      coefficient: 100
+  - isotropic_stress:
+      fn: mae
+  - anisotropic_stress:
+      fn: mae
+
+evaluation_metrics:
+  metrics:
+    energy:
+      - mae
+      - mse
+      - energy_within_threshold
+    forces:
+      - mae
+      - cosine_similarity
+    isotropic_stress:
+      - mae
+    anisotropic_stress:
+      - mae
+    stress:
+      - stress_mae_from_decomposition
+    misc:
+      - energy_forces_within_threshold
+  primary_metric: forces_mae
+
+outputs:
+  energy:
+    shape: 1
+    level: system
+  forces:
+    shape: 3
+    level: atom
+    train_on_free_atoms: True
+    eval_on_free_atoms: True
+
+  stress:
+    level: system
+    decomposition:
+      isotropic_stress:
+        irrep_dim: 0
+      anisotropic_stress:
+        irrep_dim: 2
+
+model:
+  name: gemnet_oc
+  num_spherical: 7
+  num_radial: 128
+  num_blocks: 4
+  emb_size_atom: 256
+  emb_size_edge: 512
+  emb_size_trip_in: 64
+  emb_size_trip_out: 64
+  emb_size_quad_in: 32
+  emb_size_quad_out: 32
+  emb_size_aint_in: 64
+  emb_size_aint_out: 64
+  emb_size_rbf: 16
+  emb_size_cbf: 16
+  emb_size_sbf: 32
+  num_before_skip: 2
+  num_after_skip: 2
+  num_concat: 1
+  num_atom: 3
+  num_output_afteratom: 3
+  cutoff: 12.0
+  cutoff_qint: 12.0
+  cutoff_aeaint: 12.0
+  cutoff_aint: 12.0
+  max_neighbors: 30
+  max_neighbors_qint: 8
+  max_neighbors_aeaint: 20
+  max_neighbors_aint: 1000
+  rbf:
+    name: gaussian
+  envelope:
+    name: polynomial
+    exponent: 5
+  cbf:
+    name: spherical_harmonics
+  sbf:
+    name: legendre_outer
+  extensive: True
+  output_init: HeOrthogonal
+  activation: silu
+  scale_file: configs/s2ef/all/gemnet/scaling_factors/gemnet-oc.pt
+
+  regress_forces: True
+  direct_forces: True
+  forces_coupled: False
+
+  quad_interaction: True
+  atom_edge_interaction: True
+  edge_atom_interaction: True
+  atom_interaction: True
+
+  num_elements: 100
+  num_atom_emb_layers: 2
+  num_global_out_layers: 2
+  qint_tags: [1, 2]
+  otf_graph: True
+
+optim:
+  batch_size: 4
+  eval_batch_size: 4
+  load_balancing: atoms
+  eval_every: 5000
+  num_workers: 2
+  lr_initial: 5.e-4
+  optimizer: AdamW
+  optimizer_params: {"amsgrad": True}
+  scheduler: ReduceLROnPlateau
+  mode: min
+  factor: 0.8
+  patience: 3
+  max_epochs: 80
+  ema_decay: 0.999
+  clip_grad_norm: 10
+  weight_decay: 0