pt: refact training code

deepmd/dpmodel/descriptor/make_base_descriptor.py

@@ -89,10 +89,12 @@
         @abstractmethod
         def share_params(self, base_class, shared_level, resume=False):
             """Share the parameters of self to the base_class with shared_level."""
             pass
 
         def compute_input_stats(
-            self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+            self,
+            merged: Union[Callable[[], List[dict]], List[dict]],
+            path: Optional[DPPath] = None,
         ):
             """Update mean and stddev for descriptor elements."""
             raise NotImplementedError

deepmd/dpmodel/model/base_model.py

@@ -10,6 +10,9 @@
     Type,
 )
 
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 from deepmd.utils.plugin import (
     PluginVariant,
     make_plugin_registry,
@@ -93,7 +96,7 @@ def model_output_type(self) -> str:
             """Get the output type for the model."""
 
         @abstractmethod
-        def data_requirement(self) -> dict:
+        def data_requirement(self) -> List[DataRequirementItem]:
             """Get the data requirement for the model."""
 
         @abstractmethod

deepmd/dpmodel/model/dp_model.py

@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from typing import (
+    List,
+)
+
 from deepmd.dpmodel.atomic_model import (
     DPAtomicModel,
 )
@@ -8,6 +12,9 @@
 from deepmd.dpmodel.model.base_model import (
     BaseModel,
 )
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 
 from .make_model import (
     make_model,
@@ -17,6 +24,23 @@
 # use "class" to resolve "Variable not allowed in type expression"
 @BaseModel.register("standard")
 class DPModel(make_model(DPAtomicModel), BaseModel):
-    def data_requirement(self) -> dict:
+    def data_requirement(self) -> List[DataRequirementItem]:
         """Get the data requirement for the model."""
         raise NotImplementedError
+
+    @classmethod
+    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        """Update the selection and perform neighbor statistics.
+
+        Parameters
+        ----------
+        global_jdata : dict
+            The global data, containing the training section
+        local_jdata : dict
+            The local data refer to the current class
+        """
+        local_jdata_cpy = local_jdata.copy()
+        local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel(
+            global_jdata, local_jdata["descriptor"]
+        )
+        return local_jdata_cpy

deepmd/pt/entrypoints/main.py

@@ -50,6 +50,12 @@
 from deepmd.pt.utils.multi_task import (
     preprocess_shared_params,
 )
+from deepmd.utils.argcheck import (
+    normalize,
+)
+from deepmd.utils.compat import (
+    update_deepmd_input,
+)
 from deepmd.utils.path import (
     DPPath,
 )
@@ -99,8 +105,8 @@
         type_split = False
         if model_params_single["descriptor"]["type"] in ["se_e2_a"]:
             type_split = True
         validation_dataset_params = data_dict_single.get("validation_data", None)
         validation_systems = (
             validation_dataset_params["systems"] if validation_dataset_params else None
         )
         training_systems = training_dataset_params["systems"]
@@ -118,7 +124,7 @@
             stat_file_path_single = DPPath(stat_file_path_single, "a")
 
         # validation and training data
         validation_data_single = (
             DpLoaderSet(
                 validation_systems,
                 validation_dataset_params["batch_size"],
@@ -154,7 +160,7 @@
             config["model"], config["training"], config["loss"]
         )
     else:
         train_data, validation_data, stat_file_path = {}, {}, {}
         for model_key in config["model"]["model_dict"]:
             (
                 train_data[model_key],

deepmd/pt/model/descriptor/descriptor.py

@@ -88,8 +88,10 @@
         """Returns the embedding dimension."""
         pass
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for DescriptorBlock elements."""
         raise NotImplementedError

deepmd/pt/model/descriptor/dpa1.py

@@ -147,23 +147,23 @@
         """
         return self.se_atten.mixed_types()
 
     def share_params(self, base_class, shared_level, resume=False):
         assert (
             self.__class__ == base_class.__class__
         ), "Only descriptors of the same type can share params!"
         # For DPA1 descriptors, the user-defined share-level
         # shared_level: 0
         # share all parameters in both type_embedding and se_atten
         if shared_level == 0:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
             self.se_atten.share_params(base_class.se_atten, 0, resume=resume)
         # shared_level: 1
         # share all parameters in type_embedding
         elif shared_level == 1:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
         # Other shared levels
         else:
             raise NotImplementedError
 
     @property
     def dim_out(self):
@@ -173,8 +173,10 @@
     def dim_emb(self):
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         return self.se_atten.compute_input_stats(merged, path)
 

deepmd/pt/model/descriptor/dpa2.py

@@ -291,40 +291,40 @@
         """
         return True
 
     def share_params(self, base_class, shared_level, resume=False):
         assert (
             self.__class__ == base_class.__class__
         ), "Only descriptors of the same type can share params!"
         # For DPA2 descriptors, the user-defined share-level
         # shared_level: 0
         # share all parameters in type_embedding, repinit and repformers
         if shared_level == 0:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
             self.repinit.share_params(base_class.repinit, 0, resume=resume)
             self._modules["g1_shape_tranform"] = base_class._modules[
                 "g1_shape_tranform"
             ]
             self.repformers.share_params(base_class.repformers, 0, resume=resume)
         # shared_level: 1
         # share all parameters in type_embedding and repinit
         elif shared_level == 1:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
             self.repinit.share_params(base_class.repinit, 0, resume=resume)
         # shared_level: 2
         # share all parameters in type_embedding and repformers
         elif shared_level == 2:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
             self._modules["g1_shape_tranform"] = base_class._modules[
                 "g1_shape_tranform"
             ]
             self.repformers.share_params(base_class.repformers, 0, resume=resume)
         # shared_level: 3
         # share all parameters in type_embedding
         elif shared_level == 3:
             self._modules["type_embedding"] = base_class._modules["type_embedding"]
         # Other shared levels
         else:
             raise NotImplementedError
 
     @property
     def dim_out(self):
@@ -335,11 +335,13 @@
         """Returns the embedding dimension g2."""
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         for ii, descrpt in enumerate([self.repinit, self.repformers]):
             descrpt.compute_input_stats(merged, path)
 
     def serialize(self) -> dict:
         """Serialize the obj to dict."""

deepmd/pt/model/descriptor/hybrid.py

@@ -157,13 +157,15 @@
         else:
             raise NotImplementedError
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         for ii, descrpt in enumerate(self.descriptor_list):
             # need support for hybrid descriptors
             descrpt.compute_input_stats(merged, path)
 
     def forward(
         self,

deepmd/pt/model/descriptor/repformers.py

@@ -280,22 +280,24 @@
 
         return g1, g2, h2, rot_mat.view(-1, nloc, self.dim_emb, 3), sw
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         if path is None or not path.is_dir():
             if callable(merged):
                 # only get data for once
                 sampled = merged()
             else:
                 sampled = merged
         else:
             sampled = []
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_a.py

@@ -129,26 +129,28 @@
         """
         return self.sea.mixed_types()
 
     def share_params(self, base_class, shared_level, resume=False):
         assert (
             self.__class__ == base_class.__class__
         ), "Only descriptors of the same type can share params!"
         # For SeA descriptors, the user-defined share-level
         # shared_level: 0
         # share all parameters in sea
         if shared_level == 0:
             self.sea.share_params(base_class.sea, 0, resume=resume)
         # Other shared levels
         else:
             raise NotImplementedError
 
     @property
     def dim_out(self):
         """Returns the output dimension of this descriptor."""
         return self.sea.dim_out
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         return self.sea.compute_input_stats(merged, path)
@@ -428,22 +430,24 @@
         else:
             raise KeyError(key)
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         if path is None or not path.is_dir():
             if callable(merged):
                 # only get data for once
                 sampled = merged()
             else:
                 sampled = merged
         else:
             sampled = []
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_atten.py

@@ -202,22 +202,24 @@
         """Returns the output dimension of embedding."""
         return self.get_dim_emb()
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         if path is None or not path.is_dir():
             if callable(merged):
                 # only get data for once
                 sampled = merged()
             else:
                 sampled = merged
         else:
             sampled = []
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/descriptor/se_r.py

@@ -153,50 +153,52 @@
         """
         return False
 
     def share_params(self, base_class, shared_level, resume=False):
         assert (
             self.__class__ == base_class.__class__
         ), "Only descriptors of the same type can share params!"
         # For SeR descriptors, the user-defined share-level
         # shared_level: 0
         if shared_level == 0:
             # link buffers
             if hasattr(self, "mean") and not resume:
                 # in case of change params during resume
                 base_env = EnvMatStatSe(base_class)
                 base_env.stats = base_class.stats
                 for kk in base_class.get_stats():
                     base_env.stats[kk] += self.get_stats()[kk]
                 mean, stddev = base_env()
                 if not base_class.set_davg_zero:
                     base_class.mean.copy_(torch.tensor(mean, device=env.DEVICE))
                 base_class.stddev.copy_(torch.tensor(stddev, device=env.DEVICE))
                 self.mean = base_class.mean
                 self.stddev = base_class.stddev
             # self.load_state_dict(base_class.state_dict()) # this does not work, because it only inits the model
             # the following will successfully link all the params except buffers
             for item in self._modules:
                 self._modules[item] = base_class._modules[item]
         # Other shared levels
         else:
             raise NotImplementedError
 
     def compute_input_stats(
-        self, merged: Union[Callable, List[dict]], path: Optional[DPPath] = None
+        self,
+        merged: Union[Callable[[], List[dict]], List[dict]],
+        path: Optional[DPPath] = None,
     ):
         """Update mean and stddev for descriptor elements."""
         env_mat_stat = EnvMatStatSe(self)
         if path is not None:
             path = path / env_mat_stat.get_hash()
         if path is None or not path.is_dir():
             if callable(merged):
                 # only get data for once
                 sampled = merged()
             else:
                 sampled = merged
         else:
             sampled = []
         env_mat_stat.load_or_compute_stats(sampled, path)
         self.stats = env_mat_stat.stats
         mean, stddev = env_mat_stat()
         if not self.set_davg_zero:

deepmd/pt/model/model/dipole_model.py

@@ -1,11 +1,16 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Dict,
+    List,
     Optional,
 )
 
 import torch
 
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
+
 from .dp_model import (
     DPModel,
 )
@@ -91,22 +96,24 @@
             model_predict = model_ret
         return model_predict
 
     @property
-    def data_requirement(self):
-        data_requirement = {
-            "dipole": {
-                "ndof": 3,
-                "atomic": False,
-                "must": False,
-                "high_prec": False,
-                "type_sel": self.get_sel_type(),
-            },
-            "atomic_dipole": {
-                "ndof": 3,
-                "atomic": True,
-                "must": False,
-                "high_prec": False,
-                "type_sel": self.get_sel_type(),
-            },
-        }
+    def data_requirement(self) -> List[DataRequirementItem]:
+        data_requirement = [
+            DataRequirementItem(
+                "dipole",
+                ndof=3,
+                atomic=False,
+                must=False,
+                high_prec=False,
+                type_sel=self.get_sel_type(),
+            ),
+            DataRequirementItem(
+                "atomic_dipole",
+                ndof=3,
+                atomic=True,
+                must=False,
+                high_prec=False,
+                type_sel=self.get_sel_type(),
+            ),
+        ]
         return data_requirement

deepmd/pt/model/model/dp_zbl_model.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 from typing import (
     Dict,
+    List,
     Optional,
 )
 
@@ -15,6 +16,9 @@
 from deepmd.pt.model.model.model import (
     BaseModel,
 )
+from deepmd.utils.data import (
+    DataRequirementItem,
+)
 
 from .make_model import (
     make_model,
@@ -101,39 +105,61 @@
         model_predict = model_ret
         return model_predict
 
     @property
-    def data_requirement(self):
-        data_requirement = {
-            "energy": {
-                "ndof": 1,
-                "atomic": False,
-                "must": False,
-                "high_prec": True,
-            },
-            "force": {
-                "ndof": 3,
-                "atomic": True,
-                "must": False,
-                "high_prec": False,
-            },
-            "virial": {
-                "ndof": 9,
-                "atomic": False,
-                "must": False,
-                "high_prec": False,
-            },
-            "atom_ener": {
-                "ndof": 1,
-                "atomic": True,
-                "must": False,
-                "high_prec": False,
-            },
-            "atom_pref": {
-                "ndof": 1,
-                "atomic": True,
-                "must": False,
-                "high_prec": False,
-                "repeat": 3,
-            },
-        }
+    def data_requirement(self) -> List[DataRequirementItem]:
+        data_requirement = [
+            DataRequirementItem(
+                "energy",
+                ndof=1,
+                atomic=False,
+                must=False,
+                high_prec=True,
+            ),
+            DataRequirementItem(
+                "force",
+                ndof=3,
+                atomic=True,
+                must=False,
+                high_prec=False,
+            ),
+            DataRequirementItem(
+                "virial",
+                ndof=9,
+                atomic=False,
+                must=False,
+                high_prec=False,
+            ),
+            DataRequirementItem(
+                "atom_ener",
+                ndof=1,
+                atomic=True,
+                must=False,
+                high_prec=False,
+            ),
+            DataRequirementItem(
+                "atom_pref",
+                ndof=1,
+                atomic=True,
+                must=False,
+                high_prec=False,
+                repeat=3,
+            ),
+        ]
         return data_requirement
+
+    @classmethod
+    def update_sel(cls, global_jdata: dict, local_jdata: dict):
+        """Update the selection and perform neighbor statistics.
+
+        Parameters
+        ----------
+        global_jdata : dict
+            The global data, containing the training section
+        local_jdata : dict
+            The local data refer to the current class
+        """
+        local_jdata_cpy = local_jdata.copy()
+        local_jdata_cpy["dpmodel"] = DPModel.update_sel(
+            global_jdata, local_jdata["dpmodel"]
+        )
+        return local_jdata_cpy