add mae

Author: iProzd

deepmd/pt/loss/ener.py

@@ -187,28 +187,26 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                 )
                 # more_loss['log_keys'].append('rmse_e')
             else:  # use l1 and for all atoms
+                energy_pred = energy_pred * atom_norm
+                energy_label = energy_label * atom_norm
                 l1_ener_loss = F.l1_loss(
                     energy_pred.reshape(-1),
                     energy_label.reshape(-1),
-                    reduction="sum",
+                    reduction="mean",
                 )
                 loss += pref_e * l1_ener_loss
                 more_loss["mae_e"] = self.display_if_exist(
-                    F.l1_loss(
-                        energy_pred.reshape(-1),
-                        energy_label.reshape(-1),
-                        reduction="mean",
-                    ).detach(),
+                    l1_ener_loss.detach(),
                     find_energy,
                 )
                 # more_loss['log_keys'].append('rmse_e')
-            if mae:
-                mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm
-                more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
-                mae_e_all = torch.mean(torch.abs(energy_pred - energy_label))
-                more_loss["mae_e_all"] = self.display_if_exist(
-                    mae_e_all.detach(), find_energy
-                )
+            # if mae:
+            #     mae_e = torch.mean(torch.abs(energy_pred - energy_label)) * atom_norm
+            #     more_loss["mae_e"] = self.display_if_exist(mae_e.detach(), find_energy)
+            #     mae_e_all = torch.mean(torch.abs(energy_pred - energy_label))
+            #     more_loss["mae_e_all"] = self.display_if_exist(
+            #         mae_e_all.detach(), find_energy
+            #     )
 
         if (
             (self.has_f or self.has_pf or self.relative_f or self.has_gf)
@@ -241,17 +239,17 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
                         rmse_f.detach(), find_force
                     )
                 else:
-                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="none")
+                    l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean")
                     more_loss["mae_f"] = self.display_if_exist(
-                        l1_force_loss.mean().detach(), find_force
+                        l1_force_loss.detach(), find_force
                     )
-                    l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
+                    # l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum()
                     loss += (pref_f * l1_force_loss).to(GLOBAL_PT_FLOAT_PRECISION)
-                if mae:
-                    mae_f = torch.mean(torch.abs(diff_f))
-                    more_loss["mae_f"] = self.display_if_exist(
-                        mae_f.detach(), find_force
-                    )
+                # if mae:
+                #     mae_f = torch.mean(torch.abs(diff_f))
+                #     more_loss["mae_f"] = self.display_if_exist(
+                #         mae_f.detach(), find_force
+                #     )
 
             if self.has_pf and "atom_pref" in label:
                 atom_pref = label["atom_pref"]
@@ -297,18 +295,29 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False):
         if self.has_v and "virial" in model_pred and "virial" in label:
             find_virial = label.get("find_virial", 0.0)
             pref_v = pref_v * find_virial
+            virial_label = label["virial"]
+            virial_pred = model_pred["virial"].reshape(-1, 9)
             diff_v = label["virial"] - model_pred["virial"].reshape(-1, 9)
-            l2_virial_loss = torch.mean(torch.square(diff_v))
-            if not self.inference:
-                more_loss["l2_virial_loss"] = self.display_if_exist(
-                    l2_virial_loss.detach(), find_virial
+            if not self.use_l1_all:
+                l2_virial_loss = torch.mean(torch.square(diff_v))
+                if not self.inference:
+                    more_loss["l2_virial_loss"] = self.display_if_exist(
+                        l2_virial_loss.detach(), find_virial
+                    )
+                loss += atom_norm * (pref_v * l2_virial_loss)
+                rmse_v = l2_virial_loss.sqrt() * atom_norm
+                more_loss["rmse_v"] = self.display_if_exist(
+                    rmse_v.detach(), find_virial
+                )
+            else:
+                l1_virial_loss = F.l1_loss(virial_label, virial_pred, reduction="mean")
+                more_loss["mae_v"] = self.display_if_exist(
+                    l1_virial_loss.detach(), find_virial
                 )
-            loss += atom_norm * (pref_v * l2_virial_loss)
-            rmse_v = l2_virial_loss.sqrt() * atom_norm
-            more_loss["rmse_v"] = self.display_if_exist(rmse_v.detach(), find_virial)
-            if mae:
-                mae_v = torch.mean(torch.abs(diff_v)) * atom_norm
-                more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial)
+                loss += (pref_v * l1_virial_loss).to(GLOBAL_PT_FLOAT_PRECISION)
+            # if mae:
+            #     mae_v = torch.mean(torch.abs(diff_v)) * atom_norm
+            #     more_loss["mae_v"] = self.display_if_exist(mae_v.detach(), find_virial)
 
         if self.has_ae and "atom_energy" in model_pred and "atom_ener" in label:
             atom_ener = model_pred["atom_energy"]

deepmd/utils/argcheck.py

@@ -2346,6 +2346,12 @@ def loss_ener():
     doc_relative_f = "If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label."
     doc_enable_atom_ener_coeff = "If true, the energy will be computed as \\sum_i c_i E_i. c_i should be provided by file atom_ener_coeff.npy in each data system, otherwise it's 1."
     return [
+        Argument(
+            "use_l1_all",
+            bool,
+            optional=True,
+            default=False,
+        ),
         Argument(
             "start_pref_e",
             [float, int],