deepmodeling / deepmodeling/deepmd-kit

TF backend loss: exclude mixed_type padding atoms (per-frame normalization)

Open
#5,760 0 comments 0 reactions 0 assignees View on GitHub
bug
Dominant language
Python
Stars
2k
Forks
649
Avg merge
6d 18h
Merged PRs (30d)
15

Description

Follow-up to #5738.

PR #5738 fixed the mixed_type padding-atom dilution in the training loss for the dpmodel / pt / pt_expt backends: every loss term is now normalized per frame so a padded `[3+5]`-atom batch yields the same loss/gradient as processing each frame separately and averaging. The **TensorFlow backend loss was left unchanged** and still normalizes by the padded scalar `natoms` / uses unmasked or cross-frame-pooled means, so `mixed_type` batches remain mis-normalized there.

Scope: apply the same per-frame masked normalization (per-atom masked mean for extensive/atomic terms; per-frame real-atom count for extensive energy/virial/property) to the TF backend loss, with an all-ones-mask no-op guard so non-mixed training stays bit-identical. Add grad-accumulation-invariant tests mirroring `source/tests/*/test_loss_padding.py`.

Non-mixed TF training is unaffected; only `mixed_type` batches change.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.