From e97bf965411162a143cbf6eff9857571bedb790c Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 06:14:12 -0400 Subject: [PATCH 01/12] refactor: refactor update_sel and save min_nbor_dist Fix #3525. Fix #3544. Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/descriptor/dpa1.py | 29 +++- deepmd/dpmodel/descriptor/dpa2.py | 47 ++++-- deepmd/dpmodel/descriptor/hybrid.py | 40 ++++- .../descriptor/make_base_descriptor.py | 26 +++- deepmd/dpmodel/descriptor/se_e2_a.py | 28 +++- deepmd/dpmodel/descriptor/se_r.py | 29 +++- deepmd/dpmodel/model/base_model.py | 37 ++++- deepmd/dpmodel/model/dp_model.py | 35 ++++- deepmd/dpmodel/model/spin_model.py | 4 + deepmd/dpmodel/utils/update_sel.py | 4 - deepmd/pt/entrypoints/main.py | 44 +++++- deepmd/pt/model/descriptor/dpa1.py | 29 +++- deepmd/pt/model/descriptor/dpa2.py | 47 ++++-- deepmd/pt/model/descriptor/hybrid.py | 40 ++++- deepmd/pt/model/descriptor/se_a.py | 28 +++- deepmd/pt/model/descriptor/se_r.py | 28 +++- deepmd/pt/model/model/dp_model.py | 35 ++++- deepmd/pt/model/model/dp_zbl_model.py | 31 +++- deepmd/pt/model/model/frozen.py | 31 +++- deepmd/pt/model/model/model.py | 6 + deepmd/pt/model/model/spin_model.py | 5 + deepmd/pt/utils/serialization.py | 7 +- deepmd/pt/utils/update_sel.py | 4 - deepmd/tf/descriptor/descriptor.py | 25 +++- deepmd/tf/descriptor/hybrid.py | 39 ++++- deepmd/tf/descriptor/loc_frame.py | 25 +++- deepmd/tf/descriptor/se.py | 29 +++- deepmd/tf/descriptor/se_a_mask.py | 25 +++- deepmd/tf/descriptor/se_atten.py | 29 +++- deepmd/tf/entrypoints/compress.py | 13 +- deepmd/tf/entrypoints/train.py | 18 ++- deepmd/tf/model/frozen.py | 26 +++- deepmd/tf/model/linear.py | 40 ++++- deepmd/tf/model/model.py | 44 ++++-- deepmd/tf/model/pairtab.py | 25 +++- deepmd/tf/model/pairwise_dprc.py | 29 +++- deepmd/tf/utils/update_sel.py | 15 -- deepmd/utils/update_sel.py | 140 +++++++----------- doc/development/create-a-model-pt.md | 8 +- 39 files changed, 872 insertions(+), 272 deletions(-) diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/deepmd/dpmodel/descriptor/dpa1.py index 138cab82d6..42929ca1d3 100644 --- a/deepmd/dpmodel/descriptor/dpa1.py +++ b/deepmd/dpmodel/descriptor/dpa1.py @@ -14,6 +14,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -524,18 +527,36 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist @DescriptorBlock.register("se_atten") diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 78bf174685..9243883ebe 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -15,6 +15,9 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -743,30 +746,46 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy["repinit"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repinit"], + min_nbor_dist, repinit_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repinit"]["rcut"], + local_jdata_cpy["repinit"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - local_jdata_cpy["repformer"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repformer"], + local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0] + min_nbor_dist, repformer_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repformer"]["rcut"], + local_jdata_cpy["repformer"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - return local_jdata_cpy + local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/hybrid.py b/deepmd/dpmodel/descriptor/hybrid.py index 15825ecc10..d359bf6fb7 100644 --- a/deepmd/dpmodel/descriptor/hybrid.py +++ b/deepmd/dpmodel/descriptor/hybrid.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -19,6 +20,9 @@ from deepmd.dpmodel.utils.nlist import ( nlist_distinguish_types, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -222,22 +226,42 @@ def call( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - BaseDescriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self) -> dict: return { diff --git a/deepmd/dpmodel/descriptor/make_base_descriptor.py b/deepmd/dpmodel/descriptor/make_base_descriptor.py index cba9eebe4b..a4fc8bddf9 100644 --- a/deepmd/dpmodel/descriptor/make_base_descriptor.py +++ b/deepmd/dpmodel/descriptor/make_base_descriptor.py @@ -7,12 +7,16 @@ Callable, List, Optional, + Tuple, Union, ) from deepmd.common import ( j_get_type, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -149,19 +153,33 @@ def deserialize(cls, data: dict) -> "BD": @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # call subprocess cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) setattr(BD, fwd_method_name, BD.fwd) delattr(BD, "fwd") diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 193383ac4f..6feffc114f 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -9,6 +9,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -422,15 +425,32 @@ def deserialize(cls, data: dict) -> "DescrptSeA": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 5175b91ae1..82e9589465 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -4,6 +4,9 @@ from deepmd.dpmodel.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -21,6 +24,7 @@ Any, List, Optional, + Tuple, ) from deepmd.dpmodel import ( @@ -345,15 +349,32 @@ def deserialize(cls, data: dict) -> "DescrptSeR": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index b416e57fb8..354fa58348 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -7,9 +7,14 @@ from typing import ( Any, List, + Optional, + Tuple, Type, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.plugin import ( PluginVariant, make_plugin_registry, @@ -129,12 +134,21 @@ def deserialize(cls, data: dict) -> "BaseBaseModel": raise NotImplementedError(f"Not implemented in class {cls.__name__}") model_def_script: str + """The model definition script.""" + min_nbor_dist: Optional[float] + """The minimum distance between two atoms. Used for model compression. + None when skipping neighbor statistics. + """ @abstractmethod def get_model_def_script(self) -> str: """Get the model definition script.""" pass + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum distance between two atoms.""" + return self.min_nbor_dist + @abstractmethod def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" @@ -148,22 +162,36 @@ def get_nsel(self) -> int: @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # getting model type based on fitting type model_type = local_jdata.get("type", "standard") if model_type == "standard": model_type = local_jdata.get("fitting", {}).get("type", "ener") cls = cls.get_class_by_type(model_type) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) return BaseBaseModel @@ -186,6 +214,7 @@ class BaseModel(make_base_model()): def __init__(self) -> None: self.model_def_script = "" + self.min_nbor_dist = None def get_model_def_script(self) -> str: """Get the model definition script.""" diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 37cb426ab7..1597ba0b14 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -1,26 +1,49 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, + Optional, + Tuple, +) + from deepmd.dpmodel.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) # use "class" to resolve "Variable not allowed in type expression" class DPModelCommon: @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 90e2bb3fb4..09bbabe358 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -227,6 +227,10 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() + def get_min_nbor_dist(self) -> float: + """Get the minimum neighbor distance.""" + return self.backbone_model.get_min_nbor_dist() + def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" # for C++ interface diff --git a/deepmd/dpmodel/utils/update_sel.py b/deepmd/dpmodel/utils/update_sel.py index 48463b5743..dc38a6a041 100644 --- a/deepmd/dpmodel/utils/update_sel.py +++ b/deepmd/dpmodel/utils/update_sel.py @@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # TODO: save to the model in UpdateSel.hook - pass diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index eafce67e84..62c0cefa1b 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -63,6 +63,7 @@ update_deepmd_input, ) from deepmd.utils.data_system import ( + get_data, process_systems, ) from deepmd.utils.path import ( @@ -252,21 +253,56 @@ def train(FLAGS): log.info( "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) + + type_map = config["model"].get("type_map") if not multi_task: - config["model"] = BaseModel.update_sel(config, config["model"]) + train_data = get_data( + config["training"]["training_data"], 0, type_map, None + ) + config["model"], min_nbor_dist = BaseModel.update_sel( + config, config["model"] + ) else: + assert ( + type_map is not None + ), "Data stat in multi-task mode must have available type_map! " + train_data = None + for systems in config["training"]["data_dict"]: + tmp_data = get_data( + config["training"]["data_dict"][systems]["training_data"], + 0, + type_map, + None, + ) + tmp_data.get_batch() + assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " + if train_data is None: + train_data = tmp_data + else: + train_data.system_dirs += tmp_data.system_dirs + train_data.data_systems += tmp_data.data_systems + train_data.natoms += tmp_data.natoms + train_data.natoms_vec += tmp_data.natoms_vec + train_data.default_mesh += tmp_data.default_mesh + training_jdata = deepcopy(config["training"]) training_jdata.pop("data_dict", {}) training_jdata.pop("model_prob", {}) + # considering multi-task shares the descriptor, we need a minimal min_nbor_dist + min_nbor_dist = 0 for model_item in config["model"]["model_dict"]: fake_global_jdata = { "model": deepcopy(config["model"]["model_dict"][model_item]), "training": deepcopy(config["training"]["data_dict"][model_item]), } fake_global_jdata["training"].update(training_jdata) - config["model"]["model_dict"][model_item] = BaseModel.update_sel( - fake_global_jdata, config["model"]["model_dict"][model_item] + config["model"]["model_dict"][model_item], min_nbor_dist_item = ( + BaseModel.update_sel( + fake_global_jdata, config["model"]["model_dict"][model_item] + ) ) + if min_nbor_dist_item is not None: + min_nbor_dist = min(min_nbor_dist, min_nbor_dist_item) with open(FLAGS.output, "w") as fp: json.dump(config, fp, indent=4) @@ -281,6 +317,8 @@ def train(FLAGS): FLAGS.init_frz_model, shared_links=shared_links, ) + # save min_nbor_dist + trainer.model.min_nbor_dist = min_nbor_dist trainer.run() diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 4ab39465dc..92f5cf2e15 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -27,6 +27,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -549,15 +552,33 @@ def forward( return g1, rot_mat, g2, h2, sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index 678b797e6c..b33a528721 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -37,6 +37,9 @@ from deepmd.pt.utils.utils import ( to_numpy_array, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -573,30 +576,46 @@ def forward( return g1, rot_mat, g2, h2, sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy["repinit"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repinit"], + min_nbor_dist, repinit_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repinit"]["rcut"], + local_jdata_cpy["repinit"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - local_jdata_cpy["repformer"] = update_sel.update_one_sel( - global_jdata, - local_jdata_cpy["repformer"], + local_jdata_cpy["repinit"]["nsel"] = repinit_sel[0] + min_nbor_dist, repformer_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repformer"]["rcut"], + local_jdata_cpy["repformer"]["nsel"], True, - rcut_key="rcut", - sel_key="nsel", ) - return local_jdata_cpy + local_jdata_cpy["repformer"]["nsel"] = repformer_sel[0] + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/hybrid.py b/deepmd/pt/model/descriptor/hybrid.py index bce52fd8ca..e202005f4c 100644 --- a/deepmd/pt/model/descriptor/hybrid.py +++ b/deepmd/pt/model/descriptor/hybrid.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -20,6 +21,9 @@ from deepmd.pt.utils.utils import ( to_torch_tensor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -242,22 +246,42 @@ def forward( return out_descriptor, out_gr, out_g2, out_h2, out_sw @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - BaseDescriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = BaseDescriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self) -> dict: return { diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 0035eddba6..6293351d41 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -30,6 +30,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -295,18 +298,35 @@ def t_cvt(xx): return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + return local_jdata_cpy, min_nbor_dist @DescriptorBlock.register("se_e2_a") diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 340fd0c02a..0f030733d1 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -35,6 +35,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -413,15 +416,32 @@ def t_cvt(xx): return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index fab1ff580f..d3a65db287 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -1,28 +1,51 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, + Optional, + Tuple, +) + from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) class DPModelCommon: """A base class to implement common methods for all the Models.""" @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = BaseDescriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = BaseDescriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist def get_fitting_net(self): """Get the fitting network.""" diff --git a/deepmd/pt/model/model/dp_zbl_model.py b/deepmd/pt/model/model/dp_zbl_model.py index f18e1d097f..45b76d1da6 100644 --- a/deepmd/pt/model/model/dp_zbl_model.py +++ b/deepmd/pt/model/model/dp_zbl_model.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Dict, + List, Optional, + Tuple, ) import torch @@ -12,6 +14,9 @@ from deepmd.pt.model.model.model import ( BaseModel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .dp_model import ( DPModelCommon, @@ -105,18 +110,32 @@ def forward_lower( return model_predict @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["dpmodel"] = DPModelCommon.update_sel( - global_jdata, local_jdata["dpmodel"] + local_jdata_cpy["dpmodel"], min_nbor_dist = DPModelCommon.update_sel( + train_data, type_map, local_jdata["dpmodel"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index e3dcd389bb..4059450a2d 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -5,6 +5,7 @@ Dict, List, Optional, + Tuple, ) import torch @@ -18,6 +19,9 @@ from deepmd.pt.model.model.model import ( BaseModel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @BaseModel.register("frozen") @@ -130,6 +134,11 @@ def get_model_def_script(self) -> str: # be a problem return self.model.get_model_def_script() + @torch.jit.export + def get_min_nbor_dist(self) -> float: + """Get the minimum neighbor distance.""" + return self.model.get_min_nbor_dist() + def serialize(self) -> dict: from deepmd.pt.model.model import ( get_model, @@ -156,17 +165,31 @@ def get_nsel(self) -> int: return self.model.get_nsel() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None @torch.jit.export def model_output_type(self) -> str: diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index bf97472e33..d3670737ba 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -18,6 +18,7 @@ def __init__(self, *args, **kwargs): """Construct a basic model for different tasks.""" torch.nn.Module.__init__(self) self.model_def_script = "" + self.min_nbor_dist = None def compute_or_load_stat( self, @@ -46,6 +47,11 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.model_def_script + @torch.jit.export + def get_min_nbor_dist(self) -> Optional[float]: + """Get the minimum distance between two atoms.""" + return self.min_nbor_dist + @torch.jit.export def get_ntypes(self): """Returns the number of element types.""" diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index ff932bc402..601d5fcd70 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -291,6 +291,11 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() + @torch.jit.export + def get_min_nbor_dist(self) -> float: + """Get the minimum neighbor distance.""" + return self.backbone_model.get_min_nbor_dist() + @torch.jit.export def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index c99ddbb3c6..7ab65ab426 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -50,8 +50,9 @@ def serialize_from_file(model_file: str) -> dict: "pt_version": torch.__version__, "model": model_dict, "model_def_script": model_def_script, - # TODO - "@variables": {}, + "@variables": { + "min_nbor_dist": model.get_min_nbor_dist(), + }, } return data @@ -72,4 +73,6 @@ def deserialize_to_file(model_file: str, data: dict) -> None: # JIT will happy in this way... model.model_def_script = json.dumps(data["model_def_script"]) model = torch.jit.script(model) + if "min_nbor_dist" in data.get("@variables", {}): + model.min_nbor_dist = data["@variables"]["min_nbor_dist"] torch.jit.save(model, model_file) diff --git a/deepmd/pt/utils/update_sel.py b/deepmd/pt/utils/update_sel.py index 8c2d0699f2..7f42a9f91c 100644 --- a/deepmd/pt/utils/update_sel.py +++ b/deepmd/pt/utils/update_sel.py @@ -15,7 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # TODO: save to the model in UpdateSel.hook - pass diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index fabaf78c85..2f813da731 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -26,6 +26,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.plugin import ( make_plugin_registry, ) @@ -462,19 +465,33 @@ def explicit_ntypes(self) -> bool: @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # call subprocess cls = cls.get_class_by_type(j_get_type(local_jdata, cls.__name__)) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": diff --git a/deepmd/tf/descriptor/hybrid.py b/deepmd/tf/descriptor/hybrid.py index 7c69ea5202..fe4fc2ae6a 100644 --- a/deepmd/tf/descriptor/hybrid.py +++ b/deepmd/tf/descriptor/hybrid.py @@ -17,6 +17,9 @@ from deepmd.tf.utils.spin import ( Spin, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -423,22 +426,42 @@ def explicit_ntypes(self) -> bool: return any(ii.explicit_ntypes for ii in self.descrpt_list) @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["list"] = [ - Descriptor.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["list"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["list"]: + new_sub_jdata, min_nbor_dist_ = Descriptor.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["list"] = new_list + return local_jdata_cpy, min_nbor_dist def serialize(self, suffix: str = "") -> dict: return { diff --git a/deepmd/tf/descriptor/loc_frame.py b/deepmd/tf/descriptor/loc_frame.py index 963e9bf607..e6247290c0 100644 --- a/deepmd/tf/descriptor/loc_frame.py +++ b/deepmd/tf/descriptor/loc_frame.py @@ -20,6 +20,9 @@ from deepmd.tf.utils.sess import ( run_sess, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -431,14 +434,28 @@ def init_variables( self.dstd = get_tensor_by_name_from_graph(graph, f"descrpt_attr{suffix}/t_std") @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index a14dc0a5ce..b6c23db575 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -2,6 +2,7 @@ import re from typing import ( List, + Optional, Set, Tuple, ) @@ -21,6 +22,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -150,19 +154,36 @@ def precision(self) -> tf.DType: return self.filter_precision @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # default behavior is to update sel which is a list local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + return local_jdata_cpy, min_nbor_dist def serialize_network( self, diff --git a/deepmd/tf/descriptor/se_a_mask.py b/deepmd/tf/descriptor/se_a_mask.py index e78dfba461..b79e806fca 100644 --- a/deepmd/tf/descriptor/se_a_mask.py +++ b/deepmd/tf/descriptor/se_a_mask.py @@ -24,6 +24,9 @@ from deepmd.tf.utils.network import ( embedding_net_rand_seed_shift, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .descriptor import ( Descriptor, @@ -420,14 +423,28 @@ def prod_force_virial( return force, virial, atom_virial @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ - return local_jdata + return local_jdata, None diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py index 43c38b0955..6d3cfeaa6e 100644 --- a/deepmd/tf/descriptor/se_atten.py +++ b/deepmd/tf/descriptor/se_atten.py @@ -78,6 +78,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.version import ( check_version_compatibility, ) @@ -1470,18 +1473,36 @@ def explicit_ntypes(self) -> bool: return True @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist def serialize_attention_layers( self, diff --git a/deepmd/tf/entrypoints/compress.py b/deepmd/tf/entrypoints/compress.py index 1e50d88b63..6f3ff13e3b 100644 --- a/deepmd/tf/entrypoints/compress.py +++ b/deepmd/tf/entrypoints/compress.py @@ -32,6 +32,9 @@ from deepmd.tf.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + get_data, +) from .freeze import ( freeze, @@ -115,9 +118,17 @@ def compress( log.info("stage 0: compute the min_nbor_dist") jdata = j_loader(training_script) jdata = update_deepmd_input(jdata) + + type_map = jdata["model"].get("type_map", None) + train_data = get_data( + jdata["training"]["training_data"], + 0, # not used + type_map, + None, + ) update_sel = UpdateSel() t_min_nbor_dist = update_sel.get_min_nbor_dist( - jdata, update_sel.get_rcut(jdata) + train_data, ) _check_compress_type(graph) diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index 3c4decbe8c..c654482656 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -17,6 +17,7 @@ j_loader, ) from deepmd.tf.env import ( + GLOBAL_ENER_FLOAT_PRECISION, reset_default_tf_session_config, tf, ) @@ -255,5 +256,20 @@ def update_sel(jdata): "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" ) jdata_cpy = jdata.copy() - jdata_cpy["model"] = Model.update_sel(jdata, jdata["model"]) + type_map = jdata["model"].get("type_map") + train_data = get_data( + jdata["training"]["training_data"], + 0, # not used + type_map, + None, # not used + ) + jdata_cpy["model"], min_nbor_dist = Model.update_sel( + train_data, type_map, jdata["model"] + ) + + tf.constant( + min_nbor_dist, + name="train_attr/min_nbor_dist", + dtype=GLOBAL_ENER_FLOAT_PRECISION, + ) return jdata_cpy diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index fa28b2cc58..3e296c00f2 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -8,6 +8,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -38,6 +39,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .model import ( Model, @@ -237,18 +241,32 @@ def get_type_map(self) -> list: return self.model.get_type_map() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # we don't know how to compress it, so no neighbor statistics here - return local_jdata + return local_jdata, None def serialize(self, suffix: str = "") -> dict: # try to recover the original model diff --git a/deepmd/tf/model/linear.py b/deepmd/tf/model/linear.py index 26bc382569..1bd1644e54 100644 --- a/deepmd/tf/model/linear.py +++ b/deepmd/tf/model/linear.py @@ -10,6 +10,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -27,6 +28,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from .model import ( Model, @@ -133,22 +137,42 @@ def get_type_map(self) -> list: return self.models[0].get_type_map() @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["models"] = [ - Model.update_sel(global_jdata, sub_jdata) - for sub_jdata in local_jdata["models"] - ] - return local_jdata_cpy + new_list = [] + min_nbor_dist = None + for sub_jdata in local_jdata["models"]: + new_sub_jdata, min_nbor_dist_ = Model.update_sel( + train_data, type_map, sub_jdata + ) + if min_nbor_dist_ is not None: + min_nbor_dist = min_nbor_dist_ + new_list.append(new_sub_jdata) + local_jdata_cpy["models"] = new_list + return local_jdata_cpy, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 194750a9d7..a1baf85dbc 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -11,6 +11,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -511,7 +512,12 @@ def get_feed_dict( @classmethod @abstractmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Notes @@ -520,8 +526,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class @@ -529,9 +537,11 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: ------- dict The updated local data + float + The minimum distance between two atoms """ cls = cls.get_class_by_type(local_jdata.get("type", "standard")) - return cls.update_sel(global_jdata, local_jdata) + return cls.update_sel(train_data, type_map, local_jdata) @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Model": @@ -744,21 +754,35 @@ def get_ntypes(self) -> int: return self.ntypes @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - local_jdata_cpy["descriptor"] = Descriptor.update_sel( - global_jdata, local_jdata["descriptor"] + local_jdata_cpy["descriptor"], min_nbor_dist = Descriptor.update_sel( + train_data, type_map, local_jdata["descriptor"] ) - return local_jdata_cpy + return local_jdata_cpy, min_nbor_dist @classmethod def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": diff --git a/deepmd/tf/model/pairtab.py b/deepmd/tf/model/pairtab.py index 4f71dcd76e..8c79dedea0 100644 --- a/deepmd/tf/model/pairtab.py +++ b/deepmd/tf/model/pairtab.py @@ -5,6 +5,7 @@ from typing import ( List, Optional, + Tuple, Union, ) @@ -35,6 +36,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @Model.register("pairtab") @@ -268,7 +272,12 @@ def enable_compression(self, suffix: str = "") -> None: # nothing needs to do @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Notes @@ -277,8 +286,10 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class @@ -286,9 +297,15 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: ------- dict The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, True) + min_nbor_dist, sel = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + ) + local_jdata_cpy["sel"] = sel[0] + return local_jdata_cpy, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/model/pairwise_dprc.py b/deepmd/tf/model/pairwise_dprc.py index 3d61dfd339..44e3943e12 100644 --- a/deepmd/tf/model/pairwise_dprc.py +++ b/deepmd/tf/model/pairwise_dprc.py @@ -3,6 +3,7 @@ Dict, List, Optional, + Tuple, Union, ) @@ -36,6 +37,9 @@ from deepmd.utils.data import ( DataRequirementItem, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) @Model.register("pairwise_dprc") @@ -410,20 +414,33 @@ def get_feed_dict( return feed_dict @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ # do not update sel; only find min distance - # rcut is not important here - UpdateSel().get_min_nbor_dist(global_jdata, 6.0) - return local_jdata + min_nbor_dist = UpdateSel().get_min_nbor_dist(train_data) + return local_jdata, min_nbor_dist @property def input_requirement(self) -> List[DataRequirementItem]: diff --git a/deepmd/tf/utils/update_sel.py b/deepmd/tf/utils/update_sel.py index db0420dde8..726aec4d41 100644 --- a/deepmd/tf/utils/update_sel.py +++ b/deepmd/tf/utils/update_sel.py @@ -3,12 +3,6 @@ Type, ) -from deepmd.env import ( - GLOBAL_ENER_FLOAT_PRECISION, -) -from deepmd.tf.env import ( - tf, -) from deepmd.tf.utils.neighbor_stat import ( NeighborStat, ) @@ -21,12 +15,3 @@ class UpdateSel(BaseUpdateSel): @property def neighbor_stat(self) -> Type[NeighborStat]: return NeighborStat - - def hook(self, min_nbor_dist, max_nbor_size): - # moved from traier.py as duplicated - tf.constant( - min_nbor_dist, - name="train_attr/min_nbor_dist", - dtype=GLOBAL_ENER_FLOAT_PRECISION, - ) - tf.constant(max_nbor_size, name="train_attr/max_nbor_size", dtype=tf.int32) diff --git a/deepmd/utils/update_sel.py b/deepmd/utils/update_sel.py index d1be8e8138..6feed525e5 100644 --- a/deepmd/utils/update_sel.py +++ b/deepmd/utils/update_sel.py @@ -1,14 +1,19 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging from abc import ( + ABC, abstractmethod, ) from typing import ( + List, + Optional, + Tuple, Type, + Union, ) from deepmd.utils.data_system import ( - get_data, + DeepmdDataSystem, ) from deepmd.utils.neighbor_stat import ( NeighborStat, @@ -17,32 +22,29 @@ log = logging.getLogger(__name__) -class BaseUpdateSel: +class BaseUpdateSel(ABC): """Update the sel field in the descriptor.""" def update_one_sel( self, - jdata, - descriptor, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + rcut: float, + sel: Union[int, List[int], str], mixed_type: bool = False, - rcut_key="rcut", - sel_key="sel", - ): - rcut = descriptor[rcut_key] - tmp_sel = self.get_sel( - jdata, + ) -> Tuple[float, List[int]]: + min_nbor_dist, tmp_sel = self.get_nbor_stat( + train_data, + type_map, rcut, mixed_type=mixed_type, ) - sel = descriptor[sel_key] if isinstance(sel, int): # convert to list and finnally convert back to int sel = [sel] - if self.parse_auto_sel(descriptor[sel_key]): - ratio = self.parse_auto_sel_ratio(descriptor[sel_key]) - descriptor[sel_key] = sel = [ - int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel - ] + if self.parse_auto_sel(sel): + ratio = self.parse_auto_sel_ratio(sel) + sel = [int(self.wrap_up_4(ii * ratio)) for ii in tmp_sel] else: # sel is set by user for ii, (tt, dd) in enumerate(zip(tmp_sel, sel)): @@ -54,9 +56,7 @@ def update_one_sel( "not less than %d, but you set it to %d. The accuracy" " of your model may get worse." % (ii, tt, dd) ) - if mixed_type: - descriptor[sel_key] = sum(sel) - return descriptor + return min_nbor_dist, sel def parse_auto_sel(self, sel): if not isinstance(sel, str): @@ -83,65 +83,36 @@ def parse_auto_sel_ratio(self, sel): def wrap_up_4(self, xx): return 4 * ((int(xx) + 3) // 4) - def get_sel(self, jdata, rcut, mixed_type: bool = False): - _, max_nbor_size = self.get_nbor_stat(jdata, rcut, mixed_type=mixed_type) - return max_nbor_size - - def get_rcut(self, jdata): - if jdata["model"].get("type") == "pairwise_dprc": - return max( - jdata["model"]["qm_model"]["descriptor"]["rcut"], - jdata["model"]["qmmm_model"]["descriptor"]["rcut"], - ) - descrpt_data = jdata["model"]["descriptor"] - rcut_list = [] - if descrpt_data["type"] == "hybrid": - for ii in descrpt_data["list"]: - rcut_list.append(ii["rcut"]) - else: - rcut_list.append(descrpt_data["rcut"]) - return max(rcut_list) - - def get_type_map(self, jdata): - return jdata["model"].get("type_map", None) - - def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): - # it seems that DeepmdDataSystem does not need rcut - # it's not clear why there is an argument... - # max_rcut = get_rcut(jdata) - max_rcut = rcut - type_map = self.get_type_map(jdata) - + def get_nbor_stat( + self, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + rcut: float, + mixed_type: bool = False, + ) -> Tuple[float, Union[int, List[int]]]: + """Get the neighbor statistics of the data. + + Parameters + ---------- + train_data : DeepmdDataSystem + The training data. + type_map : Optional[List[str]] + The type map. + rcut : float + The cutoff radius. + mixed_type : bool, optional + Whether to mix the types. + + Returns + ------- + min_nbor_dist : float + The minimum neighbor distance. + max_nbor_size : List[int] + The maximum neighbor size. + """ if type_map and len(type_map) == 0: type_map = None - multi_task_mode = "data_dict" in jdata["training"] - if not multi_task_mode: - train_data = get_data( - jdata["training"]["training_data"], max_rcut, type_map, None - ) - train_data.get_batch() - else: - assert ( - type_map is not None - ), "Data stat in multi-task mode must have available type_map! " - train_data = None - for systems in jdata["training"]["data_dict"]: - tmp_data = get_data( - jdata["training"]["data_dict"][systems]["training_data"], - max_rcut, - type_map, - None, - ) - tmp_data.get_batch() - assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " - if train_data is None: - train_data = tmp_data - else: - train_data.system_dirs += tmp_data.system_dirs - train_data.data_systems += tmp_data.data_systems - train_data.natoms += tmp_data.natoms - train_data.natoms_vec += tmp_data.natoms_vec - train_data.default_mesh += tmp_data.default_mesh + train_data.get_batch() data_ntypes = train_data.get_ntypes() if type_map is not None: map_ntypes = len(type_map) @@ -152,7 +123,6 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): neistat = self.neighbor_stat(ntypes, rcut, mixed_type=mixed_type) min_nbor_dist, max_nbor_size = neistat.get_stat(train_data) - self.hook(min_nbor_dist, max_nbor_size) return min_nbor_dist, max_nbor_size @@ -161,10 +131,14 @@ def get_nbor_stat(self, jdata, rcut, mixed_type: bool = False): def neighbor_stat(self) -> Type[NeighborStat]: pass - @abstractmethod - def hook(self, min_nbor_dist, max_nbor_size): - pass - - def get_min_nbor_dist(self, jdata, rcut): - min_nbor_dist, _ = self.get_nbor_stat(jdata, rcut) + def get_min_nbor_dist( + self, + train_data: DeepmdDataSystem, + ): + min_nbor_dist, _ = self.get_nbor_stat( + train_data, + None, # type_map doesn't affect min_nbor_dist + 1e-6, # we don't need the max_nbor_size + mixed_type=True, # mixed_types doesn't affect min_nbor_dist + ) return min_nbor_dist diff --git a/doc/development/create-a-model-pt.md b/doc/development/create-a-model-pt.md index 35d81b364a..c6d372b195 100644 --- a/doc/development/create-a-model-pt.md +++ b/doc/development/create-a-model-pt.md @@ -20,6 +20,7 @@ The framework-independent backend is implemented in pure NumPy, serving as a ref When creating a new descriptor, it is essential to inherit from both the {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor` class and the {py:class}`torch.nn.Module` class. Abstract methods, including {py:class}`deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor.forward`, must be implemented, while others remain optional. It is crucial to adhere to the original method arguments without any modifications. Once the implementation is complete, the next step involves registering the component with a designated key: ```py +from deepmd.utils.data_system import DeepmdDataSystem from deepmd.pt.model.descriptor.base_descriptor import ( BaseDescriptor, ) @@ -63,7 +64,12 @@ class SomeDescript(BaseDescriptor, torch.nn.Module): def deserialize(cls, data: dict) -> "SomeDescript": pass - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ): pass ``` From a397dd916d12c576446ad1d5313158510ff0ea57 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:01:46 -0400 Subject: [PATCH 02/12] fix tests Signed-off-by: Jinzhe Zeng --- deepmd/tf/entrypoints/train.py | 11 +-- source/tests/tf/test_train.py | 128 +++++++++++++++++++++------------ 2 files changed, 89 insertions(+), 50 deletions(-) diff --git a/deepmd/tf/entrypoints/train.py b/deepmd/tf/entrypoints/train.py index c654482656..d394773cf2 100755 --- a/deepmd/tf/entrypoints/train.py +++ b/deepmd/tf/entrypoints/train.py @@ -267,9 +267,10 @@ def update_sel(jdata): train_data, type_map, jdata["model"] ) - tf.constant( - min_nbor_dist, - name="train_attr/min_nbor_dist", - dtype=GLOBAL_ENER_FLOAT_PRECISION, - ) + if min_nbor_dist is not None: + tf.constant( + min_nbor_dist, + name="train_attr/min_nbor_dist", + dtype=GLOBAL_ENER_FLOAT_PRECISION, + ) return jdata_cpy diff --git a/source/tests/tf/test_train.py b/source/tests/tf/test_train.py index 3e22dc57bc..f5a57a948d 100644 --- a/source/tests/tf/test_train.py +++ b/source/tests/tf/test_train.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import random import unittest from unittest.mock import ( patch, @@ -15,6 +16,7 @@ class TestTrain(unittest.TestCase): def setUp(self) -> None: self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() return super().setUp() def test_train_parse_auto_sel(self): @@ -34,22 +36,25 @@ def test_train_parse_auto_sel_ratio(self): with self.assertRaises(RuntimeError): self.update_sel.parse_auto_sel_ratio([1, 2, 3]) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_one_sel(self, sel_mock): - sel_mock.return_value = [10, 20] - jdata = {} - descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto"} - descriptor = self.update_sel.update_one_sel(jdata, descriptor) + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") # self.assertEqual(descriptor['sel'], [11,22]) - self.assertEqual(descriptor["sel"], [12, 24]) - descriptor = {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"} - descriptor = self.update_sel.update_one_sel(jdata, descriptor) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") # self.assertEqual(descriptor['sel'], [15,30]) - self.assertEqual(descriptor["sel"], [16, 32]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_hybrid(self, sel_mock): - sel_mock.return_value = [10, 20] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -59,7 +64,8 @@ def test_update_sel_hybrid(self, sel_mock): {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, ], } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -70,24 +76,33 @@ def test_update_sel_hybrid(self, sel_mock): {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, ], } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel(self, sel_mock): - sel_mock.return_value = [10, 20] - jdata = {"model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}} + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + get_data_mock.return_value = None + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } expected_out = { - "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}} + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_auto(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -95,7 +110,8 @@ def test_update_sel_atten_auto(self, sel_mock): "sel": "auto", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -104,14 +120,17 @@ def test_update_sel_atten_auto(self, sel_mock): "sel": 28, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_int(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -119,7 +138,8 @@ def test_update_sel_atten_int(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -128,14 +148,17 @@ def test_update_sel_atten_int(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - @patch("deepmd.tf.utils.update_sel.UpdateSel.get_sel") - def test_update_sel_atten_list(self, sel_mock): - sel_mock.return_value = [25] + @patch("deepmd.tf.entrypoints.train.get_data") + @patch("deepmd.tf.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + get_data_mock.return_value = None jdata = { "model": { "descriptor": { @@ -143,7 +166,8 @@ def test_update_sel_atten_list(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -152,19 +176,23 @@ def test_update_sel_atten_list(self, sel_mock): "sel": 30, "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_loc_frame(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_loc_frame(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "descriptor": { "type": "loc_frame", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } expected_out = { "model": { @@ -172,22 +200,28 @@ def test_skip_loc_frame(self): "type": "loc_frame", "rcut": 6, } - } + }, + "training": {"training_data": {}}, } jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_frozen(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_frozen(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "type": "frozen", - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) - def test_skip_linear_frozen(self): + @patch("deepmd.tf.entrypoints.train.get_data") + def test_skip_linear_frozen(self, get_data_mock): + get_data_mock.return_value = None jdata = { "model": { "type": "linear_ener", @@ -197,15 +231,18 @@ def test_skip_linear_frozen(self): {"type": "frozen"}, {"type": "frozen"}, ], - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata) self.assertEqual(jdata, expected_out) + @patch("deepmd.tf.entrypoints.train.get_data") @patch("deepmd.tf.utils.update_sel.UpdateSel.get_min_nbor_dist") - def test_pairwise_dprc(self, sel_mock): - sel_mock.return_value = 0.5 + def test_pairwise_dprc(self, sel_mock, get_data_mock): + sel_mock.return_value = self.mock_min_nbor_dist + get_data_mock.return_value = None jdata = { "model": { "type": "pairwise_dprc", @@ -215,7 +252,8 @@ def test_pairwise_dprc(self, sel_mock): {"type": "frozen"}, {"type": "frozen"}, ], - } + }, + "training": {"training_data": {}}, } expected_out = jdata.copy() jdata = update_sel(jdata) From a8d97c95d89379f4305c0dbc9d73a100054471d7 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:10:44 -0400 Subject: [PATCH 03/12] fix pt Signed-off-by: Jinzhe Zeng --- deepmd/pt/entrypoints/main.py | 58 +++++++++++++++++++---------------- 1 file changed, 31 insertions(+), 27 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 62c0cefa1b..d1c78575c3 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -260,45 +260,49 @@ def train(FLAGS): config["training"]["training_data"], 0, type_map, None ) config["model"], min_nbor_dist = BaseModel.update_sel( - config, config["model"] + train_data, type_map, config["model"] ) else: - assert ( - type_map is not None - ), "Data stat in multi-task mode must have available type_map! " - train_data = None - for systems in config["training"]["data_dict"]: - tmp_data = get_data( - config["training"]["data_dict"][systems]["training_data"], - 0, - type_map, - None, - ) - tmp_data.get_batch() - assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " - if train_data is None: - train_data = tmp_data - else: - train_data.system_dirs += tmp_data.system_dirs - train_data.data_systems += tmp_data.data_systems - train_data.natoms += tmp_data.natoms - train_data.natoms_vec += tmp_data.natoms_vec - train_data.default_mesh += tmp_data.default_mesh - - training_jdata = deepcopy(config["training"]) - training_jdata.pop("data_dict", {}) - training_jdata.pop("model_prob", {}) # considering multi-task shares the descriptor, we need a minimal min_nbor_dist min_nbor_dist = 0 for model_item in config["model"]["model_dict"]: + training_jdata = deepcopy(config["training"]) + training_jdata.pop("data_dict", {}) + training_jdata.pop("model_prob", {}) + fake_global_jdata = { "model": deepcopy(config["model"]["model_dict"][model_item]), "training": deepcopy(config["training"]["data_dict"][model_item]), } fake_global_jdata["training"].update(training_jdata) + + assert ( + type_map is not None + ), "Data stat in multi-task mode must have available type_map! " + train_data = None + for systems in fake_global_jdata["training"]["data_dict"]: + tmp_data = get_data( + fake_global_jdata["training"]["data_dict"][systems][ + "training_data" + ], + 0, + type_map, + None, + ) + tmp_data.get_batch() + assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " + if train_data is None: + train_data = tmp_data + else: + train_data.system_dirs += tmp_data.system_dirs + train_data.data_systems += tmp_data.data_systems + train_data.natoms += tmp_data.natoms + train_data.natoms_vec += tmp_data.natoms_vec + train_data.default_mesh += tmp_data.default_mesh + config["model"]["model_dict"][model_item], min_nbor_dist_item = ( BaseModel.update_sel( - fake_global_jdata, config["model"]["model_dict"][model_item] + train_data, type_map, config["model"]["model_dict"][model_item] ) ) if min_nbor_dist_item is not None: From 33acd372267c84a31410756bb42c9a77a213d3eb Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:13:16 -0400 Subject: [PATCH 04/12] add tests for pt Signed-off-by: Jinzhe Zeng --- source/tests/pt/test_update_sel.py | 188 +++++++++++++++++++++++++++++ 1 file changed, 188 insertions(+) create mode 100644 source/tests/pt/test_update_sel.py diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py new file mode 100644 index 0000000000..8ae61e0260 --- /dev/null +++ b/source/tests/pt/test_update_sel.py @@ -0,0 +1,188 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import random +import unittest +from unittest.mock import ( + patch, +) + +from deepmd.pt.model.model.model import ( + BaseModel, +) +from deepmd.pt.utils.update_sel import ( + UpdateSel, +) + + +def update_sel(jdata): + type_map = jdata["model"].get("type_map") + train_data = None + jdata["model"], min_nbor_dist = BaseModel.update_sel( + train_data, type_map, jdata["model"] + ) + return jdata + + +class TestTrain(unittest.TestCase): + def setUp(self) -> None: + self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() + return super().setUp() + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") + # self.assertEqual(descriptor['sel'], [11,22]) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") + # self.assertEqual(descriptor['sel'], [15,30]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": "auto"}, + {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, + ], + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}, + {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, + ], + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } + expected_out = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": "auto", + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 28, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.pt.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_skip_frozen(self): + jdata = { + "model": { + "type": "frozen", + }, + "training": {"training_data": {}}, + } + expected_out = jdata.copy() + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_wrap_up_4(self): + self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) + self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4) From acb446387439d0bb434674febb58c6735115ee88 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:25:47 -0400 Subject: [PATCH 05/12] fix multi task Signed-off-by: Jinzhe Zeng --- deepmd/pt/entrypoints/main.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index d1c78575c3..44c5c7c47f 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -263,13 +263,11 @@ def train(FLAGS): train_data, type_map, config["model"] ) else: - # considering multi-task shares the descriptor, we need a minimal min_nbor_dist - min_nbor_dist = 0 + training_jdata = deepcopy(config["training"]) + training_jdata.pop("data_dict", {}) + training_jdata.pop("model_prob", {}) + min_nbor_dist = {} for model_item in config["model"]["model_dict"]: - training_jdata = deepcopy(config["training"]) - training_jdata.pop("data_dict", {}) - training_jdata.pop("model_prob", {}) - fake_global_jdata = { "model": deepcopy(config["model"]["model_dict"][model_item]), "training": deepcopy(config["training"]["data_dict"][model_item]), @@ -300,13 +298,11 @@ def train(FLAGS): train_data.natoms_vec += tmp_data.natoms_vec train_data.default_mesh += tmp_data.default_mesh - config["model"]["model_dict"][model_item], min_nbor_dist_item = ( + config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( train_data, type_map, config["model"]["model_dict"][model_item] ) ) - if min_nbor_dist_item is not None: - min_nbor_dist = min(min_nbor_dist, min_nbor_dist_item) with open(FLAGS.output, "w") as fp: json.dump(config, fp, indent=4) @@ -322,7 +318,11 @@ def train(FLAGS): shared_links=shared_links, ) # save min_nbor_dist - trainer.model.min_nbor_dist = min_nbor_dist + if not multi_task: + trainer.model.min_nbor_dist = min_nbor_dist + else: + for model_item in min_nbor_dist: + trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item] trainer.run() From 3db844c164cdaa0ade6728e2e0dcf3a76bbe6273 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:33:17 -0400 Subject: [PATCH 06/12] remove "fake_global_jdata" Signed-off-by: Jinzhe Zeng --- deepmd/pt/entrypoints/main.py | 42 +++++------------------------------ 1 file changed, 6 insertions(+), 36 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 44c5c7c47f..6654d19687 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -3,9 +3,6 @@ import json import logging import os -from copy import ( - deepcopy, -) from pathlib import ( Path, ) @@ -263,41 +260,14 @@ def train(FLAGS): train_data, type_map, config["model"] ) else: - training_jdata = deepcopy(config["training"]) - training_jdata.pop("data_dict", {}) - training_jdata.pop("model_prob", {}) min_nbor_dist = {} for model_item in config["model"]["model_dict"]: - fake_global_jdata = { - "model": deepcopy(config["model"]["model_dict"][model_item]), - "training": deepcopy(config["training"]["data_dict"][model_item]), - } - fake_global_jdata["training"].update(training_jdata) - - assert ( - type_map is not None - ), "Data stat in multi-task mode must have available type_map! " - train_data = None - for systems in fake_global_jdata["training"]["data_dict"]: - tmp_data = get_data( - fake_global_jdata["training"]["data_dict"][systems][ - "training_data" - ], - 0, - type_map, - None, - ) - tmp_data.get_batch() - assert tmp_data.get_type_map(), f"In multi-task mode, 'type_map.raw' must be defined in data systems {systems}! " - if train_data is None: - train_data = tmp_data - else: - train_data.system_dirs += tmp_data.system_dirs - train_data.data_systems += tmp_data.data_systems - train_data.natoms += tmp_data.natoms - train_data.natoms_vec += tmp_data.natoms_vec - train_data.default_mesh += tmp_data.default_mesh - + train_data = get_data( + config["training"]["data_dict"][model_item]["training_data"], + 0, + type_map, + None, + ) config["model"]["model_dict"][model_item], min_nbor_dist[model_item] = ( BaseModel.update_sel( train_data, type_map, config["model"]["model_dict"][model_item] From b67114f8971e1d24206312bd3d187780c48dc58e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:53:50 -0400 Subject: [PATCH 07/12] fix: mixed_types in several descriptors should be false Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/descriptor/se_e2_a.py | 2 +- deepmd/dpmodel/descriptor/se_r.py | 2 +- deepmd/pt/model/descriptor/se_a.py | 2 +- deepmd/pt/model/descriptor/se_r.py | 2 +- deepmd/tf/descriptor/se.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/deepmd/dpmodel/descriptor/se_e2_a.py index 6feffc114f..2b6c3a843e 100644 --- a/deepmd/dpmodel/descriptor/se_e2_a.py +++ b/deepmd/dpmodel/descriptor/se_e2_a.py @@ -451,6 +451,6 @@ def update_sel( """ local_jdata_cpy = local_jdata.copy() min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( - train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False ) return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/dpmodel/descriptor/se_r.py b/deepmd/dpmodel/descriptor/se_r.py index 82e9589465..e4a0a80657 100644 --- a/deepmd/dpmodel/descriptor/se_r.py +++ b/deepmd/dpmodel/descriptor/se_r.py @@ -375,6 +375,6 @@ def update_sel( """ local_jdata_cpy = local_jdata.copy() min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( - train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False ) return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 6293351d41..b6fe511f98 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -324,7 +324,7 @@ def update_sel( """ local_jdata_cpy = local_jdata.copy() min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( - train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False ) return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/se_r.py b/deepmd/pt/model/descriptor/se_r.py index 0f030733d1..b0a739f5e6 100644 --- a/deepmd/pt/model/descriptor/se_r.py +++ b/deepmd/pt/model/descriptor/se_r.py @@ -442,6 +442,6 @@ def update_sel( """ local_jdata_cpy = local_jdata.copy() min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( - train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False ) return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/tf/descriptor/se.py b/deepmd/tf/descriptor/se.py index b6c23db575..f5f54550f2 100644 --- a/deepmd/tf/descriptor/se.py +++ b/deepmd/tf/descriptor/se.py @@ -181,7 +181,7 @@ def update_sel( # default behavior is to update sel which is a list local_jdata_cpy = local_jdata.copy() min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( - train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], True + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False ) return local_jdata_cpy, min_nbor_dist From cedaca81a54d96a2ab2976e99c7a70df28e038d2 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 07:55:42 -0400 Subject: [PATCH 08/12] prevent unbonded min_nbor_dist when skip neighbor stat Signed-off-by: Jinzhe Zeng --- deepmd/pt/entrypoints/main.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/deepmd/pt/entrypoints/main.py b/deepmd/pt/entrypoints/main.py index 6654d19687..d3646f3b94 100644 --- a/deepmd/pt/entrypoints/main.py +++ b/deepmd/pt/entrypoints/main.py @@ -246,6 +246,7 @@ def train(FLAGS): config = normalize(config) # do neighbor stat + min_nbor_dist = None if not FLAGS.skip_neighbor_stat: log.info( "Calculate neighbor statistics... (add --skip-neighbor-stat to skip this step)" @@ -288,11 +289,12 @@ def train(FLAGS): shared_links=shared_links, ) # save min_nbor_dist - if not multi_task: - trainer.model.min_nbor_dist = min_nbor_dist - else: - for model_item in min_nbor_dist: - trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item] + if min_nbor_dist is not None: + if not multi_task: + trainer.model.min_nbor_dist = min_nbor_dist + else: + for model_item in min_nbor_dist: + trainer.model[model_item].min_nbor_dist = min_nbor_dist[model_item] trainer.run() From 0bcd410da6f1052a6a2dc85b3b30ed6d7efb81cc Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 08:09:40 -0400 Subject: [PATCH 09/12] do not save none Signed-off-by: Jinzhe Zeng --- deepmd/pt/utils/serialization.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py index 7ab65ab426..21a2a3fbda 100644 --- a/deepmd/pt/utils/serialization.py +++ b/deepmd/pt/utils/serialization.py @@ -50,10 +50,10 @@ def serialize_from_file(model_file: str) -> dict: "pt_version": torch.__version__, "model": model_dict, "model_def_script": model_def_script, - "@variables": { - "min_nbor_dist": model.get_min_nbor_dist(), - }, + "@variables": {}, } + if model.get_min_nbor_dist() is not None: + data["@variables"]["min_nbor_dist"] = model.get_min_nbor_dist() return data From 88c5219ef922d48c4c190559dd1c9ff08cd04372 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 08:11:22 -0400 Subject: [PATCH 10/12] fix type hints Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/model/spin_model.py | 2 +- deepmd/pt/model/model/frozen.py | 2 +- deepmd/pt/model/model/spin_model.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/dpmodel/model/spin_model.py b/deepmd/dpmodel/model/spin_model.py index 09bbabe358..c2cea35d27 100644 --- a/deepmd/dpmodel/model/spin_model.py +++ b/deepmd/dpmodel/model/spin_model.py @@ -227,7 +227,7 @@ def get_model_def_script(self) -> str: """Get the model definition script.""" return self.backbone_model.get_model_def_script() - def get_min_nbor_dist(self) -> float: + def get_min_nbor_dist(self) -> Optional[float]: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py index 4059450a2d..148ffaa703 100644 --- a/deepmd/pt/model/model/frozen.py +++ b/deepmd/pt/model/model/frozen.py @@ -135,7 +135,7 @@ def get_model_def_script(self) -> str: return self.model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> float: + def get_min_nbor_dist(self) -> Optional[float]: """Get the minimum neighbor distance.""" return self.model.get_min_nbor_dist() diff --git a/deepmd/pt/model/model/spin_model.py b/deepmd/pt/model/model/spin_model.py index 601d5fcd70..369a413f50 100644 --- a/deepmd/pt/model/model/spin_model.py +++ b/deepmd/pt/model/model/spin_model.py @@ -292,7 +292,7 @@ def get_model_def_script(self) -> str: return self.backbone_model.get_model_def_script() @torch.jit.export - def get_min_nbor_dist(self) -> float: + def get_min_nbor_dist(self) -> Optional[float]: """Get the minimum neighbor distance.""" return self.backbone_model.get_min_nbor_dist() From 27ea21892bc1ab670a7ff227cef22a36b44bdea3 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 28 May 2024 16:58:19 -0400 Subject: [PATCH 11/12] add tests for dpmodel Signed-off-by: Jinzhe Zeng --- .../tests/common/dpmodel/test_update_sel.py | 175 ++++++++++++++++++ source/tests/pt/test_update_sel.py | 4 +- 2 files changed, 176 insertions(+), 3 deletions(-) create mode 100644 source/tests/common/dpmodel/test_update_sel.py diff --git a/source/tests/common/dpmodel/test_update_sel.py b/source/tests/common/dpmodel/test_update_sel.py new file mode 100644 index 0000000000..a2656da7de --- /dev/null +++ b/source/tests/common/dpmodel/test_update_sel.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import random +import unittest +from unittest.mock import ( + patch, +) + +from deepmd.dpmodel.model.base_model import ( + BaseModel, +) +from deepmd.dpmodel.utils.update_sel import ( + UpdateSel, +) + + +def update_sel(jdata): + type_map = jdata["model"].get("type_map") + train_data = None + jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"]) + return jdata + + +class TestTrain(unittest.TestCase): + def setUp(self) -> None: + self.update_sel = UpdateSel() + self.mock_min_nbor_dist = random.random() + return super().setUp() + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_one_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto") + # self.assertEqual(descriptor['sel'], [11,22]) + self.assertEqual(sel, [12, 24]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + min_nbor_dist, sel = self.update_sel.update_one_sel(None, None, 6, "auto:1.5") + # self.assertEqual(descriptor['sel'], [15,30]) + self.assertEqual(sel, [16, 32]) + self.assertAlmostEqual(min_nbor_dist, self.mock_min_nbor_dist) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_hybrid(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": "auto"}, + {"type": "se_e2_a", "rcut": 6, "sel": "auto:1.5"}, + ], + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "hybrid", + "list": [ + {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}, + {"type": "se_e2_a", "rcut": 6, "sel": [16, 32]}, + ], + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [10, 20] + + jdata = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": "auto"}}, + "training": {"training_data": {}}, + } + expected_out = { + "model": {"descriptor": {"type": "se_e2_a", "rcut": 6, "sel": [12, 24]}}, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_auto(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": "auto", + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 28, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_int(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + @patch("deepmd.dpmodel.utils.update_sel.UpdateSel.get_nbor_stat") + def test_update_sel_atten_list(self, sel_mock): + sel_mock.return_value = self.mock_min_nbor_dist, [25] + + jdata = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + expected_out = { + "model": { + "descriptor": { + "type": "se_atten", + "sel": 30, + "rcut": 6, + } + }, + "training": {"training_data": {}}, + } + jdata = update_sel(jdata) + self.assertEqual(jdata, expected_out) + + def test_wrap_up_4(self): + self.assertEqual(self.update_sel.wrap_up_4(12), 3 * 4) + self.assertEqual(self.update_sel.wrap_up_4(13), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(14), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(15), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(16), 4 * 4) + self.assertEqual(self.update_sel.wrap_up_4(17), 5 * 4) diff --git a/source/tests/pt/test_update_sel.py b/source/tests/pt/test_update_sel.py index 8ae61e0260..e4eff0ebf3 100644 --- a/source/tests/pt/test_update_sel.py +++ b/source/tests/pt/test_update_sel.py @@ -16,9 +16,7 @@ def update_sel(jdata): type_map = jdata["model"].get("type_map") train_data = None - jdata["model"], min_nbor_dist = BaseModel.update_sel( - train_data, type_map, jdata["model"] - ) + jdata["model"], _ = BaseModel.update_sel(train_data, type_map, jdata["model"]) return jdata From ecb7b6671c33573480082aae5646e2f4bea0ba3a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 30 May 2024 19:48:52 -0400 Subject: [PATCH 12/12] se_e3 Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/descriptor/se_t.py | 28 ++++++++++++++++++++++++---- deepmd/pt/model/descriptor/se_t.py | 28 ++++++++++++++++++++++++---- 2 files changed, 48 insertions(+), 8 deletions(-) diff --git a/deepmd/dpmodel/descriptor/se_t.py b/deepmd/dpmodel/descriptor/se_t.py index d9b2741b3e..eac6a9640e 100644 --- a/deepmd/dpmodel/descriptor/se_t.py +++ b/deepmd/dpmodel/descriptor/se_t.py @@ -9,6 +9,9 @@ from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.path import ( DPPath, ) @@ -348,15 +351,32 @@ def deserialize(cls, data: dict) -> "DescrptSeT": return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist diff --git a/deepmd/pt/model/descriptor/se_t.py b/deepmd/pt/model/descriptor/se_t.py index db6244000d..2c8f52709f 100644 --- a/deepmd/pt/model/descriptor/se_t.py +++ b/deepmd/pt/model/descriptor/se_t.py @@ -30,6 +30,9 @@ from deepmd.pt.utils.update_sel import ( UpdateSel, ) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.env_mat_stat import ( StatItem, ) @@ -324,18 +327,35 @@ def t_cvt(xx): return obj @classmethod - def update_sel(cls, global_jdata: dict, local_jdata: dict): + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[List[str]], + local_jdata: dict, + ) -> Tuple[dict, Optional[float]]: """Update the selection and perform neighbor statistics. Parameters ---------- - global_jdata : dict - The global data, containing the training section + train_data : DeepmdDataSystem + data used to do neighbor statictics + type_map : list[str], optional + The name of each type of atoms local_jdata : dict The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms """ local_jdata_cpy = local_jdata.copy() - return UpdateSel().update_one_sel(global_jdata, local_jdata_cpy, False) + min_nbor_dist, local_jdata_cpy["sel"] = UpdateSel().update_one_sel( + train_data, type_map, local_jdata_cpy["rcut"], local_jdata_cpy["sel"], False + ) + return local_jdata_cpy, min_nbor_dist @DescriptorBlock.register("se_e3")