Skip to content

refactor(loss): extract masked per-frame reduction idioms to cut nested branching #5783

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
wanghan-iapcm merged 11 commits into from
Jul 16, 2026
Merged

refactor(loss): extract masked per-frame reduction idioms to cut nested branching #5783

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
3c71eb8
refactor(loss): add shared masked per-frame reduction helpers
Jul 13, 2026
9fd5cd8
refactor(loss): adopt masked reduction helpers in dpmodel EnergyLoss
Jul 13, 2026
9a80a1c
refactor(loss): adopt idiom-2 helper for dpmodel EnergyLoss virial
Jul 13, 2026
d38162b
refactor(loss): adopt masked reduction helpers in dpmodel EnergySpinLoss
Jul 13, 2026
56c4dc5
refactor(loss): adopt masked reduction helpers in dpmodel DOS/Tensor …
Jul 13, 2026
f327683
refactor(loss): adopt shared masked reduction helpers in pt EnergyStd…
Jul 13, 2026
fb1045a
refactor(loss): adopt shared masked reduction helpers in pt EnergySpi…
Jul 13, 2026
cd47027
refactor(loss): adopt shared masked reduction helpers in pt DOS/Tenso…
Jul 13, 2026
13b74c9
refactor(loss): tidy reduction helpers (derive nf internally, docstri…
Jul 13, 2026
bc7fc31
fix(loss): guard zero-real-atom frames in masked_atom_mean
Jul 16, 2026
97c26f5
test(loss): cover zero-real-atom frames in masked_atom_mean
Jul 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 12 additions & 22 deletions deepmd/dpmodel/loss/dos.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,10 @@
from deepmd.dpmodel.loss.loss import (
Loss,
)
from deepmd.dpmodel.loss.reduction import (
masked_atom_mean,
masked_atom_num,
)
from deepmd.utils.data import (
DataRequirementItem,
)
Expand Down Expand Up @@ -134,11 +138,9 @@ def call(
if "mask" in model_dict:
# idiom 1: per-frame masked mean, then average over frames
maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms]
nf = diff3d.shape[0]
sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1))
per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf]
per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf]
l2_local_loss_dos = xp.mean(per_frame_sum / per_frame_dof)
l2_local_loss_dos = masked_atom_mean(
xp.square(diff3d), maskf, self.numb_dos
)
else:
l2_local_loss_dos = xp.mean(xp.square(diff3d))
loss += pref_ados * l2_local_loss_dos
Expand All @@ -161,11 +163,9 @@ def call(
if "mask" in model_dict:
# idiom 1: per-frame masked mean, then average over frames
maskf = xp.astype(model_dict["mask"], diff3d.dtype) # [nf, natoms]
nf = diff3d.shape[0]
sq = xp.square(diff3d) * xp.reshape(maskf, (nf, natoms, 1))
per_frame_sum = xp.sum(xp.reshape(sq, (nf, -1)), axis=-1) # [nf]
per_frame_dof = xp.sum(maskf, axis=-1) * self.numb_dos # [nf]
l2_local_loss_cdf = xp.mean(per_frame_sum / per_frame_dof)
l2_local_loss_cdf = masked_atom_mean(
xp.square(diff3d), maskf, self.numb_dos
)
else:
l2_local_loss_cdf = xp.mean(xp.square(diff3d))
loss += pref_acdf * l2_local_loss_cdf
Expand All @@ -181,12 +181,7 @@ def call(
diff = global_pred - global_label
# idiom 3: global dos is already padding-invariant; plain mean suffices
l2_global_loss_dos = xp.mean(xp.square(diff))
if "mask" in model_dict:
atom_num = xp.mean(
xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
)
else:
atom_num = natoms
atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype)
loss += pref_dos * l2_global_loss_dos
more_loss["rmse_global_dos"] = self.display_if_exist(
xp.sqrt(l2_global_loss_dos) / atom_num, find_global
Expand All @@ -204,12 +199,7 @@ def call(
diff = global_pred_cdf - global_label_cdf
# idiom 3: global cdf is already padding-invariant; plain mean suffices
l2_global_loss_cdf = xp.mean(xp.square(diff))
if "mask" in model_dict:
atom_num = xp.mean(
xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
)
else:
atom_num = natoms
atom_num = masked_atom_num(model_dict.get("mask"), natoms, diff.dtype)
loss += pref_cdf * l2_global_loss_cdf
more_loss["rmse_global_cdf"] = self.display_if_exist(
xp.sqrt(l2_global_loss_cdf) / atom_num, find_global
Expand Down
69 changes: 25 additions & 44 deletions deepmd/dpmodel/loss/ener.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,10 @@
from deepmd.dpmodel.loss.loss import (
Loss,
)
from deepmd.dpmodel.loss.reduction import (
masked_atom_mean,
per_frame_component_mean,
)
from deepmd.utils.data import (
DataRequirementItem,
)
Expand Down Expand Up @@ -296,7 +300,7 @@ def call(
if maskf is not None:
# Idiom 2 (extensive): per-frame normalization by real-atom count.
se = xp.square(energy - energy_hat) # [nf, k]
per_frame = xp.mean(xp.reshape(se, (_nf, -1)), axis=-1) # [nf]
per_frame = per_frame_component_mean(se) # [nf]
if not self.use_huber:
loss += pref_e * xp.mean(per_frame * inv**norm_exp)
else:
Expand Down Expand Up @@ -327,9 +331,7 @@ def call(
l1_ener_loss = xp.mean(xp.abs(energy - energy_hat))
if maskf is not None:
abs_e = xp.abs(energy - energy_hat) # [nf, k]
per_frame_ae = xp.mean(
xp.reshape(abs_e, (_nf, -1)), axis=-1
) # [nf]
per_frame_ae = per_frame_component_mean(abs_e) # [nf]
l1_ener_masked = xp.mean(per_frame_ae * inv)
loss += pref_e * l1_ener_masked
more_loss["mae_e"] = self.display_if_exist(
Expand All @@ -346,8 +348,7 @@ def call(
)
if mae:
if maskf is not None:
abs_e = xp.abs(energy - energy_hat)
per_frame_ae = xp.mean(xp.reshape(abs_e, (_nf, -1)), axis=-1)
per_frame_ae = per_frame_component_mean(xp.abs(energy - energy_hat))
mae_e = xp.mean(per_frame_ae * inv)
else:
mae_e = xp.mean(xp.abs(energy - energy_hat)) * atom_norm_ener
Expand All @@ -362,10 +363,7 @@ def call(
diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3)) # [nf, nloc, 3]
maskf_col = xp.reshape(maskf, (_nf, _nloc, 1)) # [nf, nloc, 1]
# Masked MSE computed for rmse_f display regardless of use_huber.
sq_f = xp.square(diff_f_3d) * maskf_col # [nf, nloc, 3]
_pfs = xp.sum(xp.reshape(sq_f, (_nf, -1)), axis=-1) # [nf]
_pfd = xp.sum(maskf, axis=-1) * 3 # [nf]
l2_force_masked = xp.mean(_pfs / _pfd)
l2_force_masked = masked_atom_mean(xp.square(diff_f_3d), maskf, 3)
if not self.use_huber:
loss += pref_f * l2_force_masked
else:
Expand Down Expand Up @@ -435,12 +433,8 @@ def call(
elif self.loss_func == "mae":
if maskf is not None:
diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
if not self.f_use_norm:
abs_f = xp.abs(diff_f_3d) * maskf_col # [nf, nloc, 3]
per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1)
per_frame_dof = xp.sum(maskf, axis=-1) * 3
l1_force_masked = xp.mean(per_frame_sum / per_frame_dof)
l1_force_masked = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
else:
diff_3 = xp.reshape(force_hat - force, (_nf, _nloc, 3))
norm_2d = xp.reshape(
Expand Down Expand Up @@ -474,11 +468,7 @@ def call(
if mae:
if maskf is not None:
diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
abs_f = xp.abs(diff_f_3d) * maskf_col
per_frame_sum = xp.sum(xp.reshape(abs_f, (_nf, -1)), axis=-1)
per_frame_dof = xp.sum(maskf, axis=-1) * 3
mae_f = xp.mean(per_frame_sum / per_frame_dof)
mae_f = masked_atom_mean(xp.abs(diff_f_3d), maskf, 3)
else:
mae_f = xp.mean(xp.abs(diff_f))
more_loss["mae_f"] = self.display_if_exist(mae_f, find_force)
Expand All @@ -494,7 +484,7 @@ def call(
v2d = xp.reshape(virial, (_nf, 9))
v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
se_v = xp.square(v_hat_2d - v2d) # [nf, 9]
per_frame_v = xp.mean(se_v, axis=-1) # [nf]
per_frame_v = per_frame_component_mean(se_v) # [nf]
if not self.use_huber:
loss += pref_v * xp.mean(per_frame_v * inv**norm_exp)
else:
Expand Down Expand Up @@ -526,8 +516,9 @@ def call(
if maskf is not None:
v2d = xp.reshape(virial, (_nf, 9))
v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
abs_v = xp.abs(v_hat_2d - v2d) # [nf, 9]
per_frame_v = xp.mean(abs_v, axis=-1) # [nf]
per_frame_v = per_frame_component_mean(
xp.abs(v_hat_2d - v2d)
) # [nf]
l1_virial_masked = xp.mean(per_frame_v * inv)
loss += pref_v * l1_virial_masked
more_loss["mae_v"] = self.display_if_exist(
Expand All @@ -546,8 +537,7 @@ def call(
if maskf is not None:
v2d = xp.reshape(virial, (_nf, 9))
v_hat_2d = xp.reshape(virial_hat, (_nf, 9))
abs_v = xp.abs(v_hat_2d - v2d)
per_frame_v = xp.mean(abs_v, axis=-1)
per_frame_v = per_frame_component_mean(xp.abs(v_hat_2d - v2d))
mae_v = xp.mean(per_frame_v * inv)
else:
mae_v = (
Expand All @@ -565,10 +555,10 @@ def call(
# Idiom 1 (per-atom masked mean, ncomp=1).
ae_2d = xp.reshape(atom_ener, (_nf, _nloc))
ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc))
sq_ae = xp.square(ae_hat_2d - ae_2d) * maskf # [nf, nloc]
per_frame_sum = xp.sum(sq_ae, axis=-1) # [nf]
per_frame_dof = xp.sum(maskf, axis=-1) # [nf]
l2_ae_masked = xp.mean(per_frame_sum / per_frame_dof)
l2_ae_masked = masked_atom_mean(
xp.square(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
)
if not self.use_huber:
loss += pref_ae * l2_ae_masked
else:
Expand Down Expand Up @@ -609,10 +599,9 @@ def call(
if maskf is not None:
ae_2d = xp.reshape(atom_ener, (_nf, _nloc))
ae_hat_2d = xp.reshape(atom_ener_hat, (_nf, _nloc))
abs_ae = xp.abs(ae_hat_2d - ae_2d) * maskf # [nf, nloc]
per_frame_sum = xp.sum(abs_ae, axis=-1) # [nf]
per_frame_dof = xp.sum(maskf, axis=-1) # [nf]
l1_ae_masked = xp.mean(per_frame_sum / per_frame_dof)
l1_ae_masked = masked_atom_mean(
xp.abs(ae_hat_2d - ae_2d)[:, :, None], maskf, 1
)
loss += pref_ae * l1_ae_masked
more_loss["mae_ae"] = self.display_if_exist(
l1_ae_masked, find_atom_ener
Expand All @@ -637,13 +626,9 @@ def call(
# Idiom 1 with pref weight (ncomp=3).
diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3))
maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
sq_pf = xp.square(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3]
per_frame_sum = xp.sum(
xp.reshape(sq_pf, (_nf, -1)), axis=-1
) # [nf]
per_frame_dof = xp.sum(maskf, axis=-1) * 3 # [nf]
l2_pf_masked = xp.mean(per_frame_sum / per_frame_dof)
l2_pf_masked = masked_atom_mean(
xp.square(diff_f_3d) * pf_3d, maskf, 3
)
loss += pref_pf * l2_pf_masked
more_loss["rmse_pf"] = self.display_if_exist(
xp.sqrt(l2_pf_masked), find_atom_pref
Expand All @@ -660,11 +645,7 @@ def call(
if maskf is not None:
diff_f_3d = xp.reshape(diff_f, (_nf, _nloc, 3))
pf_3d = xp.reshape(atom_pref, (_nf, _nloc, 3))
maskf_col = xp.reshape(maskf, (_nf, _nloc, 1))
abs_pf = xp.abs(diff_f_3d) * pf_3d * maskf_col # [nf, nloc, 3]
per_frame_sum = xp.sum(xp.reshape(abs_pf, (_nf, -1)), axis=-1)
per_frame_dof = xp.sum(maskf, axis=-1) * 3
l1_pf_masked = xp.mean(per_frame_sum / per_frame_dof)
l1_pf_masked = masked_atom_mean(xp.abs(diff_f_3d) * pf_3d, maskf, 3)
loss += pref_pf * l1_pf_masked
more_loss["mae_pf"] = self.display_if_exist(
l1_pf_masked, find_atom_pref
Expand Down
Loading
Loading