feat(pt): add plugin for data modifier #4661
Conversation
…-kit into devel-modifier-plugin
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| box = data["box"][:get_nframes, :] | ||
| atype = data["atype"][:get_nframes, :] | ||
| atype = atype[0] | ||
| nframes = coord.shape[0] |
Check notice
Code scanning / CodeQL
Unused local variable Note
📝 Walkthrough""" WalkthroughThe pull request introduces a public data modifier API by adding two new files in the modifier package. A new class, Changes
Sequence Diagram(s)sequenceDiagram
participant Trainer
participant Modifier
participant Data
Trainer->>Modifier: get_data_modifier(modifier_params)
Trainer->>Data: Fetch training/validation sample
Trainer->>Modifier: modify_data(sample)
Modifier->>Modifier: Process data (slice inputs, call forward)
Modifier-->>Trainer: Return modified data
Suggested labels
Suggested reviewers
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 Pylint (3.3.7)deepmd/pt/train/training.pyNo files to lint: exiting. 📜 Recent review detailsConfiguration used: CodeRabbit UI 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms (29)
✨ 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:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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 (2)
deepmd/pt/modifier/base_modifier.py (2)
41-44: Consider simplifying the box assignment.
A ternary operator can reduce verbosity here:- if data["box"] is None: - box = None - else: - box = data["box"][:get_nframes, :] + box = None if data["box"] is None else data["box"][:get_nframes, :]🧰 Tools
🪛 Ruff (0.8.2)
41-44: Use ternary operator
box = None if data["box"] is None else data["box"][:get_nframes, :]instead ofif-else-blockReplace
if-else-block withbox = None if data["box"] is None else data["box"][:get_nframes, :](SIM108)
47-47: Remove or use thenframesvariable.
Currently,nframes = coord.shape[0]is not used, which may confuse future maintainers.🧰 Tools
🪛 Ruff (0.8.2)
47-47: Local variable
nframesis assigned to but never usedRemove assignment to unused variable
nframes(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
deepmd/pt/modifier/__init__.py(1 hunks)deepmd/pt/modifier/base_modifier.py(1 hunks)deepmd/pt/train/training.py(7 hunks)deepmd/pt/utils/stat.py(1 hunks)
🧰 Additional context used
🧬 Code Definitions (2)
deepmd/pt/modifier/__init__.py (1)
deepmd/pt/modifier/base_modifier.py (1) (1)
BaseModifier(9-56)
deepmd/pt/train/training.py (3)
deepmd/pt/modifier/base_modifier.py (2) (2)
BaseModifier(9-56)modify_data(14-56)deepmd/pd/train/training.py (1) (1)
get_additional_data_requirement(1163-1187)deepmd/pd/utils/stat.py (1) (1)
make_stat_input(40-85)
🪛 Ruff (0.8.2)
deepmd/pt/modifier/base_modifier.py
41-44: Use ternary operator box = None if data["box"] is None else data["box"][:get_nframes, :] instead of if-else-block
Replace if-else-block with box = None if data["box"] is None else data["box"][:get_nframes, :]
(SIM108)
47-47: Local variable nframes is assigned to but never used
Remove assignment to unused variable nframes
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz)
- GitHub Check: Test C++ (false)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Analyze (python)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz)
- GitHub Check: Test C++ (true)
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Test Python (1, 3.9)
🔇 Additional comments (11)
deepmd/pt/modifier/__init__.py (1)
1-8: Good job exposing the BaseModifier API.
This new__init__.pycleanly re-exports theBaseModifierclass and ensures users can import it directly fromdeepmd.pt.modifier.deepmd/pt/utils/stat.py (2)
50-51: Conditional logging is well-handled.
Only logging whennbatches > 0helps keep logs cleaner in scenarios where no batches are processed.
56-59: Logic for handling nbatches == -1 is clear and correct.
This new condition ensures the entire dataset is used whennbatchesis -1. No issues found.deepmd/pt/train/training.py (8)
39-41: Import of BaseModifier is appropriate.
This import makes the newly introduced functionality available where needed.
140-149: Modifier parameter handling is well-structured.
The assertion preventing usage in multi-task scenarios is clear and avoids incompatible configurations.
231-231: Defaultingmodifierto None is appropriate.
Makes the modifier usage optional without complicating the training interface.
239-250: Verify data modification logic.
Applyingmodifier.modify_datato every system might lead to repeated transformations ifsingle_model_statis called multiple times. Confirm this matches your intended workflow.
345-345: Single-model signature usage is consistent.
Passingmodifier=self.modifierensures the same modifier instance is applied throughout the training flow.
384-384: Multi-task signature usage is consistent.
Again, passingmodifier=self.modifierallows uniform data processing across tasks if needed.
1075-1081: Storing thedata_modifierstate is a good idea.
Consider providing a loading mechanism in the future so thatdata_modifiercan be restored automatically.
1389-1400: Factory function for data modifiers looks good.
Encapsulates logic for dynamically obtaining modifier classes, making the code more extensible.
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4661 +/- ##
==========================================
- Coverage 84.80% 84.76% -0.05%
==========================================
Files 698 700 +2
Lines 67798 67849 +51
Branches 3542 3542
==========================================
+ Hits 57494 57509 +15
- Misses 9171 9205 +34
- Partials 1133 1135 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
Pull Request Overview
This PR introduces a new plugin for modifying data in the PyTorch backend and integrates it into the training workflow. Key changes include:
- Creation of a new BaseModifier class and registration in the modifier package.
- Integration of the data modifier into the training process, including saving its state.
- Minor adjustments to the statistics data preparation in the utils module.
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| deepmd/pt/modifier/init.py | Exposes BaseModifier for external use. |
| deepmd/pt/modifier/base_modifier.py | Adds a new BaseModifier class for data modification. |
| deepmd/pt/train/training.py | Integrates the data modifier into training data preparation and model saving. |
| deepmd/pt/utils/stat.py | Tweaks logging and batch calculation in the statistics utility. |
Comments suppressed due to low confidence (2)
deepmd/pt/modifier/base_modifier.py:9
- Ensure that make_base_modifier() returns a valid class to use for multiple inheritance with torch.nn.Module. If it does not, consider revising the inheritance structure or renaming for clarity.
class BaseModifier(torch.nn.Module, make_base_modifier()):
deepmd/pt/modifier/base_modifier.py:40
- The variable get_nframes is explicitly set to None, which will slice the full array; if a limit on the number of frames was intended, assign get_nframes an appropriate value.
coord = data["coord"][:get_nframes, :]
| # modifier for the training data | ||
| modifier_params = model_params.get("modifier", None) | ||
| if modifier_params is not None: | ||
| assert self.multi_task is False, ( |
There was a problem hiding this comment.
NotImplementedError is preferred to assert.
| if self.modifier is not None: | ||
| save_dict["data_modifier"] = self.modifier.state_dict() |
There was a problem hiding this comment.
I see this is saved, but how is it recovered?
There was a problem hiding this comment.
I plan to restore it by adapting wrapper. Do you have any other idea for this?
| # modify data | ||
| if modifier is not None: | ||
| log.info(f"Using {modifier.modifier_type} as data modifier") | ||
| for _data in [_training_data, _validation_data]: | ||
| if _data is not None: | ||
| all_sampled = make_stat_input( | ||
| _data.systems, | ||
| _data.dataloaders, | ||
| -1, | ||
| ) | ||
| for sampled in all_sampled: | ||
| modifier.modify_data(sampled) |
There was a problem hiding this comment.
I noticed that this modification only affects the data statistics. It still uses the original data in get_data for actual training.
There was a problem hiding this comment.
I don't think so. modifier.modify_data(sampled) in line 267 has directly modify the training data sampled.
There was a problem hiding this comment.
The sampled data is only used for data statistics, while it uses the original data from get_data function in Line 1092.
Summary by CodeRabbit
New Features
Refactor