fix(pt_expt): reuse the stored min_nbor_dist and batch the neighbor statistics

deepmd/main.py

@@ -668,6 +668,13 @@ def main_parser() -> argparse.ArgumentParser:
         default=None,
         help="The training script of the input frozen model",
     )
+    parser_compress.add_argument(
+        "--recompute-min-nbor-dist",
+        action="store_true",
+        help="(Supported backend: PyTorch Exportable) Ignore the minimal neighbor "
+        "distance saved in the model and recompute it from the training data. "
+        "Requires -t,--training-script",
+    )
     parser_compress.add_argument(
         "--head",
         "--model-branch",

deepmd/pt_expt/entrypoints/compress.py

@@ -2,6 +2,9 @@
 """Compress a pt_expt model (.pte) by tabulating embedding nets."""
 
 import logging
+from typing import (
+    Any,
+)
 
 from deepmd.pt_expt.utils.serialization import (
     deserialize_to_file,
@@ -11,13 +14,41 @@
 log = logging.getLogger(__name__)
 
 
+def _read_saved_min_nbor_dist(model: Any, model_dict: dict) -> tuple[float | None, str]:
+    """Read the stored minimal neighbor distance and where it was read from.
+
+    ``@variables`` is the cross-backend location of this value: it is written
+    by :mod:`deepmd.pt.utils.serialization` and read back by the PyTorch and
+    Paddle backends, so a ``.pt2`` produced by ``dp convert-backend`` carries
+    the value there rather than inside the serialized model dict.
+
+    Returns
+    -------
+    float or None
+        The stored minimal neighbor distance, None if the model has none.
+    str
+        Human-readable description of where the value was read from.
+    """
+    min_nbor_dist = model.get_min_nbor_dist()
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model"
+    min_nbor_dist = model_dict.get("min_nbor_dist")
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model file"
+    min_nbor_dist = (model_dict.get("@variables") or {}).get("min_nbor_dist")
+    if min_nbor_dist is not None:
+        return float(min_nbor_dist), "the model file (@variables)"
+    return None, ""
+
+
 def enable_compression(
     input_file: str,
     output: str,
     stride: float = 0.01,
     extrapolate: int = 5,
     check_frequency: int = -1,
     training_script: str | None = None,
+    recompute_min_nbor_dist: bool = False,
 ) -> None:
     """Compress a .pte model by tabulating embedding nets.
 
@@ -36,6 +67,9 @@ def enable_compression(
     training_script : str or None
         Path to training script, used to compute min_nbor_dist if not
         stored in the model.
+    recompute_min_nbor_dist : bool
+        Ignore the min_nbor_dist stored in the model and recompute it from
+        the training data. Requires training_script.
     """
     from deepmd.pt_expt.model.model import (
         BaseModel,
@@ -46,20 +80,34 @@ def enable_compression(
     model = BaseModel.deserialize(model_dict["model"])
 
     # 2. Get or compute min_nbor_dist
-    min_nbor_dist = model.get_min_nbor_dist()
-    if min_nbor_dist is None:
-        min_nbor_dist = model_dict.get("min_nbor_dist")
-    if min_nbor_dist is None:
-        log.info(
-            "Minimal neighbor distance is not saved in the model, "
-            "compute it from the training data."
-        )
-        if training_script is None:
-            raise ValueError(
-                "The model does not have a minimum neighbor distance, "
-                "so the training script and data must be provided "
-                "(via -t,--training-script)."
+    if recompute_min_nbor_dist:
+        min_nbor_dist, source = None, ""
+    else:
+        min_nbor_dist, source = _read_saved_min_nbor_dist(model, model_dict)
+    if min_nbor_dist is not None:
+        log.info(f"Minimal neighbor distance read from {source}: {min_nbor_dist:f}")
+    else:
+        if recompute_min_nbor_dist:
+            log.info(
+                "Recompute the minimal neighbor distance from the training data, "
+                "ignoring the one saved in the model."
+            )
+            if training_script is None:
+                raise ValueError(
+                    "Recomputing the minimal neighbor distance requires the "
+                    "training script and data (via -t,--training-script)."
+                )
+        else:
+            log.info(
+                "Minimal neighbor distance is not saved in the model, "
+                "compute it from the training data."
             )
+            if training_script is None:
+                raise ValueError(
+                    "The model does not have a minimum neighbor distance, "
+                    "so the training script and data must be provided "
+                    "(via -t,--training-script)."
+                )
         from deepmd.common import (
             j_loader,
         )

deepmd/pt_expt/entrypoints/main.py

@@ -928,6 +928,7 @@ def main(args: list[str] | argparse.Namespace | None = None) -> None:
             extrapolate=FLAGS.extrapolate,
             check_frequency=FLAGS.frequency,
             training_script=FLAGS.training_script,
+            recompute_min_nbor_dist=FLAGS.recompute_min_nbor_dist,
         )
     else:
         raise RuntimeError(

deepmd/pt_expt/utils/neighbor_stat.py

@@ -7,6 +7,9 @@
 import torch
 
 from deepmd.dpmodel.utils.neighbor_stat import NeighborStatOP as NeighborStatOPDP
+from deepmd.pt.utils.auto_batch_size import (
+    AutoBatchSize,
+)
 from deepmd.pt_expt.common import (
     torch_module,
 )
@@ -28,6 +31,12 @@ class NeighborStatOP(NeighborStatOPDP):
 class NeighborStat(BaseNeighborStat):
     """Neighbor statistics using torch on DEVICE.
 
+    The statistics are evaluated one frame batch at a time. The intermediate
+    tensor is of shape ``[nframes, nloc, nall, 3]``, so processing a whole set
+    at once needs hundreds of GiB for a large set; :class:`AutoBatchSize`
+    keeps the batch within the available device memory, as the PyTorch,
+    Paddle, JAX and TensorFlow backends already do.
+
     Parameters
     ----------
     ntypes : int
@@ -46,6 +55,7 @@ def __init__(
     ) -> None:
         super().__init__(ntypes, rcut, mixed_type)
         self.op = NeighborStatOP(ntypes, rcut, mixed_type)
+        self.auto_batch_size = AutoBatchSize()
 
     def iterator(
         self, data: DeepmdDataSystem
@@ -65,7 +75,10 @@ def iterator(
             for jj in data.data_systems[ii].dirs:
                 data_set = data.data_systems[ii]
                 data_set_data = data_set._load_set(jj)
-                minrr2, max_nnei = self._execute(
+                minrr2, max_nnei = self.auto_batch_size.execute_all(
+                    self._execute,
+                    data_set_data["coord"].shape[0],
+                    data_set.get_natoms(),
                     data_set_data["coord"],
                     data_set_data["type"],
                     data_set_data["box"] if data_set.pbc else None,

source/tests/pt_expt/test_compress_min_nbor_dist.py

@@ -0,0 +1,70 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+"""Tests for reading the stored minimal neighbor distance in pt_expt compress."""
+
+from typing import (
+    Any,
+)
+
+import pytest
+
+from deepmd.main import (
+    parse_args,
+)
+from deepmd.pt_expt.entrypoints.compress import (
+    _read_saved_min_nbor_dist,
+)
+
+
+class _FakeModel:
+    """Stand-in exposing only the getter that the reader calls."""
+
+    def __init__(self, min_nbor_dist: float | None) -> None:
+        self._min_nbor_dist = min_nbor_dist
+
+    def get_min_nbor_dist(self) -> float | None:
+        return self._min_nbor_dist
+
+
+@pytest.mark.parametrize(
+    ("model_min_nbor_dist", "model_dict", "expected"),
+    [
+        # the model buffer wins over both metadata locations
+        (
+            1.5,
+            {"min_nbor_dist": 9.0, "@variables": {"min_nbor_dist": 8.0}},
+            (1.5, "the model"),
+        ),
+        # the top-level key, written by pt_expt's own compress
+        (
+            None,
+            {"min_nbor_dist": 2.5, "@variables": {"min_nbor_dist": 8.0}},
+            (2.5, "the model file"),
+        ),
+        # "@variables", the cross-backend location that dp convert-backend writes
+        (
+            None,
+            {"@variables": {"min_nbor_dist": 3.5}},
+            (3.5, "the model file (@variables)"),
+        ),
+        # nothing stored anywhere
+        (None, {}, (None, "")),
+        (None, {"@variables": {}}, (None, "")),
+        (None, {"@variables": None}, (None, "")),
+    ],
+)
+def test_read_saved_min_nbor_dist(
+    model_min_nbor_dist: float | None,
+    model_dict: dict[str, Any],
+    expected: tuple[float | None, str],
+) -> None:
+    """Every storage location is honored, in order of precedence."""
+    assert _read_saved_min_nbor_dist(_FakeModel(model_min_nbor_dist), model_dict) == (
+        expected
+    )
+
+
+def test_recompute_min_nbor_dist_flag_defaults_to_false() -> None:
+    """The compress parser exposes the flag and leaves it off by default."""
+    args = ["--pt-expt", "compress", "-i", "in.pt2", "-o", "out.pt2"]
+    assert parse_args(args).recompute_min_nbor_dist is False
+    assert parse_args([*args, "--recompute-min-nbor-dist"]).recompute_min_nbor_dist