From 30f1051556fd01c282099086b65bf08749c7bb17 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 19 Feb 2024 19:39:34 -0500 Subject: [PATCH 1/7] checkpoint Signed-off-by: Jinzhe Zeng --- .../dpmodel/atomic_model/dp_atomic_model.py | 1 + deepmd/dpmodel/fitting/__init__.py | 4 + deepmd/dpmodel/model/make_model.py | 5 +- .../pt/model/atomic_model/dp_atomic_model.py | 1 + deepmd/tf/descriptor/descriptor.py | 2 +- deepmd/tf/fit/ener.py | 4 +- deepmd/tf/fit/fitting.py | 71 ++++++- deepmd/tf/model/model.py | 128 ++++++++++++- source/tests/consistent/common.py | 14 +- source/tests/consistent/model/__init__.py | 1 + source/tests/consistent/model/common.py | 65 +++++++ source/tests/consistent/model/test_ener.py | 177 ++++++++++++++++++ 12 files changed, 454 insertions(+), 19 deletions(-) create mode 100644 source/tests/consistent/model/__init__.py create mode 100644 source/tests/consistent/model/common.py create mode 100644 source/tests/consistent/model/test_ener.py diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index abf66cc3fa..220c072765 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -13,6 +13,7 @@ DescrptSeA, ) from deepmd.dpmodel.fitting import ( # noqa # TODO: should import all fittings! + EnergyFittingNet, InvarFitting, ) from deepmd.dpmodel.output_def import ( diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 2da752eaa7..2ee244a695 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -2,6 +2,9 @@ from .dipole_fitting import ( DipoleFitting, ) +from .ener_fitting import ( + EnergyFittingNet, +) from .invar_fitting import ( InvarFitting, ) @@ -13,4 +16,5 @@ "InvarFitting", "make_base_fitting", "DipoleFitting", + "EnergyFittingNet", ] diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index e44c2e9701..eb7e57747d 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -6,6 +6,9 @@ import numpy as np +from deepmd.dpmodel.common import ( + NativeOP, +) from deepmd.dpmodel.output_def import ( ModelOutputDef, ) @@ -45,7 +48,7 @@ def make_model(T_AtomicModel): """ - class CM(T_AtomicModel): + class CM(T_AtomicModel, NativeOP): def __init__( self, *args, diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index e6dc395500..154c6425c3 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -17,6 +17,7 @@ DescrptSeA, ) from deepmd.pt.model.task.ener import ( # noqa # TODO: should import all fittings! + EnergyFittingNet, InvarFitting, ) from deepmd.pt.utils.utils import ( diff --git a/deepmd/tf/descriptor/descriptor.py b/deepmd/tf/descriptor/descriptor.py index 1a73d3c273..768e233245 100644 --- a/deepmd/tf/descriptor/descriptor.py +++ b/deepmd/tf/descriptor/descriptor.py @@ -530,7 +530,7 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": The deserialized descriptor """ if cls is Descriptor: - return Descriptor.get_class_by_input(data).deserialize(data) + return Descriptor.get_class_by_input(data).deserialize(data, suffix=suffix) raise NotImplementedError("Not implemented in class %s" % cls.__name__) def serialize(self, suffix: str = "") -> dict: diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 074856ea6c..53eeda5f7f 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -929,7 +929,7 @@ def get_loss(self, loss: dict, lr) -> Loss: raise RuntimeError("unknown loss type") @classmethod - def deserialize(cls, data: dict, suffix: str): + def deserialize(cls, data: dict, suffix: str = ""): """Deserialize the model. Parameters @@ -956,7 +956,7 @@ def deserialize(cls, data: dict, suffix: str): fitting.aparam_inv_std = data["@variables"]["aparam_inv_std"] return fitting - def serialize(self, suffix: str) -> dict: + def serialize(self, suffix: str = "") -> dict: """Serialize the model. Returns diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index 2307fb957d..d121563fe4 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -50,16 +50,33 @@ class SomeFitting(Fitting): """ return Fitting.__plugins.register(key) + @classmethod + def get_class_by_input(cls, data: dict) -> "Fitting": + """Get the fitting class by the input data. + + Parameters + ---------- + data : dict + The input data + + Returns + ------- + Fitting + The fitting class + """ + try: + fitting_type = data["type"] + except KeyError: + raise KeyError("the type of fitting should be set by `type`") + if fitting_type in Fitting.__plugins.plugins: + cls = Fitting.__plugins.plugins[fitting_type] + else: + raise RuntimeError("Unknown descriptor type: " + fitting_type) + return cls + def __new__(cls, *args, **kwargs): if cls is Fitting: - try: - fitting_type = kwargs["type"] - except KeyError: - raise KeyError("the type of fitting should be set by `type`") - if fitting_type in Fitting.__plugins.plugins: - cls = Fitting.__plugins.plugins[fitting_type] - else: - raise RuntimeError("Unknown descriptor type: " + fitting_type) + cls = cls.get_class_by_input(kwargs) return super().__new__(cls) @property @@ -110,6 +127,44 @@ def get_loss(self, loss: dict, lr) -> Loss: the loss function """ + @classmethod + def deserialize(cls, data: dict, suffix: str = "") -> "Fitting": + """Deserialize the fitting. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Parameters + ---------- + data : dict + The serialized data + suffix : str, optional + Name suffix to identify this fitting + + Returns + ------- + Fitting + The deserialized fitting + """ + if cls is Fitting: + return Fitting.get_class_by_input(data).deserialize(data, suffix=suffix) + raise NotImplementedError("Not implemented in class %s" % cls.__name__) + + def serialize(self, suffix: str = "") -> dict: + """Serialize the fitting. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Returns + ------- + dict + The serialized data + suffix : str, optional + Name suffix to identify this fitting + """ + raise NotImplementedError("Not implemented in class %s" % self.__name__) + def serialize_network( self, ntypes: int, diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index ac970e0b53..4e93eda4d5 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy from abc import ( ABC, abstractmethod, @@ -13,6 +14,9 @@ Union, ) +from deepmd.dpmodel.fitting.dipole_fitting import ( + DipoleFitting, +) from deepmd.tf.descriptor.descriptor import ( Descriptor, ) @@ -20,9 +24,18 @@ GLOBAL_TF_FLOAT_PRECISION, tf, ) +from deepmd.tf.fit.dos import ( + DOSFitting, +) +from deepmd.tf.fit.ener import ( + EnerFitting, +) from deepmd.tf.fit.fitting import ( Fitting, ) +from deepmd.tf.fit.polar import ( + PolarFittingSeA, +) from deepmd.tf.loss.loss import ( Loss, ) @@ -565,6 +578,44 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: cls = cls.get_class_by_input(local_jdata) return cls.update_sel(global_jdata, local_jdata) + @classmethod + def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": + """Deserialize the model. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Parameters + ---------- + data : dict + The serialized data + suffix : str, optional + Name suffix to identify this descriptor + + Returns + ------- + Descriptor + The deserialized descriptor + """ + if cls is Descriptor: + return Descriptor.get_class_by_input(data).deserialize(data) + raise NotImplementedError("Not implemented in class %s" % cls.__name__) + + def serialize(self, suffix: str = "") -> dict: + """Serialize the model. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Returns + ------- + dict + The serialized data + suffix : str, optional + Name suffix to identify this descriptor + """ + raise NotImplementedError("Not implemented in class %s" % self.__name__) + class StandardModel(Model): """Standard model, which must contain a descriptor and a fitting. @@ -594,16 +645,21 @@ def __new__(cls, *args, **kwargs): ) if cls is StandardModel: - fitting_type = kwargs["fitting_net"]["type"] + if isinstance(kwargs["fitting_net"], dict): + fitting_type = Fitting.get_class_by_input(kwargs["fitting_net"]["type"]) + elif isinstance(kwargs["fitting_net"], Fitting): + fitting_type = fitting_type + else: + raise RuntimeError("get unknown fitting type when building model") # init model # infer model type by fitting_type - if fitting_type == "ener": + if isinstance(fitting_type, EnerFitting): cls = EnerModel - elif fitting_type == "dos": + elif isinstance(fitting_type, DOSFitting): cls = DOSModel - elif fitting_type == "dipole": + elif isinstance(fitting_type, DipoleFitting): cls = DipoleModel - elif fitting_type == "polar": + elif isinstance(fitting_type, PolarFittingSeA): cls = PolarModel else: raise RuntimeError("get unknown fitting type when building model") @@ -730,3 +786,65 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): global_jdata, local_jdata["descriptor"] ) return local_jdata_cpy + + @classmethod + def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": + """Deserialize the model. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Parameters + ---------- + data : dict + The serialized data + suffix : str, optional + Name suffix to identify this descriptor + + Returns + ------- + Descriptor + The deserialized descriptor + """ + data = copy.deepcopy(data) + + data["descriptor"]["type"] = { + "DescrptSeA": "se_e2_a", + }[data.pop("descriptor_name")] + data["fitting"]["type"] = { + "EnergyFittingNet": "ener", + }[data.pop("fitting_name")] + descriptor = Descriptor.deserialize(data.pop("descriptor"), suffix=suffix) + fitting = Fitting.deserialize(data.pop("fitting"), suffix=suffix) + return cls( + descriptor=descriptor, + fitting_net=fitting, + **data, + ) + + def serialize(self, suffix: str = "") -> dict: + """Serialize the model. + + There is no suffix in a native DP model, but it is important + for the TF backend. + + Returns + ------- + dict + The serialized data + suffix : str, optional + Name suffix to identify this descriptor + """ + if self.typeebd is not None: + raise NotImplementedError("type embedding is not supported") + if self.spin is not None: + raise NotImplementedError("spin is not supported") + return { + "type_map": self.type_map, + "descriptor": self.descrpt.serialize(suffix=suffix), + "fitting": self.fitting.serialize(suffix=suffix), + "descriptor_name": self.descrpt.__class__.__name__, + "fitting_name": {"EnerFitting": "EnergyFittingNet"}[ + self.fitting.__class__.__name__ + ], + } diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index 45c2114b24..90cd337a16 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -16,6 +16,7 @@ List, Optional, Tuple, + Union, ) from uuid import ( uuid4, @@ -68,7 +69,7 @@ class CommonTest(ABC): """Native DP model class.""" pt_class: ClassVar[Optional[type]] """PyTorch model class.""" - args: ClassVar[Optional[List[Argument]]] + args: ClassVar[Optional[Union[Argument, List[Argument]]]] """Arguments that maps to the `data`.""" skip_dp: ClassVar[bool] = False """Whether to skip the native DP model.""" @@ -91,9 +92,18 @@ def init_backend_cls(self, cls) -> Any: if self.args is None: data = self.data else: - base = Argument("arg", dict, sub_fields=self.args) + if isinstance(self.args, list): + base = Argument("arg", dict, sub_fields=self.args) + elif isinstance(self.args, Argument): + base = self.args + else: + raise ValueError("Invalid type for args") data = base.normalize_value(self.data, trim_pattern="_*") base.check_value(data, strict=True) + return self.pass_data_to_cls(cls, data) + + def pass_data_to_cls(self, cls, data) -> Any: + """Pass data to the class.""" return cls(**data, **self.addtional_data) @abstractmethod diff --git a/source/tests/consistent/model/__init__.py b/source/tests/consistent/model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/consistent/model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py new file mode 100644 index 0000000000..dcd8a340e4 --- /dev/null +++ b/source/tests/consistent/model/common.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +from deepmd.common import ( + make_default_mesh, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, +) + +if INSTALLED_PT: + import torch + + from deepmd.pt.utils.env import DEVICE as PT_DEVICE +if INSTALLED_TF: + from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, + tf, + ) + + +class ModelTest: + """Useful utilities for model tests.""" + + def build_tf_model(self, obj, natoms, coords, atype, box, suffix): + t_coord = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, [None, None, None], name="i_coord" + ) + t_type = tf.placeholder(tf.int32, [None, None], name="i_type") + t_natoms = tf.placeholder(tf.int32, natoms.shape, name="i_natoms") + t_box = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None, 9], name="i_box") + t_mesh = tf.placeholder(tf.int32, [None], name="i_mesh") + ret = obj.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + {}, + suffix=suffix, + ) + return [ret["energy"], ret["atom_ener"]], { + t_coord: coords, + t_type: atype, + t_natoms: natoms, + t_box: box, + t_mesh: make_default_mesh(True, False), + } + + def eval_dp_model(self, dp_obj: Any, natoms, coords, atype, box) -> Any: + return dp_obj(coords, atype, box=box) + + def eval_pt_model(self, pt_obj: Any, natoms, coords, atype, box) -> Any: + return { + kk: vv.detach().cpu().numpy() if torch.is_tensor(vv) else vv + for kk, vv in pt_obj( + torch.from_numpy(coords).to(PT_DEVICE), + torch.from_numpy(atype).to(PT_DEVICE), + box=torch.from_numpy(box).to(PT_DEVICE), + ).items() + } diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py new file mode 100644 index 0000000000..1069768e31 --- /dev/null +++ b/source/tests/consistent/model/test_ener.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, + Tuple, +) + +import numpy as np + +from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeA as DescrptSeADP +from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnerFittingDP +from deepmd.dpmodel.model.dp_model import DPModel as EnergyModelDP +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, + CommonTest, +) +from .common import ( + ModelTest, +) + +if INSTALLED_PT: + from deepmd.pt.model.model import ( + get_model, + ) + from deepmd.pt.model.model.ener_model import EnergyModel as EnergyModelPT + +else: + EnergyModelPT = None +if INSTALLED_TF: + from deepmd.tf.model.ener import EnerModel as EnergyModelTF +else: + EnergyModelTF = None +from deepmd.utils.argcheck import ( + model_args, +) + + +# @parameterized( +# (True, False), # resnet_dt +# (True, False), # type_one_side +# ([], [[0, 1]]), # excluded_types +# ) +class TestSeA(CommonTest, ModelTest, unittest.TestCase): + @property + def data(self) -> dict: + return { + "type_map": ["O", "H"], + "descriptor": { + "type": "se_e2_a", + "sel": [20, 20], + "rcut_smth": 0.50, + "rcut": 6.00, + "neuron": [ + 3, + 6, + ], + "resnet_dt": False, + "axis_neuron": 2, + "precision": "float64", + "type_one_side": True, + "seed": 1, + }, + "fitting_net": { + "neuron": [ + 5, + 5, + ], + "resnet_dt": True, + "precision": "float64", + "seed": 1, + }, + } + + tf_class = EnergyModelTF + dp_class = EnergyModelDP + pt_class = EnergyModelPT + args = model_args() + + def pass_data_to_cls(self, cls, data) -> Any: + """Pass data to the class.""" + data = data.copy() + if cls is EnergyModelDP: + # should not do it here... + data["descriptor"].pop("type") + data["fitting_net"].pop("type") + descriptor = DescrptSeADP( + **data["descriptor"], + ) + fitting = EnerFittingDP( + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + **data["fitting_net"], + ) + return cls( + descriptor=descriptor, + fitting=fitting, + type_map=data["type_map"], + ) + elif cls is EnergyModelPT: + return get_model(data) + return cls(**data, **self.addtional_data) + + def setUp(self): + CommonTest.setUp(self) + + self.ntypes = 2 + self.coords = np.array( + [ + 12.83, + 2.56, + 2.18, + 12.09, + 2.87, + 2.74, + 00.25, + 3.32, + 1.68, + 3.36, + 3.00, + 1.81, + 3.51, + 2.51, + 2.60, + 4.27, + 3.22, + 1.56, + ], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, -1, 3) + self.atype = np.array([0, 1, 1, 0, 1, 1], dtype=np.int32).reshape(1, -1) + self.box = np.array( + [13.0, 0.0, 0.0, 0.0, 13.0, 0.0, 0.0, 0.0, 13.0], + dtype=GLOBAL_NP_FLOAT_PRECISION, + ).reshape(1, 9) + self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + + def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: + return self.build_tf_model( + obj, + self.natoms, + self.coords, + self.atype, + self.box, + suffix, + ) + + def eval_dp(self, dp_obj: Any) -> Any: + return self.eval_dp_model( + dp_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def eval_pt(self, pt_obj: Any) -> Any: + return self.eval_pt_model( + pt_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + + def extract_ret(self, ret: Any, backend) -> Tuple[np.ndarray, ...]: + if backend is self.RefBackend.DP: + return (ret["energy_redu"].ravel(), ret["energy"]) + elif backend is self.RefBackend.PT: + return (ret["energy"].ravel(), ret["atom_energy"]) + elif backend is self.RefBackend.TF: + return (ret[0], ret[1]) + raise ValueError(f"Unknown backend: {backend}") From e8e4f99a338f88fbe85cab7d0179fcfba95a643b Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 19 Feb 2024 19:47:27 -0500 Subject: [PATCH 2/7] atype must be sorted Signed-off-by: Jinzhe Zeng --- source/tests/consistent/model/test_ener.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index 1069768e31..cbf7ca89c7 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -139,6 +139,11 @@ def setUp(self): ).reshape(1, 9) self.natoms = np.array([6, 6, 2, 4], dtype=np.int32) + # TF requires the atype to be sort + idx_map = np.argsort(self.atype.ravel()) + self.atype = self.atype[:, idx_map] + self.coords = self.coords[:, idx_map] + def build_tf(self, obj: Any, suffix: str) -> Tuple[list, dict]: return self.build_tf_model( obj, @@ -168,10 +173,11 @@ def eval_pt(self, pt_obj: Any) -> Any: ) def extract_ret(self, ret: Any, backend) -> Tuple[np.ndarray, ...]: + # shape not matched. ravel... if backend is self.RefBackend.DP: - return (ret["energy_redu"].ravel(), ret["energy"]) + return (ret["energy_redu"].ravel(), ret["energy"].ravel()) elif backend is self.RefBackend.PT: - return (ret["energy"].ravel(), ret["atom_energy"]) + return (ret["energy"].ravel(), ret["atom_energy"].ravel()) elif backend is self.RefBackend.TF: - return (ret[0], ret[1]) + return (ret[0].ravel(), ret[1].ravel()) raise ValueError(f"Unknown backend: {backend}") From 5f74ac1dba5eedb4922532925f781a43d7456d19 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 19 Feb 2024 20:30:42 -0500 Subject: [PATCH 3/7] fix errors Signed-off-by: Jinzhe Zeng --- deepmd/tf/fit/fitting.py | 3 ++- deepmd/tf/model/model.py | 12 ++++++------ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/deepmd/tf/fit/fitting.py b/deepmd/tf/fit/fitting.py index d121563fe4..458765f7c1 100644 --- a/deepmd/tf/fit/fitting.py +++ b/deepmd/tf/fit/fitting.py @@ -6,6 +6,7 @@ from typing import ( Callable, List, + Type, ) from deepmd.dpmodel.utils.network import ( @@ -51,7 +52,7 @@ class SomeFitting(Fitting): return Fitting.__plugins.register(key) @classmethod - def get_class_by_input(cls, data: dict) -> "Fitting": + def get_class_by_input(cls, data: dict) -> Type["Fitting"]: """Get the fitting class by the input data. Parameters diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 4e93eda4d5..60759a2953 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -646,20 +646,20 @@ def __new__(cls, *args, **kwargs): if cls is StandardModel: if isinstance(kwargs["fitting_net"], dict): - fitting_type = Fitting.get_class_by_input(kwargs["fitting_net"]["type"]) + fitting_type = Fitting.get_class_by_input(kwargs["fitting_net"]) elif isinstance(kwargs["fitting_net"], Fitting): - fitting_type = fitting_type + fitting_type = type(kwargs["fitting_net"]) else: raise RuntimeError("get unknown fitting type when building model") # init model # infer model type by fitting_type - if isinstance(fitting_type, EnerFitting): + if issubclass(fitting_type, EnerFitting): cls = EnerModel - elif isinstance(fitting_type, DOSFitting): + elif issubclass(fitting_type, DOSFitting): cls = DOSModel - elif isinstance(fitting_type, DipoleFitting): + elif issubclass(fitting_type, DipoleFitting): cls = DipoleModel - elif isinstance(fitting_type, PolarFittingSeA): + elif issubclass(fitting_type, PolarFittingSeA): cls = PolarModel else: raise RuntimeError("get unknown fitting type when building model") From 9da94efe61c2b8dc7634e08981fa1400ab9d9678 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 19 Feb 2024 20:55:07 -0500 Subject: [PATCH 4/7] fix dipole fitting Signed-off-by: Jinzhe Zeng --- deepmd/tf/model/model.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index 60759a2953..c2def1677f 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -14,9 +14,6 @@ Union, ) -from deepmd.dpmodel.fitting.dipole_fitting import ( - DipoleFitting, -) from deepmd.tf.descriptor.descriptor import ( Descriptor, ) @@ -24,6 +21,9 @@ GLOBAL_TF_FLOAT_PRECISION, tf, ) +from deepmd.tf.fit.dipole import ( + DipoleFittingSeA, +) from deepmd.tf.fit.dos import ( DOSFitting, ) @@ -651,13 +651,14 @@ def __new__(cls, *args, **kwargs): fitting_type = type(kwargs["fitting_net"]) else: raise RuntimeError("get unknown fitting type when building model") + print(fitting_type) # init model # infer model type by fitting_type if issubclass(fitting_type, EnerFitting): cls = EnerModel elif issubclass(fitting_type, DOSFitting): cls = DOSModel - elif issubclass(fitting_type, DipoleFitting): + elif issubclass(fitting_type, DipoleFittingSeA): cls = DipoleModel elif issubclass(fitting_type, PolarFittingSeA): cls = PolarModel From 79371fb03bea5b32b2008a91cae97bb9e9171a40 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 23:07:01 -0500 Subject: [PATCH 5/7] use deepmd.pt.utils.utils Signed-off-by: Jinzhe Zeng --- source/tests/consistent/model/common.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index dcd8a340e4..294edec1d6 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -13,9 +13,8 @@ ) if INSTALLED_PT: - import torch - - from deepmd.pt.utils.env import DEVICE as PT_DEVICE + from deepmd.pt.utils.utils import to_numpy_array as torch_to_numpy + from deepmd.pt.utils.utils import to_torch_tensor as numpy_to_torch if INSTALLED_TF: from deepmd.tf.env import ( GLOBAL_TF_FLOAT_PRECISION, @@ -56,10 +55,10 @@ def eval_dp_model(self, dp_obj: Any, natoms, coords, atype, box) -> Any: def eval_pt_model(self, pt_obj: Any, natoms, coords, atype, box) -> Any: return { - kk: vv.detach().cpu().numpy() if torch.is_tensor(vv) else vv + kk: torch_to_numpy(vv) for kk, vv in pt_obj( - torch.from_numpy(coords).to(PT_DEVICE), - torch.from_numpy(atype).to(PT_DEVICE), - box=torch.from_numpy(box).to(PT_DEVICE), + numpy_to_torch(coords), + numpy_to_torch(atype), + box=numpy_to_torch(box), ).items() } From ab080f77807af3dd129844dc1e289affdf2ae8ab Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 23:14:55 -0500 Subject: [PATCH 6/7] rewrite the typing of to_numpy_array and to_torch_tensor Signed-off-by: Jinzhe Zeng --- deepmd/pt/utils/utils.py | 29 +++++++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/utils/utils.py b/deepmd/pt/utils/utils.py index d6621f7b4c..852c42cd0c 100644 --- a/deepmd/pt/utils/utils.py +++ b/deepmd/pt/utils/utils.py @@ -2,6 +2,7 @@ from typing import ( Callable, Optional, + overload, ) import numpy as np @@ -51,9 +52,19 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: raise RuntimeError(f"activation function {self.activation} not supported") +@overload +def to_numpy_array(xx: torch.Tensor) -> np.ndarray: + ... + + +@overload +def to_numpy_array(xx: None) -> None: + ... + + def to_numpy_array( - xx: torch.Tensor, -) -> np.ndarray: + xx, +): if xx is None: return None assert xx is not None @@ -67,9 +78,19 @@ def to_numpy_array( return xx.detach().cpu().numpy().astype(prec) +@overload +def to_torch_tensor(xx: np.ndarray) -> torch.Tensor: + ... + + +@overload +def to_torch_tensor(xx: None) -> None: + ... + + def to_torch_tensor( - xx: np.ndarray, -) -> torch.Tensor: + xx, +): if xx is None: return None assert xx is not None From 83fd5b27b4a31ab83476f61d42924976d29cba02 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 20 Feb 2024 23:31:31 -0500 Subject: [PATCH 7/7] move get_model to dpmodel Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/model/dp_model.py | 5 ++- deepmd/dpmodel/model/model.py | 41 ++++++++++++++++++++++ source/tests/consistent/model/test_ener.py | 33 +++-------------- 3 files changed, 50 insertions(+), 29 deletions(-) create mode 100644 deepmd/dpmodel/model/model.py diff --git a/deepmd/dpmodel/model/dp_model.py b/deepmd/dpmodel/model/dp_model.py index 6196bcfe87..c2c40b40ba 100644 --- a/deepmd/dpmodel/model/dp_model.py +++ b/deepmd/dpmodel/model/dp_model.py @@ -7,4 +7,7 @@ make_model, ) -DPModel = make_model(DPAtomicModel) + +# use "class" to resolve "Variable not allowed in type expression" +class DPModel(make_model(DPAtomicModel)): + pass diff --git a/deepmd/dpmodel/model/model.py b/deepmd/dpmodel/model/model.py new file mode 100644 index 0000000000..4a6e269f25 --- /dev/null +++ b/deepmd/dpmodel/model/model.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_e2_a import ( + DescrptSeA, +) +from deepmd.dpmodel.fitting.ener_fitting import ( + EnergyFittingNet, +) +from deepmd.dpmodel.model.dp_model import ( + DPModel, +) + + +def get_model(data: dict) -> DPModel: + """Get a DPModel from a dictionary. + + Parameters + ---------- + data : dict + The data to construct the model. + """ + descriptor_type = data["descriptor"].pop("type") + fitting_type = data["fitting_net"].pop("type") + if descriptor_type == "se_e2_a": + descriptor = DescrptSeA( + **data["descriptor"], + ) + else: + raise ValueError(f"Unknown descriptor type {descriptor_type}") + if fitting_type == "ener": + fitting = EnergyFittingNet( + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + **data["fitting_net"], + ) + else: + raise ValueError(f"Unknown fitting type {fitting_type}") + return DPModel( + descriptor=descriptor, + fitting=fitting, + type_map=data["type_map"], + ) diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index cbf7ca89c7..b3aa778ca0 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -7,9 +7,8 @@ import numpy as np -from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeA as DescrptSeADP -from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnerFittingDP from deepmd.dpmodel.model.dp_model import DPModel as EnergyModelDP +from deepmd.dpmodel.model.model import get_model as get_model_dp from deepmd.env import ( GLOBAL_NP_FLOAT_PRECISION, ) @@ -24,9 +23,7 @@ ) if INSTALLED_PT: - from deepmd.pt.model.model import ( - get_model, - ) + from deepmd.pt.model.model import get_model as get_model_pt from deepmd.pt.model.model.ener_model import EnergyModel as EnergyModelPT else: @@ -40,12 +37,7 @@ ) -# @parameterized( -# (True, False), # resnet_dt -# (True, False), # type_one_side -# ([], [[0, 1]]), # excluded_types -# ) -class TestSeA(CommonTest, ModelTest, unittest.TestCase): +class TestEner(CommonTest, ModelTest, unittest.TestCase): @property def data(self) -> dict: return { @@ -85,24 +77,9 @@ def pass_data_to_cls(self, cls, data) -> Any: """Pass data to the class.""" data = data.copy() if cls is EnergyModelDP: - # should not do it here... - data["descriptor"].pop("type") - data["fitting_net"].pop("type") - descriptor = DescrptSeADP( - **data["descriptor"], - ) - fitting = EnerFittingDP( - ntypes=descriptor.get_ntypes(), - dim_descrpt=descriptor.get_dim_out(), - **data["fitting_net"], - ) - return cls( - descriptor=descriptor, - fitting=fitting, - type_map=data["type_map"], - ) + return get_model_dp(data) elif cls is EnergyModelPT: - return get_model(data) + return get_model_pt(data) return cls(**data, **self.addtional_data) def setUp(self):