feat(pt): Support atomic property fitting #4642
Conversation
for more information, see https://pre-commit.ci
📝 WalkthroughWalkthroughThis pull request updates the loss computation workflow for property predictions. The changes modify the Changes
Sequence Diagram(s)sequenceDiagram
participant T as Training Module
participant L as PropertyLoss
participant A as Argcheck Utility
T->>L: Instantiate PropertyLoss with new parameters
L->>L: Validate parameters & initialize scaling factors
T->>A: Retrieve loss configuration parameters
L->>L: Compute loss in forward() using prefactors
L->>T: Return loss value and updated metrics
Possibly related PRs
Suggested labels
Suggested reviewers
✨ Finishing Touches
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
deepmd/utils/argcheck.py (1)
2745-2798: Enhance the clarity of usage and defaults for the new property parameters.The documentation and defaults for
start_pref_property,limit_pref_property,start_pref_aproperty, andlimit_pref_apropertylook good. However, consider clarifying in the docstrings how the behavior changes if a user sets one of these parameters to zero (e.g., disabling partial training on that axis or phasing it out over time). Providing explicit examples can help users avoid unintentional zero weighting.deepmd/pt/loss/property.py (2)
26-32: Guard against zero or negative starter_learning_rate.By design, the code uses
learning_rate / self.starter_learning_rateto compute a scaling coefficient. Consider adding an assertion thatstarter_learning_rate > 0.0to prevent potential division by zero or negative weighting scenarios.
79-85: Fix typographical error and validate edge cases.
- There is a small spelling error in the assertion message: "Can not assian zero weight..." → "Can not assign zero weight...".
- Consider clarifying whether zero is allowed if a user deliberately wants to disable both property and aproperty after some schedule, or strictly forbid it under all circumstances.
- "Can not assian zero weight both to `pref` and `pref_atomic`" + "Can not assign zero weight both to `pref` and `pref_atomic`"
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (3)
deepmd/pt/loss/property.py(5 hunks)deepmd/pt/train/training.py(1 hunks)deepmd/utils/argcheck.py(2 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (16)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test C++ (false)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Analyze (python)
🔇 Additional comments (8)
deepmd/pt/train/training.py (1)
1269-1269: Properly added learning rate for property lossThe change provides consistent handling of the learning rate parameter for the property loss type, aligning it with the implementation of other loss types in the function. This is an essential component for supporting atomic property fitting.
deepmd/pt/loss/property.py (7)
65-69: Looks correct.Assigning the new instance variables is straightforward and consistent with the incoming constructor parameters.
86-91: Verify partial usage scenarios for property vs. aproperty.The assertion ensures that at least one of property or atomic property is enabled. However, if users set
start_pref_property != 0.0butlimit_pref_property = 0.0, the code will consider that property “disabled” for the entire training, because the boolean check requires both to be nonzero. Verify that this logic matches intended partial usage behavior and does not confuse users expecting a gradual fade to zero.
125-135: Check division by zero for starter_learning_rate.When computing the coefficient (
coef = learning_rate / self.starter_learning_rate), a zero or near-zerostarter_learning_ratecould lead to NaN or infinite values. Ensure that the caller or constructor either enforces a positivestarter_learning_rateor handles this edge case gracefully.
160-220: Property loss logic looks good.The scaling factor application, standardization, and metric handling appear correct. No immediate performance or correctness issues noted.
221-285: Atomic property loss logic looks good.This block correctly mirrors the property case for atomic properties. The approach to standardization and separate metrics is cleanly implemented.
292-302: Appropriate label requirement for atomic property.Adding a data requirement item for the atomic variant of the property ensures correct hooking of labels. The design is consistent with the new atomic property logic.
303-311: Label requirement for standard property fits well.Retaining a separate requirement item for the non-atomic property preserves backward compatibility and clarity. Implementation looks good.
See discussion #4603
One may want to use DeepMD-kit to fit per-atom properties, such as charge, magnetization for each atom.
Summary by CodeRabbit