refactor(pt_expt): use model API for inference, consistent file naming

.gitignore

@@ -62,6 +62,8 @@ test_dp_test_*.out
 
 # Training and model output files
 *.pth
+*.pte
+*.pt2
 *.ckpt*
 checkpoint
 lcurve.out

deepmd/pt_expt/entrypoints/main.py

@@ -259,7 +259,7 @@ def freeze(
     m.eval()
 
     model_dict = m.serialize()
-    deserialize_to_file(output, {"model": model_dict}, model_params=model_params)
+    deserialize_to_file(output, {"model": model_dict, "model_def_script": model_params})
     log.info("Saved frozen model to %s", output)
 
 
@@ -344,7 +344,7 @@ def change_bias(
         )
 
         model_to_change = BaseModel.deserialize(pte_data["model"])
-        model_params = None
+        model_params = pte_data.get("model_def_script")
     else:
         raise RuntimeError(
             "The model provided must be a checkpoint file with a .pt extension "
@@ -440,7 +440,9 @@ def change_bias(
             )
         )
         model_dict = model_to_change.serialize()
-        deserialize_to_file(output_path, {"model": model_dict})
+        deserialize_to_file(
+            output_path, {"model": model_dict, "model_def_script": model_params}
+        )
     log.info(f"Saved model to {output_path}")
 
 

deepmd/pt_expt/infer/deep_eval.py

@@ -16,7 +16,6 @@
     communicate_extended_output,
 )
 from deepmd.dpmodel.output_def import (
-    FittingOutputDef,
     ModelOutputDef,
     OutputVariableCategory,
     OutputVariableDef,
@@ -59,28 +58,6 @@
     import ase.neighborlist
 
 
-def _reconstruct_model_output_def(metadata: dict) -> ModelOutputDef:
-    """Reconstruct ModelOutputDef from stored fitting_output_defs metadata."""
-    var_defs = []
-    for vd in metadata["fitting_output_defs"]:
-        var_defs.append(
-            OutputVariableDef(
-                name=vd["name"],
-                shape=vd["shape"],
-                reducible=vd["reducible"],
-                r_differentiable=vd["r_differentiable"],
-                c_differentiable=vd["c_differentiable"],
-                atomic=vd["atomic"],
-                category=vd["category"],
-                r_hessian=vd["r_hessian"],
-                magnetic=vd["magnetic"],
-                intensive=vd["intensive"],
-            )
-        )
-    fitting_output_def = FittingOutputDef(var_defs)
-    return ModelOutputDef(fitting_output_def)
-
-
 class DeepEval(DeepEvalBackend):
     """PyTorch Exportable backend implementation of DeepEval.
 
@@ -124,9 +101,6 @@ def __init__(
         else:
             self._load_pte(model_file)
 
-        # Reconstruct the model output def from stored fitting output defs
-        self._model_output_def = _reconstruct_model_output_def(self.metadata)
-
         if isinstance(auto_batch_size, bool):
             if auto_batch_size:
                 self.auto_batch_size = AutoBatchSize()
@@ -139,14 +113,30 @@ def __init__(
         else:
             raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize")
 
+    def _init_from_model_json(self, model_json_str: str) -> None:
+        """Deserialize model.json and derive model API from the dpmodel instance."""
+        from deepmd.pt_expt.model.model import (
+            BaseModel,
+        )
+        from deepmd.pt_expt.utils.serialization import (
+            _json_to_numpy,
+        )
+
+        model_dict = json.loads(model_json_str)
+        model_dict = _json_to_numpy(model_dict)
+        self._dpmodel = BaseModel.deserialize(model_dict["model"])
+        self.rcut = self._dpmodel.get_rcut()
+        self.type_map = self._dpmodel.get_type_map()
+        self._model_output_def = ModelOutputDef(self._dpmodel.atomic_output_def())
+
     def _load_pte(self, model_file: str) -> None:
         """Load a .pte (torch.export) model file."""
-        extra_files = {"model_def_script.json": ""}
+        extra_files = {"model.json": "", "model_def_script.json": ""}
         exported = torch.export.load(model_file, extra_files=extra_files)
         self.exported_module = exported.module()
-        self.metadata = json.loads(extra_files["model_def_script.json"])
-        self.rcut = self.metadata["rcut"]
-        self.type_map = self.metadata["type_map"]
+        self._init_from_model_json(extra_files["model.json"])
+        mds = extra_files["model_def_script.json"]
+        self._model_def_script = json.loads(mds) if mds else {}
 
     def _load_pt2(self, model_file: str) -> None:
         """Load a .pt2 (AOTInductor) model file."""
@@ -159,16 +149,17 @@ def _load_pt2(self, model_file: str) -> None:
         # Read metadata from the .pt2 ZIP archive
         with zipfile.ZipFile(model_file, "r") as zf:
             names = zf.namelist()
-            for required in ("extra/model_def_script.json", "extra/output_keys.json"):
-                if required not in names:
-                    raise ValueError(
-                        f"Invalid .pt2 file '{model_file}': missing '{required}'"
-                    )
-            self.metadata = json.loads(zf.read("extra/model_def_script.json"))
-            self._output_keys = json.loads(zf.read("extra/output_keys.json"))
+            if "extra/model.json" not in names:
+                raise ValueError(
+                    f"Invalid .pt2 file '{model_file}': missing 'extra/model.json'"
+                )
+            model_json_str = zf.read("extra/model.json").decode("utf-8")
+            mds = ""
+            if "extra/model_def_script.json" in names:
+                mds = zf.read("extra/model_def_script.json").decode("utf-8")
 
-        self.rcut = self.metadata["rcut"]
-        self.type_map = self.metadata["type_map"]
+        self._init_from_model_json(model_json_str)
+        self._model_def_script = json.loads(mds) if mds else {}
 
         # Load the AOTInductor model package (.pt2 ZIP archive).
         # Uses torch._inductor.aoti_load_package (private API, stable since PyTorch 2.6).
@@ -189,16 +180,16 @@ def get_type_map(self) -> list[str]:
 
     def get_dim_fparam(self) -> int:
         """Get the number (dimension) of frame parameters of this DP."""
-        return self.metadata["dim_fparam"]
+        return self._dpmodel.get_dim_fparam()
 
     def get_dim_aparam(self) -> int:
         """Get the number (dimension) of atomic parameters of this DP."""
-        return self.metadata["dim_aparam"]
+        return self._dpmodel.get_dim_aparam()
 
     @property
     def model_type(self) -> type["DeepEvalWrapper"]:
         """The the evaluator of the model type."""
-        model_output_type = self.metadata["model_output_type"]
+        model_output_type = self._dpmodel.model_output_type()
         if "energy" in model_output_type:
             return DeepPot
         elif "dos" in model_output_type:
@@ -219,7 +210,7 @@ def get_sel_type(self) -> list[int]:
         to the result of the model.
         If returning an empty list, all atom types are selected.
         """
-        return self.metadata["sel_type"]
+        return self._dpmodel.get_sel_type()
 
     def get_numb_dos(self) -> int:
         """Get the number of DOS."""
@@ -364,8 +355,8 @@ def _build_nlist_native(
         nframes = coords.shape[0]
         natoms = coords.shape[1]
         rcut = self.rcut
-        sel = self.metadata["sel"]
-        mixed_types = self.metadata["mixed_types"]
+        sel = self._dpmodel.get_sel()
+        mixed_types = self._dpmodel.mixed_types()
 
         if cells is not None:
             box_input = cells.reshape(nframes, 3, 3)
@@ -476,8 +467,8 @@ def _build_nlist_ase_single(
         nlist : np.ndarray, shape (nloc, nsel)
         mapping : np.ndarray, shape (nall,)
         """
-        sel = self.metadata["sel"]
-        mixed_types = self.metadata["mixed_types"]
+        sel = self._dpmodel.get_sel()
+        mixed_types = self._dpmodel.mixed_types()
         nsel = sum(sel)
 
         natoms = positions.shape[0]
@@ -703,8 +694,8 @@ def _get_output_shape(
             raise RuntimeError("unknown category")
 
     def get_model_def_script(self) -> dict:
-        """Get model definition script."""
-        return self.metadata
+        """Get model definition script (training config)."""
+        return self._model_def_script
 
     def get_model(self) -> torch.nn.Module:
         """Get the exported model module.

deepmd/pt_expt/utils/finetune.py

@@ -33,11 +33,7 @@ def _load_model_params(finetune_model: str) -> dict[str, Any]:
         )
 
         data = serialize_from_file(finetune_model)
-        # Prefer embedded model_params (full config); fall back to
-        # a minimal dict with just type_map for older .pte files.
-        if "model_params" in data:
-            return data["model_params"]
-        return {"type_map": data["model"]["type_map"]}
+        return data["model_def_script"]
     else:
         state_dict = torch.load(finetune_model, map_location=DEVICE, weights_only=True)
         if "model" in state_dict:

deepmd/pt_expt/utils/serialization.py

@@ -165,9 +165,17 @@ def _build_dynamic_shapes(
 
 
 def _collect_metadata(model: torch.nn.Module) -> dict:
-    """Collect metadata from the model for storage in .pte extra_files."""
-    # Serialize the fitting output definitions so that ModelOutputDef
-    # can be reconstructed at inference time without loading the full model.
+    """Collect metadata from the model for C++ inference.
+
+    This metadata is stored as ``metadata.json`` in .pt2 archives and as
+    ``model_def_script.json`` (legacy) in .pte archives.  C++ reads these
+    flat JSON fields because compiling model API methods as AOTInductor
+    entry points is impractical (~12 s per trivial function) and string
+    outputs (``get_type_map``) cannot be expressed as tensor I/O.
+
+    The ``fitting_output_defs`` list is also included so that
+    ``ModelOutputDef`` can be reconstructed without loading the full model.
+    """
     fitting_output_def = model.atomic_output_def()
     fitting_output_defs = []
     for vdef in fitting_output_def.get_data().values():
@@ -189,11 +197,9 @@ def _collect_metadata(model: torch.nn.Module) -> dict:
         "type_map": model.get_type_map(),
         "rcut": model.get_rcut(),
         "sel": model.get_sel(),
-        "model_output_type": model.model_output_type(),
         "dim_fparam": model.get_dim_fparam(),
         "dim_aparam": model.get_dim_aparam(),
         "mixed_types": model.mixed_types(),
-        "sel_type": model.get_sel_type(),
         "has_default_fparam": model.has_default_fparam(),
         "default_fparam": model.get_default_fparam(),
         "fitting_output_defs": fitting_output_defs,
@@ -214,8 +220,8 @@ def serialize_from_file(model_file: str) -> dict:
     -------
     dict
         The serialized model data.  If the archive contains
-        ``model_params.json``, it is included under the
-        ``"model_params"`` key.
+        ``model_def_script.json`` (training config), it is included
+        under the ``"model_def_script"`` key.
     """
     if model_file.endswith(".pt2"):
         return _serialize_from_file_pt2(model_file)
@@ -225,12 +231,14 @@ def serialize_from_file(model_file: str) -> dict:
 
 def _serialize_from_file_pte(model_file: str) -> dict:
     """Serialize a .pte model file to a dictionary."""
-    extra_files = {"model.json": "", "model_params.json": ""}
+    extra_files = {"model.json": "", "model_def_script.json": ""}
     torch.export.load(model_file, extra_files=extra_files)
     model_dict = json.loads(extra_files["model.json"])
     model_dict = _json_to_numpy(model_dict)
-    if extra_files["model_params.json"]:
-        model_dict["model_params"] = json.loads(extra_files["model_params.json"])
+    if extra_files["model_def_script.json"]:
+        model_dict["model_def_script"] = json.loads(
+            extra_files["model_def_script.json"]
+        )
     return model_dict
 
 
@@ -247,20 +255,21 @@ def _serialize_from_file_pt2(model_file: str) -> dict:
                 f"Invalid .pt2 file '{model_file}': missing 'extra/model.json'"
             )
         model_json = zf.read("extra/model.json").decode("utf-8")
-        model_params_json = ""
-        if "extra/model_params.json" in zf.namelist():
-            model_params_json = zf.read("extra/model_params.json").decode("utf-8")
+        model_def_script_json = ""
+        if "extra/model_def_script.json" in zf.namelist():
+            model_def_script_json = zf.read("extra/model_def_script.json").decode(
+                "utf-8"
+            )
     model_dict = json.loads(model_json)
     model_dict = _json_to_numpy(model_dict)
-    if model_params_json:
-        model_dict["model_params"] = json.loads(model_params_json)
+    if model_def_script_json:
+        model_dict["model_def_script"] = json.loads(model_def_script_json)
     return model_dict
 
 
 def deserialize_to_file(
     model_file: str,
     data: dict,
-    model_params: dict | None = None,
     model_json_override: dict | None = None,
 ) -> None:
     """Deserialize a dictionary to a .pte or .pt2 model file.
@@ -275,19 +284,18 @@ def deserialize_to_file(
     data : dict
         The dictionary to be deserialized (same format as dpmodel's
         serialize output, with "model" and optionally "model_def_script" keys).
-    model_params : dict or None
-        Original model config (the dict passed to ``get_model``).
-        If provided, embedded in the .pte so that ``--use-pretrain-script``
-        can extract descriptor/fitting params at finetune time.
+        If ``data["model_def_script"]`` is present, it is embedded in the
+        output so that ``--use-pretrain-script`` can extract descriptor/fitting
+        params at finetune time.
     model_json_override : dict or None
         If provided, this dict is stored in model.json instead of ``data``.
         Used by ``dp compress`` to store the compressed model dict while
         tracing the uncompressed model (make_fx cannot trace custom ops).
     """
     if model_file.endswith(".pt2"):
-        _deserialize_to_file_pt2(model_file, data, model_json_override, model_params)
+        _deserialize_to_file_pt2(model_file, data, model_json_override)
     else:
-        _deserialize_to_file_pte(model_file, data, model_json_override, model_params)
+        _deserialize_to_file_pte(model_file, data, model_json_override)
 
 
 def _trace_and_export(
@@ -387,19 +395,19 @@ def _deserialize_to_file_pte(
     model_file: str,
     data: dict,
     model_json_override: dict | None = None,
-    model_params: dict | None = None,
 ) -> None:
     """Deserialize a dictionary to a .pte model file."""
-    exported, metadata, data_for_json, _output_keys = _trace_and_export(
+    exported, metadata, data_for_json, output_keys = _trace_and_export(
         data, model_json_override
     )
 
+    model_def_script = data.get("model_def_script") or {}
+    metadata["output_keys"] = output_keys
     extra_files = {
-        "model_def_script.json": json.dumps(metadata),
+        "metadata.json": json.dumps(metadata),
+        "model_def_script.json": json.dumps(model_def_script),
         "model.json": json.dumps(data_for_json, separators=(",", ":")),
     }
-    if model_params is not None:
-        extra_files["model_params.json"] = json.dumps(model_params)
 
     torch.export.save(exported, model_file, extra_files=extra_files)
 
@@ -408,7 +416,6 @@ def _deserialize_to_file_pt2(
     model_file: str,
     data: dict,
     model_json_override: dict | None = None,
-    model_params: dict | None = None,
 ) -> None:
     """Deserialize a dictionary to a .pt2 model file (AOTInductor).
 
@@ -430,12 +437,12 @@ def _deserialize_to_file_pt2(
     aoti_compile_and_package(exported, package_path=model_file)
 
     # Embed metadata into the .pt2 ZIP archive
+    model_def_script = data.get("model_def_script") or {}
+    metadata["output_keys"] = output_keys
     with zipfile.ZipFile(model_file, "a") as zf:
-        zf.writestr("extra/model_def_script.json", json.dumps(metadata))
-        zf.writestr("extra/output_keys.json", json.dumps(output_keys))
+        zf.writestr("extra/metadata.json", json.dumps(metadata))
+        zf.writestr("extra/model_def_script.json", json.dumps(model_def_script))
         zf.writestr(
             "extra/model.json",
             json.dumps(data_for_json, separators=(",", ":")),
         )
-        if model_params is not None:
-            zf.writestr("extra/model_params.json", json.dumps(model_params))