Merge pull request #100 from chaitjo/patch-1

Fix unweighted EnergyForcesLoss
ACEsuit · May 24, 2023 · 538b03c · 538b03c
2 parents f6d48af + 3655a91
commit 538b03c
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 29 deletions.
diff --git a/mace/modules/__init__.py b/mace/modules/__init__.py
@@ -21,7 +21,6 @@
 )
 from .loss import (
     DipoleSingleLoss,
-    EnergyForcesLoss,
     WeightedEnergyForcesDipoleLoss,
     WeightedEnergyForcesLoss,
     WeightedEnergyForcesStressLoss,
@@ -87,7 +86,6 @@
     "ScaleShiftBOTNet",
     "AtomicDipolesMACE",
     "EnergyDipolesMACE",
-    "EnergyForcesLoss",
     "WeightedEnergyForcesLoss",
     "WeightedForcesLoss",
     "WeightedEnergyForcesVirialsLoss",

diff --git a/mace/modules/loss.py b/mace/modules/loss.py
@@ -77,30 +77,6 @@ def weighted_mean_squared_error_dipole(ref: Batch, pred: TensorDict) -> torch.Te
     # return torch.mean(torch.square((torch.reshape(ref['dipole'], pred["dipole"].shape) - pred['dipole']) / num_atoms))  # []
 
 
-class EnergyForcesLoss(torch.nn.Module):
-    def __init__(self, energy_weight=1.0, forces_weight=1.0) -> None:
-        super().__init__()
-        self.register_buffer(
-            "energy_weight",
-            torch.tensor(energy_weight, dtype=torch.get_default_dtype()),
-        )
-        self.register_buffer(
-            "forces_weight",
-            torch.tensor(forces_weight, dtype=torch.get_default_dtype()),
-        )
-
-    def forward(self, ref: Batch, pred: TensorDict) -> torch.Tensor:
-        return self.energy_weight * mean_squared_error_energy(
-            ref, pred
-        ) + self.forces_weight * mean_squared_error_forces(ref, pred)
-
-    def __repr__(self):
-        return (
-            f"{self.__class__.__name__}(energy_weight={self.energy_weight:.3f}, "
-            f"forces_weight={self.forces_weight:.3f})"
-        )
-
-
 class WeightedEnergyForcesLoss(torch.nn.Module):
     def __init__(self, energy_weight=1.0, forces_weight=1.0) -> None:
         super().__init__()

diff --git a/scripts/run_train.py b/scripts/run_train.py
@@ -189,9 +189,8 @@ def main() -> None:
             dipole_weight=args.dipole_weight,
         )
     else:
-        loss_fn = modules.EnergyForcesLoss(
-            energy_weight=args.energy_weight, forces_weight=args.forces_weight
-        )
+        # Unweighted Energy and Forces loss by default
+        loss_fn = modules.WeightedEnergyForcesLoss(energy_weight=1.0, forces_weight=1.0)
     logging.info(loss_fn)
 
     if args.compute_avg_num_neighbors: