feat(pt): full validation support lmdb format

Author: OutisLi

deepmd/dpmodel/utils/__init__.py

@@ -10,6 +10,7 @@
     DistributedSameNlocBatchSampler,
     LmdbDataReader,
     LmdbTestData,
+    LmdbTestDataNlocView,
     SameNlocBatchSampler,
     is_lmdb,
     make_neighbor_stat_data,
@@ -58,6 +59,7 @@
     "FittingNet",
     "LmdbDataReader",
     "LmdbTestData",
+    "LmdbTestDataNlocView",
     "NativeLayer",
     "NativeNet",
     "NetworkCollection",

deepmd/dpmodel/utils/lmdb_data.py

@@ -1286,6 +1286,25 @@ def add(
             "dtype": dtype,
         }
 
+    def add_data_requirement(self, data_requirement: list[DataRequirementItem]) -> None:
+        """Register expected keys from ``DataRequirementItem`` objects.
+
+        Mirrors :meth:`LmdbDataReader.add_data_requirement` so the same
+        requirement list can be forwarded to both the training reader and
+        the full-validation test data.
+        """
+        for item in data_requirement:
+            self.add(
+                item["key"],
+                ndof=item["ndof"],
+                atomic=item["atomic"],
+                must=item["must"],
+                high_prec=item["high_prec"],
+                repeat=item["repeat"],
+                default=item["default"],
+                dtype=item["dtype"],
+            )
+
     def _resolve_dtype(self, key: str) -> np.dtype:
         """Resolve target dtype for a key using registered requirements."""
         if key in self._requirements:
