diff --git a/deepmd/main.py b/deepmd/main.py index 1206b2fa0e..92377ed5e6 100644 --- a/deepmd/main.py +++ b/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", diff --git a/deepmd/pt_expt/entrypoints/compress.py b/deepmd/pt_expt/entrypoints/compress.py index 3d9bb6506a..92dfa947a9 100644 --- a/deepmd/pt_expt/entrypoints/compress.py +++ b/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,6 +14,33 @@ 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, @@ -18,6 +48,7 @@ def enable_compression( 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, ) diff --git a/deepmd/pt_expt/entrypoints/main.py b/deepmd/pt_expt/entrypoints/main.py index 66ddd1de79..5e9f6addf3 100644 --- a/deepmd/pt_expt/entrypoints/main.py +++ b/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( diff --git a/deepmd/pt_expt/utils/neighbor_stat.py b/deepmd/pt_expt/utils/neighbor_stat.py index 850bf86def..36f4788cc0 100644 --- a/deepmd/pt_expt/utils/neighbor_stat.py +++ b/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, diff --git a/source/tests/pt_expt/test_compress_min_nbor_dist.py b/source/tests/pt_expt/test_compress_min_nbor_dist.py new file mode 100644 index 0000000000..7f090d2fe3 --- /dev/null +++ b/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