Skip to content

fix(loss): exclude mixed_type padding atoms from the training loss #5738

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 17 commits into from
Jul 12, 2026
Merged

fix(loss): exclude mixed_type padding atoms from the training loss #5738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
17 commits
Select commit Hold shift + click to select a range
f737342
fix(loss): recover the per-atom mask in the pt loss so mixed_type gho…
Jul 4, 2026
7bacaee
fix(loss): per-frame normalize dos/tensor losses to exclude mixed_typ…
Jul 4, 2026
29786a1
fix(loss): per-frame normalize the energy loss to exclude mixed_type …
Jul 4, 2026
82c71fd
test(loss): cover atom_ener Huber invariant and mask the rmse_f displ…
Jul 4, 2026
9e8f441
fix(loss): normalize extensive property loss by per-frame real atom c…
Jul 4, 2026
968ae3c
fix(loss): per-frame normalize the spin energy loss to exclude mixed_…
Jul 4, 2026
423be70
fix(loss): per-frame normalize the spin MAE loss paths for mixed_type…
Jul 4, 2026
4381d74
fix(loss): pass axis as a keyword to xp.sum in the property loss
Jul 5, 2026
3a342d8
fix(loss): per-frame normalize the spin atom-energy loss for mixed_ty…
Jul 5, 2026
4990db2
test(loss): pin cpu device in the property torch-path test
Jul 5, 2026
e7e617f
fix(loss): drop dead None stores flagged by CodeQL
Jul 9, 2026
68aa05c
test(loss): xfail-track force_mag MAE frame-normalization debt
Jul 9, 2026
b143f98
docs(loss): reference tracking issues for mixed_type padding follow-ups
Jul 10, 2026
b578129
ci: re-trigger readthedocs build (prior build reaped by RTD resource …
Jul 11, 2026
24fe235
test(pt): run PropertyLoss padding tests on env.DEVICE
Jul 12, 2026
374e24a
fix(loss): make ener_spin force_mag MAE batch-size independent
Jul 12, 2026
1fff149
Merge upstream/master into fix-loss-mixed-type-padding
Jul 12, 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
50 changes: 26 additions & 24 deletions deepmd/dpmodel/loss/dos.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,15 +130,17 @@ def call(
local_label = xp.reshape(
label_dict["atom_dos"], (-1, natoms, self.numb_dos)
)
diff = xp.reshape(local_pred - local_label, (-1, self.numb_dos))
diff3d = local_pred - local_label # [nf, natoms, numb_dos]
if "mask" in model_dict:
mask = xp.reshape(model_dict["mask"], (-1,))
mask_float = xp.astype(mask, diff.dtype)
diff = diff * mask_float[:, None]
n_valid = xp.sum(mask_float)
l2_local_loss_dos = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos)
# 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)
else:
l2_local_loss_dos = xp.mean(xp.square(diff))
l2_local_loss_dos = xp.mean(xp.square(diff3d))
loss += pref_ados * l2_local_loss_dos
more_loss["rmse_local_dos"] = self.display_if_exist(
xp.sqrt(l2_local_loss_dos), find_local
Expand All @@ -155,15 +157,17 @@ def call(
xp.reshape(label_dict["atom_dos"], (-1, natoms, self.numb_dos)),
axis=-1,
)
diff = xp.reshape(local_pred_cdf - local_label_cdf, (-1, self.numb_dos))
diff3d = local_pred_cdf - local_label_cdf # [nf, natoms, numb_dos]
if "mask" in model_dict:
mask = xp.reshape(model_dict["mask"], (-1,))
mask_float = xp.astype(mask, diff.dtype)
diff = diff * mask_float[:, None]
n_valid = xp.sum(mask_float)
l2_local_loss_cdf = xp.sum(xp.square(diff)) / (n_valid * self.numb_dos)
# 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)
else:
l2_local_loss_cdf = xp.mean(xp.square(diff))
l2_local_loss_cdf = xp.mean(xp.square(diff3d))
loss += pref_acdf * l2_local_loss_cdf
more_loss["rmse_local_cdf"] = self.display_if_exist(
xp.sqrt(l2_local_loss_cdf), find_local
Expand All @@ -175,15 +179,14 @@ def call(
global_pred = xp.reshape(model_dict["dos"], (-1, self.numb_dos))
global_label = xp.reshape(label_dict["dos"], (-1, self.numb_dos))
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.sum(model_dict["mask"], axis=-1, keepdims=True)
l2_global_loss_dos = xp.mean(
xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num)
atom_num = xp.mean(
xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
Comment thread
wanghan-iapcm marked this conversation as resolved.
)
atom_num = xp.mean(xp.astype(atom_num, diff.dtype))
else:
atom_num = natoms
l2_global_loss_dos = xp.mean(xp.square(diff))
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 @@ -199,15 +202,14 @@ def call(
xp.reshape(label_dict["dos"], (-1, self.numb_dos)), axis=-1
)
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.sum(model_dict["mask"], axis=-1, keepdims=True)
l2_global_loss_cdf = xp.mean(
xp.sum(xp.square(diff) * atom_num, axis=0) / xp.sum(atom_num)
atom_num = xp.mean(
xp.astype(xp.sum(model_dict["mask"], axis=-1), diff.dtype)
)
atom_num = xp.mean(xp.astype(atom_num, diff.dtype))
else:
atom_num = natoms
l2_global_loss_cdf = xp.mean(xp.square(diff))
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
Loading
Loading