feat(pt/pd): add size option to dp show (deepmodeling#4783)

<!-- 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 -->

Author: iProzd

deepmd/entrypoints/show.py

@@ -63,3 +63,9 @@ def show(
         else:
             fitting_net = model_params["fitting_net"]
             log.info(f"The fitting_net parameter is {fitting_net}")
+    if "size" in ATTRIBUTES:
+        size_dict = model.get_model_size()
+        log_prefix = " for a single branch model" if model_is_multi_task else ""
+        log.info(f"Parameter counts{log_prefix}:")
+        for k in sorted(size_dict):
+            log.info(f"Parameters in {k}: {size_dict[k]:,}")

deepmd/infer/deep_eval.py

@@ -291,6 +291,10 @@ def get_model_def_script(self) -> dict:
         """Get model definition script."""
         raise NotImplementedError("Not implemented in this backend.")
 
+    def get_model_size(self) -> dict:
+        """Get model parameter count."""
+        raise NotImplementedError("Not implemented in this backend.")
+
 
 class DeepEval(ABC):
     """High-level Deep Evaluator interface.
@@ -560,3 +564,7 @@ def get_ntypes_spin(self) -> int:
     def get_model_def_script(self) -> dict:
         """Get model definition script."""
         return self.deep_eval.get_model_def_script()
+
+    def get_model_size(self) -> dict:
+        """Get model parameter count."""
+        return self.deep_eval.get_model_size()

deepmd/main.py

@@ -851,7 +851,7 @@ def main_parser() -> argparse.ArgumentParser:
     )
     parser_show.add_argument(
         "ATTRIBUTES",
-        choices=["model-branch", "type-map", "descriptor", "fitting-net"],
+        choices=["model-branch", "type-map", "descriptor", "fitting-net", "size"],
         nargs="+",
     )
     return parser

deepmd/pd/infer/deep_eval.py

@@ -482,6 +482,30 @@ def get_model_def_script(self) -> str:
         """Get model definition script."""
         return self.model_def_script
 
+    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
+        )
+        sum_param_fit = sum(
+            params_dict[k].numel()
+            for k in params_dict.keys()
+            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,
+        }
+
     def eval_descriptor(
         self,
         coords: np.ndarray,

deepmd/pt/infer/deep_eval.py

@@ -624,6 +624,30 @@ def get_model_def_script(self) -> str:
         """Get model definition script."""
         return self.model_def_script
 
+    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
+        )
+        sum_param_fit = sum(
+            params_dict[k].numel()
+            for k in params_dict.keys()
+            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,
+        }
+
     def eval_descriptor(
         self,
         coords: np.ndarray,

doc/model/index.rst

@@ -28,3 +28,4 @@ Model
    pairtab
    change-bias
    precision
+   show-model-info

doc/model/show-model-info.md

@@ -0,0 +1,84 @@
+# Use `dp show` to show the model information
+
+The `dp show` command is designed to display essential information about a trained model checkpoint or frozen model file.
+This utility helps to understand the model's architecture, configuration, and parameter statistics in both single-task and multi-task settings.
+
+## Command Syntax
+
+```bash
+dp --pt show <INPUT> <ATTRIBUTES...>
+```
+
+- `<INPUT>`: Path to the model checkpoint file or frozen model file.
+- `<ATTRIBUTES>`: One or more information categories to display. Supported values are:
+
+  - `model-branch`: Shows available branches for multi-task models.
+  - `type-map`: Shows the type mapping used by the model.
+  - `descriptor`: Displays the model descriptor parameters.
+  - `fitting-net`: Displays parameters of the fitting network.
+  - `size`: (Supported Backends: PyTorch and PaddlePaddle) Shows the parameter counts for various components.
+
+## Example Usage
+
+```bash
+# For a multi-task model (model.pt)
+dp show model.pt model-branch type-map descriptor fitting-net size
+
+# For a single-task frozen model (frozen_model.pth)
+dp show frozen_model.pth type-map descriptor fitting-net size
+```
+
+## Output Description
+
+Depending on the provided attributes and the model type, the output includes:
+
+- **Model Type**
+
+  - Logs whether the loaded model is a _singletask_ or _multitask_ model.
+
+- **model-branch**
+
+  - _Only available for multitask models._
+  - Lists all available model branches and the special `"RANDOM"` branch, which refers to a randomly initialized fitting net.
+
+- **type-map**
+
+  - For multitask models: Shows the type map for each branch.
+  - For singletask models: Shows the model's type map.
+
+- **descriptor**
+
+  - For multitask models: Displays the descriptor parameter for each branch.
+  - For singletask models: Displays the descriptor parameter.
+
+- **fitting-net**
+
+  - For multitask models: Shows the fitting network parameters for each branch.
+  - For singletask models: Shows the fitting network parameters.
+
+- **size**
+
+  - Prints the number of parameters for each component (`descriptor`, `fitting-net`, etc.), as well as the total parameter count.
+
+## Example Output
+
+For a singletask model, the output might look like:
+
+```
+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,441
+```
+
+For a multitask model, if `model-branch` is selected, it will additionally display available branches:
+
+```
+This is a multitask model
+Available model branches are ['branch1', 'branch2', 'RANDOM'], where 'RANDOM' means using a randomly initialized fitting net.
+...
+```

source/tests/pt/test_dp_show.py

@@ -47,35 +47,43 @@ def setUp(self) -> None:
 
     def test_checkpoint(self) -> None:
         INPUT = "model.pt"
-        ATTRIBUTES = "type-map descriptor fitting-net"
+        ATTRIBUTES = "type-map descriptor fitting-net size"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-3]
+        assert "This is a singletask model" in results[-8]
+        assert "The type_map is ['O', 'H', 'Au']" in results[-7]
         assert (
             "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-2]
+        ) in results[-6]
         assert (
             "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-1]
+            in results[-5]
         )
+        assert "Parameter counts:" in results[-4]
+        assert "Parameters in descriptor: 19,350" in results[-3]
+        assert "Parameters in fitting-net: 119,091" in results[-2]
+        assert "Parameters in total: 138,441" in results[-1]
 
     def test_frozen_model(self) -> None:
         INPUT = "frozen_model.pth"
-        ATTRIBUTES = "type-map descriptor fitting-net"
+        ATTRIBUTES = "type-map descriptor fitting-net size"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'Au']" in results[-3]
+        assert "This is a singletask model" in results[-8]
+        assert "The type_map is ['O', 'H', 'Au']" in results[-7]
         assert (
             "{'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut': 4.0"
-        ) in results[-2]
+        ) in results[-6]
         assert (
             "The fitting_net parameter is {'neuron': [24, 24, 24], 'resnet_dt': True, 'seed': 1}"
-            in results[-1]
+            in results[-5]
         )
+        assert "Parameter counts:" in results[-4]
+        assert "Parameters in descriptor: 19,350" in results[-3]
+        assert "Parameters in fitting-net: 119,091" in results[-2]
+        assert "Parameters in total: 138,441" in results[-1]
 
     def test_checkpoint_error(self) -> None:
         INPUT = "model.pt"
@@ -144,54 +152,62 @@ def setUp(self) -> None:
 
     def test_checkpoint(self) -> None:
         INPUT = "model.ckpt.pt"
-        ATTRIBUTES = "model-branch type-map descriptor fitting-net"
+        ATTRIBUTES = "model-branch type-map descriptor fitting-net size"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a multitask model" in results[-8]
+        assert "This is a multitask model" in results[-12]
         assert (
             "Available model branches are ['model_1', 'model_2', 'RANDOM'], "
             "where 'RANDOM' means using a randomly initialized fitting net."
-            in results[-7]
+            in results[-11]
         )
-        assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-6]
-        assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-5]
+        assert "The type_map of branch model_1 is ['O', 'H', 'B']" in results[-10]
+        assert "The type_map of branch model_2 is ['O', 'H', 'B']" in results[-9]
         assert (
             "model_1"
             and "'type': 'se_e2_a'"
             and "'sel': [46, 92, 4]"
             and "'rcut_smth': 0.5"
-        ) in results[-4]
+        ) in results[-8]
         assert (
             "model_2"
             and "'type': 'se_e2_a'"
             and "'sel': [46, 92, 4]"
             and "'rcut_smth': 0.5"
-        ) in results[-3]
+        ) in results[-7]
         assert (
             "The fitting_net parameter of branch model_1 is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-2]
+            in results[-6]
         )
         assert (
             "The fitting_net parameter of branch model_2 is {'neuron': [9, 8, 7], 'seed': 1111}"
-            in results[-1]
+            in results[-5]
         )
+        assert "Parameter counts for a single branch model:" in results[-4]
+        assert "Parameters in descriptor: 19,350" in results[-3]
+        assert "Parameters in fitting-net: 4,860" in results[-2]
+        assert "Parameters in total: 24,210" in results[-1]
 
     def test_frozen_model(self) -> None:
         INPUT = "frozen_model.pth"
-        ATTRIBUTES = "type-map descriptor fitting-net"
+        ATTRIBUTES = "type-map descriptor fitting-net size"
         with redirect_stderr(io.StringIO()) as f:
             run_dp(f"dp --pt show {INPUT} {ATTRIBUTES}")
         results = f.getvalue().split("\n")[:-1]
-        assert "This is a singletask model" in results[-4]
-        assert "The type_map is ['O', 'H', 'B']" in results[-3]
+        assert "This is a singletask model" in results[-8]
+        assert "The type_map is ['O', 'H', 'B']" in results[-7]
         assert (
             "'type': 'se_e2_a'" and "'sel': [46, 92, 4]" and "'rcut_smth': 0.5"
-        ) in results[-2]
+        ) in results[-6]
         assert (
             "The fitting_net parameter is {'neuron': [1, 2, 3], 'seed': 678}"
-            in results[-1]
+            in results[-5]
         )
+        assert "Parameter counts:" in results[-4]
+        assert "Parameters in descriptor: 19,350" in results[-3]
+        assert "Parameters in fitting-net: 4,860" in results[-2]
+        assert "Parameters in total: 24,210" in results[-1]
 
     def tearDown(self) -> None:
         for f in os.listdir("."):