feat(pt/pd): add size option to dp show#4783
Merged
iProzd merged 2 commits intodeepmodeling:develfrom Jun 10, 2025
Merged
Conversation
Collaborator
Author
|
We need a proper place to add docs for dp show. |
Contributor
📝 Walkthrough## Walkthrough
Support for displaying model parameter counts was added to the "show" command. This includes backend implementations to compute parameter counts, updates to the command-line interface to accept a new "size" attribute, logic to display the counts, and corresponding test updates to verify the new output. Documentation was also added to describe the new feature.
## Changes
| File(s) | Change Summary |
|--------------------------------------------------------------|-----------------------------------------------------------------------------------------------------------------------|
| `deepmd/entrypoints/show.py` | Added logic to handle "size" attribute: retrieves and logs model parameter counts. |
| `deepmd/infer/deep_eval.py` | Added `get_model_size` abstract method to backend interface and concrete method to high-level interface. |
| `deepmd/pt/infer/deep_eval.py`, `deepmd/pd/infer/deep_eval.py`| Implemented `get_model_size` method to compute and return parameter counts by category. |
| `deepmd/main.py` | Extended "show" command's valid attributes to include "size". |
| `source/tests/pt/test_dp_show.py` | Updated tests to include "size" attribute and verify parameter count output. |
| `doc/model/index.rst` | Added "show-model-info" entry to model documentation index. |
| `doc/model/show-model-info.md` | Added new documentation describing the `dp show` command and the "size" attribute usage. |
## Sequence Diagram(s)
```mermaid
sequenceDiagram
participant User
participant CLI
participant ShowEntrypoint
participant DeepEval
participant Backend
User->>CLI: dp show size <model>
CLI->>ShowEntrypoint: parse arguments (attribute="size")
ShowEntrypoint->>DeepEval: get_model_size()
DeepEval->>Backend: get_model_size()
Backend-->>DeepEval: {descriptor, fitting-net, total}
DeepEval-->>ShowEntrypoint: {descriptor, fitting-net, total}
ShowEntrypoint->>CLI: log parameter counts
CLI-->>User: display parameter countsSuggested labels
Suggested reviewers
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
deepmd/pt/infer/deep_eval.py (1)
627-651: Optimize dictionary key iteration for better performance.The implementation is functionally correct and follows good practices for parameter counting. However, the static analysis hints highlight a minor performance optimization opportunity.
Apply this optimization to improve performance:
- sum_param_des = sum( - params[k].numel() - for k in params.keys() - if "descriptor" in k and "mean" not in k and "stddev" not in k - ) - sum_param_fit = sum( - params[k].numel() - for k in params.keys() - if "fitting" in k and "_networks" not in k - ) + sum_param_des = sum( + params[k].numel() + for k in params + if "descriptor" in k and "mean" not in k and "stddev" not in k + ) + sum_param_fit = sum( + params[k].numel() + for k in params + if "fitting" in k and "_networks" not in k + )🧰 Tools
🪛 Ruff (0.11.9)
639-639: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
644-644: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/pd/infer/deep_eval.py (1)
485-509: Implement parameter counting with Paddle-specific utilities and optimize dictionary access.The implementation is functionally correct and consistent with the PyTorch backend. However, there are a couple of improvements to consider:
- Use the Paddle-specific
numelutility function fromdeepmd.pd.utils.decompfor better compatibility with Paddle's dynamic/static modes.- Optimize dictionary iteration as suggested by static analysis.
Apply this diff to improve the implementation:
+from deepmd.pd.utils.decomp import ( + numel, +) def get_model_size(self) -> dict: """Get model parameter count. Returns ------- dict A dictionary containing the number of parameters in the model. The keys are 'descriptor', 'fitting_net', and 'total'. """ params = self.dp.state_dict() sum_param_des = sum( - params[k].numel() - for k in params.keys() + numel(params[k]) + for k in params if "descriptor" in k and "mean" not in k and "stddev" not in k ) sum_param_fit = sum( - params[k].numel() - for k in params.keys() + numel(params[k]) + for k in params if "fitting" in k and "_networks" not in k ) return { "descriptor": sum_param_des, "fitting-net": sum_param_fit, "total": sum_param_des + sum_param_fit, }🧰 Tools
🪛 Ruff (0.11.9)
497-497: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
502-502: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (6)
deepmd/entrypoints/show.py(1 hunks)deepmd/infer/deep_eval.py(2 hunks)deepmd/main.py(1 hunks)deepmd/pd/infer/deep_eval.py(1 hunks)deepmd/pt/infer/deep_eval.py(1 hunks)source/tests/pt/test_dp_show.py(2 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
deepmd/infer/deep_eval.py (6)
deepmd/pd/infer/deep_eval.py (1)
get_model_size(485-509)deepmd/pt/infer/deep_eval.py (1)
get_model_size(627-651)deepmd/backend/dpmodel.py (1)
deep_eval(63-75)deepmd/backend/pytorch.py (1)
deep_eval(69-79)deepmd/backend/backend.py (1)
deep_eval(157-165)deepmd/backend/tensorflow.py (1)
deep_eval(78-88)
deepmd/entrypoints/show.py (3)
deepmd/pd/infer/deep_eval.py (1)
get_model_size(485-509)deepmd/infer/deep_eval.py (2)
get_model_size(294-296)get_model_size(568-570)deepmd/pt/infer/deep_eval.py (1)
get_model_size(627-651)
source/tests/pt/test_dp_show.py (2)
source/tests/pt/common.py (1)
run_dp(27-49)source/tests/pd/test_dp_show.py (2)
test_frozen_model(71-85)test_frozen_model(192-206)
deepmd/pd/infer/deep_eval.py (4)
deepmd/infer/deep_eval.py (2)
get_model_size(294-296)get_model_size(568-570)deepmd/pt/infer/deep_eval.py (1)
get_model_size(627-651)deepmd/pd/train/wrapper.py (1)
state_dict(199-203)deepmd/pd/utils/decomp.py (1)
numel(130-134)
🪛 Ruff (0.11.9)
deepmd/pt/infer/deep_eval.py
639-639: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
644-644: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
deepmd/pd/infer/deep_eval.py
497-497: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
502-502: Use key in dict instead of key in dict.keys()
Remove .keys()
(SIM118)
⏰ Context from checks skipped due to timeout of 90000ms (29)
- GitHub Check: Test Python (6, 3.9)
- GitHub Check: Test Python (5, 3.9)
- GitHub Check: Test Python (2, 3.12)
- GitHub Check: Test Python (4, 3.9)
- GitHub Check: Test Python (3, 3.9)
- GitHub Check: Test Python (5, 3.12)
- GitHub Check: Test Python (6, 3.12)
- GitHub Check: Test Python (3, 3.12)
- GitHub Check: Test Python (4, 3.12)
- GitHub Check: Test Python (2, 3.9)
- GitHub Check: Test Python (1, 3.9)
- GitHub Check: Test Python (1, 3.12)
- GitHub Check: Build wheels for cp311-win_amd64
- GitHub Check: Build wheels for cp311-macosx_x86_64
- GitHub Check: Build wheels for cp310-manylinux_aarch64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cuda120, cuda)
- GitHub Check: Build C++ (rocm, rocm)
- GitHub Check: Build wheels for cp311-macosx_arm64
- GitHub Check: Build wheels for cp311-manylinux_x86_64
- GitHub Check: Build C++ (cpu, cpu)
- GitHub Check: Build C++ (clang, clang)
- GitHub Check: Build C++ (cuda, cuda)
- GitHub Check: Analyze (python)
- GitHub Check: Analyze (c-cpp)
- GitHub Check: Test C++ (true)
- GitHub Check: Test C++ (false)
- 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)
🔇 Additional comments (8)
deepmd/entrypoints/show.py (1)
66-71: LGTM! Clean implementation of parameter count display.The size attribute handling follows the established pattern in the file and includes thoughtful touches like thousands separators for readability and appropriate messaging for multitask models.
deepmd/main.py (1)
854-854: LGTM! Correct CLI interface update.The addition of "size" to the choices list properly enables the new functionality in the command-line interface.
deepmd/infer/deep_eval.py (2)
294-296: LGTM! Proper abstract interface definition.The abstract method follows the established pattern and provides a clear contract for backend implementations.
568-570: LGTM! Consistent delegation pattern.The delegation to the backend implementation follows the same pattern as other methods like
get_model_def_script(), maintaining architectural consistency.source/tests/pt/test_dp_show.py (4)
50-66: Comprehensive test coverage for the new size attribute functionality.The test updates appropriately verify the new "size" attribute functionality with proper assertions for parameter counts. The changes include adjusting assertion indices and adding specific checks for parameter count output lines.
70-86: Consistent test pattern for frozen model verification.The frozen model test follows the same pattern as the checkpoint test, ensuring consistent verification of the size attribute across different model formats.
155-190: Thorough multi-task model testing with appropriate parameter count verification.The multi-task model test correctly verifies parameter counts for a single branch model, which is the expected behavior when showing size information for multi-task models.
194-210: Complete test coverage for multi-task frozen model scenario.The test appropriately handles the multi-task to single-task conversion scenario when freezing a specific branch, with correct parameter count assertions.
njzjz
reviewed
Jun 5, 2025
Contributor
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (3)
deepmd/pd/infer/deep_eval.py (1)
485-507: Apply code style improvements suggested by static analysis.The implementation is functionally correct, but there are unnecessary
.keys()calls that can be simplified.Apply this diff to improve code style:
def get_model_size(self) -> dict: """Get model parameter count. Returns ------- dict A dictionary containing the number of parameters in the model. The keys are 'descriptor', 'fitting_net', and 'total'. """ params_dict = dict(self.dp.named_parameters()) sum_param_des = sum( - params_dict[k].numel() for k in params_dict.keys() if "descriptor" in k + params_dict[k].numel() for k in params_dict if "descriptor" in k ) sum_param_fit = sum( params_dict[k].numel() - for k in params_dict.keys() + for k in params_dict if "fitting" in k and "_networks" not in k ) return { "descriptor": sum_param_des, "fitting-net": sum_param_fit, "total": sum_param_des + sum_param_fit, }🧰 Tools
🪛 Ruff (0.11.9)
496-496: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
500-500: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
deepmd/pt/infer/deep_eval.py (1)
627-649: Apply code style improvements for consistency with PD backend.The implementation correctly mirrors the PaddlePaddle backend implementation, but the same style improvements apply here.
Apply this diff to improve code style:
def get_model_size(self) -> dict: """Get model parameter count. Returns ------- dict A dictionary containing the number of parameters in the model. The keys are 'descriptor', 'fitting_net', and 'total'. """ params_dict = dict(self.dp.named_parameters()) sum_param_des = sum( - params_dict[k].numel() for k in params_dict.keys() if "descriptor" in k + params_dict[k].numel() for k in params_dict if "descriptor" in k ) sum_param_fit = sum( params_dict[k].numel() - for k in params_dict.keys() + for k in params_dict if "fitting" in k and "_networks" not in k ) return { "descriptor": sum_param_des, "fitting-net": sum_param_fit, "total": sum_param_des + sum_param_fit, }🧰 Tools
🪛 Ruff (0.11.9)
638-638: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
642-642: Use
key in dictinstead ofkey in dict.keys()Remove
.keys()(SIM118)
doc/model/show-model-info.md (1)
67-76: Add language specifications to code blocks for better documentation.The documentation is comprehensive and well-written. Consider adding language specifications to the output example code blocks for better syntax highlighting and accessibility.
Apply this diff to improve the code blocks:
For a singletask model, the output might look like: -``` +```text This is a singletask model The type_map is ['O', 'H', 'Au'] The descriptor parameter is {'type': 'se_e2_a', 'sel': [46, 92, 4], 'rcut': 4.0} The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1} Parameter counts: Parameters in descriptor: 19,350 Parameters in fitting-net: 119,091 Parameters in total: 138,441For a multitask model, if
model-branchis selected, it will additionally display available branches:-
+text
This is a multitask model
Available model branches are ['branch1', 'branch2', 'RANDOM'], where 'RANDOM' means using a randomly initialized fitting net.
...Also applies to: 80-84 <details> <summary>🧰 Tools</summary> <details> <summary>🪛 markdownlint-cli2 (0.17.2)</summary> 67-67: Fenced code blocks should have a language specified null (MD040, fenced-code-language) </details> </details> </blockquote></details> </blockquote></details> <details> <summary>📜 Review details</summary> **Configuration used: CodeRabbit UI** **Review profile: CHILL** **Plan: Pro** <details> <summary>📥 Commits</summary> Reviewing files that changed from the base of the PR and between 3b84fd40be7434eb031f68f2d4fd294387ce58c0 and a264eb8dd65e73974675836b90697fd7c57ca104. </details> <details> <summary>📒 Files selected for processing (5)</summary> * `deepmd/pd/infer/deep_eval.py` (1 hunks) * `deepmd/pt/infer/deep_eval.py` (1 hunks) * `doc/model/index.rst` (1 hunks) * `doc/model/show-model-info.md` (1 hunks) * `source/tests/pt/test_dp_show.py` (2 hunks) </details> <details> <summary>✅ Files skipped from review due to trivial changes (1)</summary> * doc/model/index.rst </details> <details> <summary>🚧 Files skipped from review as they are similar to previous changes (1)</summary> * source/tests/pt/test_dp_show.py </details> <details> <summary>🧰 Additional context used</summary> <details> <summary>🧬 Code Graph Analysis (2)</summary> <details> <summary>deepmd/pd/infer/deep_eval.py (3)</summary><blockquote> <details> <summary>deepmd/infer/deep_eval.py (2)</summary> * `get_model_size` (294-296) * `get_model_size` (568-570) </details> <details> <summary>deepmd/pt/infer/deep_eval.py (1)</summary> * `get_model_size` (627-649) </details> <details> <summary>deepmd/pd/utils/decomp.py (1)</summary> * `numel` (130-134) </details> </blockquote></details> <details> <summary>deepmd/pt/infer/deep_eval.py (2)</summary><blockquote> <details> <summary>deepmd/infer/deep_eval.py (2)</summary> * `get_model_size` (294-296) * `get_model_size` (568-570) </details> <details> <summary>deepmd/pd/infer/deep_eval.py (1)</summary> * `get_model_size` (485-507) </details> </blockquote></details> </details><details> <summary>🪛 Ruff (0.11.9)</summary> <details> <summary>deepmd/pd/infer/deep_eval.py</summary> 496-496: Use `key in dict` instead of `key in dict.keys()` Remove `.keys()` (SIM118) --- 500-500: Use `key in dict` instead of `key in dict.keys()` Remove `.keys()` (SIM118) </details> <details> <summary>deepmd/pt/infer/deep_eval.py</summary> 638-638: Use `key in dict` instead of `key in dict.keys()` Remove `.keys()` (SIM118) --- 642-642: Use `key in dict` instead of `key in dict.keys()` Remove `.keys()` (SIM118) </details> </details> <details> <summary>🪛 LanguageTool</summary> <details> <summary>doc/model/show-model-info.md</summary> [misspelling] ~4-~4: This word is normally spelled as one. Context: ...eter statistics in both single-task and multi-task settings. ## Command Syntax ```bash d... (EN_COMPOUNDS_MULTI_TASK) --- [uncategorized] ~12-~12: Loose punctuation mark. Context: ...<INPUT> <ATTRIBUTES...> ``` - `<INPUT>`: Path to the model checkpoint file or fr... (UNLIKELY_OPENING_PUNCTUATION) --- [uncategorized] ~13-~13: Loose punctuation mark. Context: ...e or frozen model file. - `<ATTRIBUTES>`: One or more information categories to d... (UNLIKELY_OPENING_PUNCTUATION) --- [uncategorized] ~15-~15: Loose punctuation mark. Context: ...upported values are: - `model-branch`: Shows available branches for multi-task... (UNLIKELY_OPENING_PUNCTUATION) --- [misspelling] ~15-~15: This word is normally spelled as one. Context: ...l-branch`: Shows available branches for multi-task models. - `type-map`: Shows the type ... (EN_COMPOUNDS_MULTI_TASK) --- [uncategorized] ~16-~16: Loose punctuation mark. Context: ...es for multi-task models. - `type-map`: Shows the type mapping used by the mode... (UNLIKELY_OPENING_PUNCTUATION) --- [uncategorized] ~17-~17: Loose punctuation mark. Context: ...ping used by the model. - `descriptor`: Displays the model descriptor parameter... (UNLIKELY_OPENING_PUNCTUATION) --- [uncategorized] ~18-~18: Loose punctuation mark. Context: ...descriptor parameters. - `fitting-net`: Displays parameters of the fitting netw... (UNLIKELY_OPENING_PUNCTUATION) </details> </details> <details> <summary>🪛 markdownlint-cli2 (0.17.2)</summary> <details> <summary>doc/model/show-model-info.md</summary> 67-67: Fenced code blocks should have a language specified null (MD040, fenced-code-language) --- 80-80: Fenced code blocks should have a language specified null (MD040, fenced-code-language) </details> </details> </details> <details> <summary>⏰ Context from checks skipped due to timeout of 90000ms (29)</summary> * 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: Test Python (4, 3.12) * GitHub Check: Build wheels for cp310-manylinux_aarch64 * GitHub Check: Test Python (4, 3.9) * GitHub Check: Build wheels for cp311-win_amd64 * GitHub Check: Test Python (3, 3.12) * GitHub Check: Build wheels for cp311-macosx_arm64 * GitHub Check: Test Python (3, 3.9) * GitHub Check: Build C++ (clang, clang) * GitHub Check: Build wheels for cp311-macosx_x86_64 * GitHub Check: Test Python (2, 3.12) * GitHub Check: Build C++ (rocm, rocm) * GitHub Check: Build wheels for cp311-manylinux_x86_64 * GitHub Check: Test Python (2, 3.9) * GitHub Check: Analyze (python) * GitHub Check: Build C++ (cuda120, cuda) * GitHub Check: Test C++ (false) * GitHub Check: Build wheels for cp311-manylinux_x86_64 * GitHub Check: Test Python (1, 3.12) * GitHub Check: Build C library (2.14, >=2.5.0rc0,<2.15, libdeepmd_c_cu11.tar.gz) * GitHub Check: Build C++ (cuda, cuda) * GitHub Check: Test C++ (true) * GitHub Check: Test Python (1, 3.9) * GitHub Check: Build C++ (cpu, cpu) * GitHub Check: Build C library (2.18, libdeepmd_c.tar.gz) * GitHub Check: Analyze (c-cpp) </details> </details> <!-- This is an auto-generated comment by CodeRabbit for review status -->
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## devel #4783 +/- ##
==========================================
- Coverage 84.80% 84.79% -0.01%
==========================================
Files 698 698
Lines 67787 67818 +31
Branches 3542 3542
==========================================
+ Hits 57486 57506 +20
- Misses 9168 9177 +9
- Partials 1133 1135 +2 ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
njzjz
approved these changes
Jun 9, 2025
Chengqian-Zhang
approved these changes
Jun 9, 2025
ChiahsinChu
pushed a commit
to ChiahsinChu/deepmd-kit
that referenced
this pull request
Dec 17, 2025
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit - **New Features** - Added the ability to display model parameter counts using the "size" attribute in the "show" command. - Users can now view detailed parameter counts for descriptor, fitting-net, and total parameters for both single-task and multi-task models. - **Tests** - Updated tests to verify correct output of model parameter counts when using the new "size" attribute. - **Documentation** - Added new documentation for the "show" command detailing usage and output for various model attributes, including the new "size" attribute. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
New Features
Tests
Documentation