Merge branch 'master' into feat-pt-expt-linear

wanghan-iapcm · web-flow · commit f8eff039fa7a · 2026-03-18T15:37:50.000+08:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -30,7 +30,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.15.5
+    rev: v0.15.6
     hooks:
       - id: ruff
         args: ["--fix"]
@@ -39,14 +39,6 @@ repos:
       - id: ruff-format
         exclude: ^source/3rdparty
         types_or: [python, pyi, jupyter]
-  - repo: https://github.com/pycqa/flake8
-    # flake8 cannot autofix
-    rev: "7.3.0"
-    hooks:
-      - id: flake8
-        additional_dependencies:
-          - torchfix==0.7.0
-          - flake8-pyproject==1.2.3
   # numpydoc
   - repo: https://github.com/Carreau/velin
     rev: 0.0.12
@@ -70,7 +62,7 @@ repos:
           - mdformat-gfm-alerts==2.0.0
   # C++
   - repo: https://github.com/pre-commit/mirrors-clang-format
-    rev: v22.1.0
+    rev: v22.1.1
     hooks:
       - id: clang-format
         exclude: ^(source/3rdparty|source/lib/src/gpu/cudart/.+\.inc|.+\.ipynb$|source/tests/infer/.+\.json$)
diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py
@@ -245,7 +245,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py
@@ -162,7 +162,7 @@ def enable_compression(
             table_stride_2: float = 0.1,
             check_frequency: int = -1,
         ) -> None:
-            """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+            """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
             Parameters
             ----------
diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py
@@ -600,7 +600,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py
@@ -264,7 +264,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py
@@ -579,7 +579,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py
@@ -237,7 +237,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py
@@ -262,7 +262,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py
@@ -364,7 +364,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py
@@ -289,7 +289,7 @@ def enable_compression(
         table_stride_2: float = 0.1,
         check_frequency: int = -1,
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py
@@ -160,6 +160,64 @@ def train(
     trainer.run()
 
 
+def freeze(
+    model: str,
+    output: str = "frozen_model.pte",
+    head: str | None = None,
+) -> None:
+    """Freeze a pt_expt checkpoint into a .pte exported model.
+
+    Parameters
+    ----------
+    model : str
+        Path to the checkpoint file (.pt).
+    output : str
+        Path for the output .pte file.
+    head : str or None
+        Head to freeze in multi-task mode (not yet supported).
+    """
+    import torch
+
+    from deepmd.pt_expt.model.get_model import (
+        get_model,
+    )
+    from deepmd.pt_expt.train.wrapper import (
+        ModelWrapper,
+    )
+    from deepmd.pt_expt.utils.env import (
+        DEVICE,
+    )
+    from deepmd.pt_expt.utils.serialization import (
+        deserialize_to_file,
+    )
+
+    state_dict = torch.load(model, map_location=DEVICE, weights_only=True)
+    if "model" in state_dict:
+        state_dict = state_dict["model"]
+
+    extra_state = state_dict.get("_extra_state")
+    if not isinstance(extra_state, dict) or "model_params" not in extra_state:
+        raise ValueError(
+            f"Unsupported checkpoint format at '{model}': missing "
+            "'_extra_state.model_params' in model state dict."
+        )
+    model_params = extra_state["model_params"]
+
+    if head is not None and "model_dict" in model_params:
+        raise NotImplementedError(
+            "Multi-task freeze is not yet supported for the pt_expt backend."
+        )
+
+    m = get_model(model_params)
+    wrapper = ModelWrapper(m)
+    wrapper.load_state_dict(state_dict)
+    m.eval()
+
+    model_dict = m.serialize()
+    deserialize_to_file(output, {"model": model_dict})
+    log.info("Saved frozen model to %s", output)
+
+
 def main(args: list[str] | argparse.Namespace | None = None) -> None:
     """Entry point for the pt_expt backend CLI.
 
@@ -195,6 +253,28 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None:
             skip_neighbor_stat=FLAGS.skip_neighbor_stat,
             output=FLAGS.output,
         )
+    elif FLAGS.command == "freeze":
+        if Path(FLAGS.checkpoint_folder).is_dir():
+            checkpoint_path = Path(FLAGS.checkpoint_folder)
+            # pt_expt training saves a symlink "model.ckpt.pt" → latest ckpt
+            default_ckpt = checkpoint_path / "model.ckpt.pt"
+            if default_ckpt.exists():
+                FLAGS.model = str(default_ckpt)
+            else:
+                raise FileNotFoundError(
+                    f"Cannot find checkpoint in '{checkpoint_path}'. "
+                    "Expected 'model.ckpt.pt' (created by pt_expt training)."
+                )
+        else:
+            model_path = Path(FLAGS.checkpoint_folder)
+            if not model_path.exists():
+                raise FileNotFoundError(
+                    f"Checkpoint path '{model_path}' does not exist."
+                )
+            FLAGS.model = str(model_path)
+        if not FLAGS.output.endswith((".pte", ".pt2")):
+            FLAGS.output = str(Path(FLAGS.output).with_suffix(".pte"))
+        freeze(model=FLAGS.model, output=FLAGS.output, head=FLAGS.head)
     else:
         raise RuntimeError(
             f"Unsupported command '{FLAGS.command}' for the pt_expt backend."
diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py
@@ -140,7 +140,7 @@ def compute_input_stats(
         input_dict: dict[str, list[np.ndarray]],
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be
+        """Compute the statistics (avg and std) of the training data. The input will be
         normalized by the statistics.
 
         Parameters
@@ -228,7 +228,7 @@ def enable_compression(
         check_frequency: int = -1,
         suffix: str = "",
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the
         training data.
 
         Parameters
diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py
@@ -144,7 +144,7 @@ def compute_input_stats(
         real_natoms_vec: list | None = None,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
@@ -183,22 +183,22 @@ def compute_input_stats(
             )
 
     def merge_input_stats(self, stat_dict: dict[str, float]) -> None:
-        """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
+        """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd.
 
         Parameters
         ----------
         stat_dict
-                The dict of statisitcs computed from compute_input_stats, including:
+                The dict of statistics computed from compute_input_stats, including:
             sumr
-                    The sum of radial statisitcs.
+                    The sum of radial statistics.
             suma
-                    The sum of relative coord statisitcs.
+                    The sum of relative coord statistics.
             sumn
                     The sum of neighbor numbers.
             sumr2
-                    The sum of square of radial statisitcs.
+                    The sum of square of radial statistics.
             suma2
-                    The sum of square of relative coord statisitcs.
+                    The sum of square of relative coord statistics.
         """
         for ii in self.descrpt_list:
             ii.merge_input_stats(stat_dict)
@@ -315,7 +315,7 @@ def enable_compression(
         check_frequency: int = -1,
         suffix: str = "",
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the
         training data.
 
         Parameters
diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py
@@ -164,7 +164,7 @@ def compute_input_stats(
         input_dict: dict,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
diff --git a/deepmd/tf/descriptor/se_a.py b/deepmd/tf/descriptor/se_a.py
@@ -355,7 +355,7 @@ def compute_input_stats(
         input_dict: dict,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
@@ -401,22 +401,22 @@ def compute_input_stats(
             self.merge_input_stats(stat_dict)
 
     def merge_input_stats(self, stat_dict: dict[str, Any]) -> None:
-        """Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
+        """Merge the statistics computed from compute_input_stats to obtain the self.davg and self.dstd.
 
         Parameters
         ----------
         stat_dict
-                The dict of statisitcs computed from compute_input_stats, including:
+                The dict of statistics computed from compute_input_stats, including:
             sumr
-                    The sum of radial statisitcs.
+                    The sum of radial statistics.
             suma
-                    The sum of relative coord statisitcs.
+                    The sum of relative coord statistics.
             sumn
                     The sum of neighbor numbers.
             sumr2
-                    The sum of square of radial statisitcs.
+                    The sum of square of radial statistics.
             suma2
-                    The sum of square of relative coord statisitcs.
+                    The sum of square of relative coord statistics.
         """
         all_davg = []
         all_dstd = []
@@ -452,7 +452,7 @@ def enable_compression(
         check_frequency: int = -1,
         suffix: str = "",
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/tf/descriptor/se_a_ef.py b/deepmd/tf/descriptor/se_a_ef.py
@@ -176,7 +176,7 @@ def compute_input_stats(
         input_dict: dict,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py
@@ -234,7 +234,7 @@ def compute_input_stats(
         input_dict: dict,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py
@@ -347,7 +347,7 @@ def compute_input_stats(
         real_natoms_vec: list | None = None,
         **kwargs: Any,
     ) -> None:
-        """Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.
+        """Compute the statistics (avg and std) of the training data. The input will be normalized by the statistics.
 
         Parameters
         ----------
@@ -432,7 +432,7 @@ def enable_compression(
         suffix: str = "",
         tebd_suffix: str = "",
     ) -> None:
-        """Receive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
@@ -2362,7 +2362,7 @@ def enable_compression(
         suffix: str = "",
         tebd_suffix: str = "",
     ) -> None:
-        """Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
+        """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data.
 
         Parameters
         ----------
diff --git a/deepmd/tf/descriptor/se_r.py b/deepmd/tf/descriptor/se_r.py
diff --git a/deepmd/tf/descriptor/se_t.py b/deepmd/tf/descriptor/se_t.py
diff --git a/source/tests/pt_expt/test_dp_freeze.py b/source/tests/pt_expt/test_dp_freeze.py
diff --git a/source/tests/pt_expt/test_dp_test.py b/source/tests/pt_expt/test_dp_test.py