@@ -1444,6 +1463,28 @@ def _stack_frames(
         return result
 
 
+class LmdbTestDataNlocView:
+    """Thin wrapper exposing a fixed-``nloc`` view of :class:`LmdbTestData`.
+
+    The underlying :class:`LmdbTestData` groups frames by atom count. This
+    view fixes one ``nloc`` group, so ``get_test()`` returns only the frames
+    with that atom count and all other attributes (``pbc``, ``mixed_type``,
+    …) are forwarded to the underlying object. It lets downstream consumers
+    that expect a ``DeepmdData``-style system (one fixed natoms per
+    ``get_test()``) work on mixed-nloc LMDB datasets without changes.
+    """
+
+    def __init__(self, lmdb_test_data: "LmdbTestData", nloc: int) -> None:
+        self._inner = lmdb_test_data
+        self._nloc = nloc
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._inner, name)
+
+    def get_test(self) -> dict[str, Any]:
+        return self._inner.get_test(nloc=self._nloc)
+
+
 def merge_lmdb(
     src_paths: list[str],
     dst_path: str,

deepmd/entrypoints/test.py

@@ -17,6 +17,7 @@
 )
 from deepmd.dpmodel.utils.lmdb_data import (
     LmdbTestData,
+    LmdbTestDataNlocView,
     is_lmdb,
 )
 from deepmd.infer.deep_dipole import (
@@ -77,24 +78,6 @@
 log = logging.getLogger(__name__)
 
 
-class _LmdbTestDataNlocView:
-    """Thin wrapper that makes LmdbTestData.get_test() return a specific nloc group.
-
-    Delegates all attributes to the underlying LmdbTestData, but get_test()
-    returns only frames with the specified nloc.
-    """
-
-    def __init__(self, lmdb_test_data: LmdbTestData, nloc: int) -> None:
-        self._inner = lmdb_test_data
-        self._nloc = nloc
-
-    def __getattr__(self, name: str) -> Any:
-        return getattr(self._inner, name)
-
-    def get_test(self) -> dict:
-        return self._inner.get_test(nloc=self._nloc)
-
-
 def test(
     *,
     model: str,
@@ -221,7 +204,7 @@ def test(
             for nloc_val in nloc_keys:
                 label = f"{system} [nloc={nloc_val}]" if len(nloc_keys) > 1 else system
                 # Create a thin wrapper that returns only this nloc group
-                data_items.append((_LmdbTestDataNlocView(lmdb_data, nloc_val), label))
+                data_items.append((LmdbTestDataNlocView(lmdb_data, nloc_val), label))
         else:
             data = DeepmdData(
                 system,

deepmd/pt/train/validation.py

@@ -24,6 +24,10 @@
 import torch.distributed as dist
 
 from deepmd.dpmodel.common import PRECISION_DICT as NP_PRECISION_DICT
+from deepmd.dpmodel.utils.lmdb_data import (
+    LmdbTestData,
+    LmdbTestDataNlocView,
+)
 from deepmd.pt.utils.auto_batch_size import (
     AutoBatchSize,
 )
@@ -35,6 +39,9 @@
     GLOBAL_PT_FLOAT_PRECISION,
     RESERVED_PRECISION_DICT,
 )
+from deepmd.pt.utils.lmdb_dataset import (
+    LmdbDataset,
+)
 from deepmd.pt.utils.utils import (
     to_torch_tensor,
 )
@@ -55,6 +62,10 @@
 log = logging.getLogger(__name__)
 
 if TYPE_CHECKING:
+    from collections.abc import (
+        Iterator,
+    )
+
     from deepmd.utils.data import (
         DeepmdData,
     )
@@ -229,6 +240,12 @@ def __init__(
         if self.rank == 0:
             self._initialize_best_checkpoints(restart_training=restart_training)
 
+        # Lazily-populated full test snapshot for LMDB validation. Mixed-nloc
+        # LMDB datasets cannot be stacked as a single (nframes, natoms*3)
+        # tensor, so we materialize frames grouped by nloc the first time
+        # full validation runs and reuse the snapshot on subsequent calls.
+        self._lmdb_test_data: LmdbTestData | None = None
+
     def should_run(self, display_step: int) -> bool:
         """Check whether the current step should trigger full validation."""
         if not self.enabled or self.start_step is None:
@@ -348,14 +365,10 @@ def evaluate_all_systems(self) -> dict[str, float]:
         if torch.cuda.is_available():
             torch.cuda.empty_cache()
 
-        system_metrics = []
-        for dataset in self.validation_data.systems:
-            if not isinstance(dataset, DeepmdDataSetForLoader):
-                raise TypeError(
-                    "Full validation expects each dataset in validation_data.systems "
-                    f"to be DeepmdDataSetForLoader, got {type(dataset)!r}."
-                )
-            system_metrics.append(self._evaluate_system(dataset.data_system))
+        system_metrics = [
+            self._evaluate_system(data_system)
+            for data_system in self._iter_validation_data_systems()
+        ]
 
         aggregated = weighted_average([metric for metric in system_metrics if metric])
         return {
@@ -364,6 +377,54 @@ def evaluate_all_systems(self) -> dict[str, float]:
             if metric_key in aggregated
         }
 
+    def _iter_validation_data_systems(self) -> Iterator[Any]:
+        """Yield ``DeepmdData``-like systems to evaluate in this run.
+
+        - For ``DpLoaderSet``-style validation data, each entry in
+          ``validation_data.systems`` is a :class:`DeepmdDataSetForLoader`,
+          and we forward its underlying ``DeepmdData`` instance.
+        - For ``LmdbDataset`` validation data, we lazily materialize a
+          :class:`LmdbTestData` snapshot (cached across calls) and yield one
+          :class:`LmdbTestDataNlocView` per ``nloc`` group, so mixed-nloc
+          frames can be stacked and evaluated group by group.
+        """
+        validation_data = self.validation_data
+        if isinstance(validation_data, LmdbDataset):
+            lmdb_test_data = self._get_lmdb_test_data_snapshot(validation_data)
+            for nloc in sorted(lmdb_test_data.nloc_groups.keys()):
+                yield LmdbTestDataNlocView(lmdb_test_data, nloc)
+            return
+
+        for dataset in validation_data.systems:
+            if not isinstance(dataset, DeepmdDataSetForLoader):
+                raise TypeError(
+                    "Full validation expects each dataset in validation_data.systems "
+                    f"to be DeepmdDataSetForLoader, got {type(dataset)!r}."
+                )
+            yield dataset.data_system
+
+    def _get_lmdb_test_data_snapshot(self, lmdb_dataset: LmdbDataset) -> LmdbTestData:
+        """Build (once) and return the cached LMDB test snapshot.
+
+        Reuses the ``type_map`` and previously-registered
+        ``DataRequirementItem`` entries on ``LmdbDataset._reader`` so that
+        the full-validation snapshot sees exactly the same fields and
+        dtypes as training batches.
+        """
+        if self._lmdb_test_data is not None:
+            return self._lmdb_test_data
+
+        reader = lmdb_dataset._reader
+        self._lmdb_test_data = LmdbTestData(
+            lmdb_dataset.lmdb_path,
+            type_map=list(reader._type_map),
+            shuffle_test=False,
+        )
+        data_requirements = list(reader._data_requirements.values())
+        if data_requirements:
+            self._lmdb_test_data.add_data_requirement(data_requirements)
+        return self._lmdb_test_data
+
     def _evaluate_system(
         self, data_system: DeepmdData
     ) -> dict[str, tuple[float, float]]: