EquiformerV2 + DeNS model and trainer (#880)

* add density metrics * update trainer & loss * interleave atoms in loss * fix call to keys * add rmse to evaluation metrics * fix linting. * per_atom_loss fix * fix test * Equiformer DeNS model and trainer * fix linting. * lint * lint again * add type hints * empty cuda cache and remove db closing * type hints * add missing args to docstring * add return type hints * rename dens heads * move use_densoising to heads * abstract denoising targets * update omat24 dens config * fix imports * fix trainer --------- Co-authored-by: lbluque <[email protected]> Co-authored-by: Luis Barroso-Luque <[email protected]> Co-authored-by: Brandon <[email protected]>
FAIR-Chem · Nov 26, 2024 · e11e78e · e11e78e
1 parent 72614bf
commit e11e78e
Show file tree

Hide file tree

Showing 4 changed files with 1,456 additions and 7 deletions.
diff --git a/configs/omat24/mptrj/eqV2_31M_dens_mptrj.yml b/configs/omat24/mptrj/eqV2_31M_dens_mptrj.yml
@@ -148,17 +148,16 @@ model:
 
     use_force_encoding:                   True
     use_noise_schedule_sigma_encoding:    False
-    use_denoising_energy:                 True
-    use_denoising_stress:                 False
-
 
   heads:
     energy:
-      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSEnergyHead
+      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSScalarHead
+      use_denoising:                True
     forces:
-      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSForceHead
+      module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSVectorHead
     stress:
       module: fairchem.core.models.equiformer_v2.equiformer_v2_dens.DeNSRank2Head
       output_name: stress
       use_source_target_embedding:   True
       decompose:                     True
+      use_denoising:                 False