fix finetune RMSE and memory issue (#2860)

Fix #2472. The previous implementation tried to allocate a (N, N) array and the RMSE result was actually MAE instead. --------- Signed-off-by: Jinzhe Zeng <[email protected]> Signed-off-by: Han Wang <[email protected]> Co-authored-by: Han Wang <[email protected]> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
deepmodeling · Sep 25, 2023 · 67f30e0 · 67f30e0
1 parent ffe10f9
commit 67f30e0
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py
@@ -856,9 +856,13 @@ def change_energy_bias(
             delta_bias = np.linalg.lstsq(type_numbs, bias_diff, rcond=None)[0]
             unbias_e = energy_predict + type_numbs @ delta_bias
             atom_numbs = type_numbs.sum(-1)
-            rmse_ae = (
-                np.sqrt(np.square(unbias_e - energy_ground_truth)) / atom_numbs
-            ).mean()
+            rmse_ae = np.sqrt(
+                np.mean(
+                    np.square(
+                        (unbias_e.ravel() - energy_ground_truth.ravel()) / atom_numbs
+                    )
+                )
+            )
             self.bias_atom_e[idx_type_map] += delta_bias.reshape(-1)
             log.info(
                 f"RMSE of atomic energy after linear regression is: {rmse_ae} eV/atom."