diff --git a/deepmd/backend/backend.py b/deepmd/backend/backend.py index 179b2e556a..f1ef4cb52a 100644 --- a/deepmd/backend/backend.py +++ b/deepmd/backend/backend.py @@ -141,6 +141,8 @@ class Feature(Flag): """Support Deep Eval backend.""" NEIGHBOR_STAT = auto() """Support neighbor statistics.""" + IO = auto() + """Support IO hook.""" name: ClassVar[str] = "Unknown" """The formal name of the backend. @@ -199,3 +201,27 @@ def neighbor_stat(self) -> Type["NeighborStat"]: The neighbor statistics of the backend. """ pass + + @property + @abstractmethod + def serialize_hook(self) -> Callable[[str], dict]: + """The serialize hook to convert the model file to a dictionary. + + Returns + ------- + Callable[[str], dict] + The serialize hook of the backend. + """ + pass + + @property + @abstractmethod + def deserialize_hook(self) -> Callable[[str, dict], None]: + """The deserialize hook to convert the dictionary to a model file. + + Returns + ------- + Callable[[str, dict], None] + The deserialize hook of the backend. + """ + pass diff --git a/deepmd/backend/dpmodel.py b/deepmd/backend/dpmodel.py index 8745ca6d5a..e5c09349cf 100644 --- a/deepmd/backend/dpmodel.py +++ b/deepmd/backend/dpmodel.py @@ -33,7 +33,9 @@ class DPModelBackend(Backend): name = "DPModel" """The formal name of the backend.""" - features: ClassVar[Backend.Feature] = Backend.Feature.NEIGHBOR_STAT + features: ClassVar[Backend.Feature] = ( + Backend.Feature.NEIGHBOR_STAT | Backend.Feature.IO + ) """The features of the backend.""" suffixes: ClassVar[List[str]] = [".dp"] """The suffixes of the backend.""" @@ -84,3 +86,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]: ) return NeighborStat + + @property + def serialize_hook(self) -> Callable[[str], dict]: + """The serialize hook to convert the model file to a dictionary. + + Returns + ------- + Callable[[str], dict] + The serialize hook of the backend. + """ + from deepmd.dpmodel.utils.network import ( + load_dp_model, + ) + + return load_dp_model + + @property + def deserialize_hook(self) -> Callable[[str, dict], None]: + """The deserialize hook to convert the dictionary to a model file. + + Returns + ------- + Callable[[str, dict], None] + The deserialize hook of the backend. + """ + from deepmd.dpmodel.utils.network import ( + save_dp_model, + ) + + return save_dp_model diff --git a/deepmd/backend/pytorch.py b/deepmd/backend/pytorch.py index 4c0b0699f9..676694172b 100644 --- a/deepmd/backend/pytorch.py +++ b/deepmd/backend/pytorch.py @@ -38,6 +38,7 @@ class TensorFlowBackend(Backend): Backend.Feature.ENTRY_POINT | Backend.Feature.DEEP_EVAL | Backend.Feature.NEIGHBOR_STAT + | Backend.Feature.IO ) """The features of the backend.""" suffixes: ClassVar[List[str]] = [".pth", ".pt"] @@ -93,3 +94,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]: ) return NeighborStat + + @property + def serialize_hook(self) -> Callable[[str], dict]: + """The serialize hook to convert the model file to a dictionary. + + Returns + ------- + Callable[[str], dict] + The serialize hook of the backend. + """ + from deepmd.pt.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file + + @property + def deserialize_hook(self) -> Callable[[str, dict], None]: + """The deserialize hook to convert the dictionary to a model file. + + Returns + ------- + Callable[[str, dict], None] + The deserialize hook of the backend. + """ + from deepmd.pt.utils.serialization import ( + deserialize_to_file, + ) + + return deserialize_to_file diff --git a/deepmd/backend/tensorflow.py b/deepmd/backend/tensorflow.py index 80569afa97..15b03ee7c8 100644 --- a/deepmd/backend/tensorflow.py +++ b/deepmd/backend/tensorflow.py @@ -38,6 +38,7 @@ class TensorFlowBackend(Backend): Backend.Feature.ENTRY_POINT | Backend.Feature.DEEP_EVAL | Backend.Feature.NEIGHBOR_STAT + | Backend.Feature.IO ) """The features of the backend.""" suffixes: ClassVar[List[str]] = [".pb"] @@ -102,3 +103,33 @@ def neighbor_stat(self) -> Type["NeighborStat"]: ) return NeighborStat + + @property + def serialize_hook(self) -> Callable[[str], dict]: + """The serialize hook to convert the model file to a dictionary. + + Returns + ------- + Callable[[str], dict] + The serialize hook of the backend. + """ + from deepmd.tf.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file + + @property + def deserialize_hook(self) -> Callable[[str, dict], None]: + """The deserialize hook to convert the dictionary to a model file. + + Returns + ------- + Callable[[str, dict], None] + The deserialize hook of the backend. + """ + from deepmd.tf.utils.serialization import ( + deserialize_to_file, + ) + + return deserialize_to_file diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 8c826c8771..c0a62c9a3e 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -6,6 +6,9 @@ import copy import itertools import json +from datetime import ( + datetime, +) from typing import ( ClassVar, Dict, @@ -54,6 +57,8 @@ def traverse_model_dict(model_obj, callback: callable, is_variable: bool = False elif isinstance(model_obj, list): for ii, vv in enumerate(model_obj): model_obj[ii] = traverse_model_dict(vv, callback, is_variable=is_variable) + elif model_obj is None: + return model_obj elif is_variable: model_obj = callback(model_obj) return model_obj @@ -79,7 +84,8 @@ def __call__(self): return self.count -def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] = None): +# TODO: should be moved to otherwhere... +def save_dp_model(filename: str, model_dict: dict) -> None: """Save a DP model to a file in the native format. Parameters @@ -88,15 +94,9 @@ def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] = The filename to save to. model_dict : dict The model dict to save. - extra_info : dict, optional - Extra meta information to save. """ model_dict = model_dict.copy() variable_counter = Counter() - if extra_info is not None: - extra_info = extra_info.copy() - else: - extra_info = {} with h5py.File(filename, "w") as f: model_dict = traverse_model_dict( model_dict, @@ -105,10 +105,11 @@ def save_dp_model(filename: str, model_dict: dict, extra_info: Optional[dict] = ).name, ) save_dict = { - "model": model_dict, "software": "deepmd-kit", "version": __version__, - **extra_info, + # use UTC+0 time + "time": str(datetime.utcnow()), + **model_dict, } f.attrs["json"] = json.dumps(save_dict, separators=(",", ":")) diff --git a/deepmd/entrypoints/convert_backend.py b/deepmd/entrypoints/convert_backend.py new file mode 100644 index 0000000000..39967d565c --- /dev/null +++ b/deepmd/entrypoints/convert_backend.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.backend.backend import ( + Backend, +) + + +def convert_backend( + *, # Enforce keyword-only arguments + INPUT: str, + OUTPUT: str, + **kwargs, +) -> None: + """Convert a model file from one backend to another. + + Parameters + ---------- + INPUT : str + The input model file. + INPUT : str + The output model file. + """ + inp_backend: Backend = Backend.detect_backend_by_model(INPUT)() + out_backend: Backend = Backend.detect_backend_by_model(OUTPUT)() + inp_hook = inp_backend.serialize_hook + out_hook = out_backend.deserialize_hook + data = inp_hook(INPUT) + out_hook(OUTPUT, data) diff --git a/deepmd/entrypoints/main.py b/deepmd/entrypoints/main.py index 9a03ac5e45..9f05b9a530 100644 --- a/deepmd/entrypoints/main.py +++ b/deepmd/entrypoints/main.py @@ -12,6 +12,9 @@ from deepmd.backend.suffix import ( format_model_suffix, ) +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) from deepmd.entrypoints.doc import ( doc_train_input, ) @@ -76,5 +79,7 @@ def main(args: argparse.Namespace): neighbor_stat(**dict_args) elif args.command == "gui": start_dpgui(**dict_args) + elif args.command == "convert-backend": + convert_backend(**dict_args) else: raise ValueError(f"Unknown command: {args.command}") diff --git a/deepmd/main.py b/deepmd/main.py index 98f5ab0c6b..4d2d62ed14 100644 --- a/deepmd/main.py +++ b/deepmd/main.py @@ -721,6 +721,23 @@ def main_parser() -> argparse.ArgumentParser: "to the network on both IPv4 and IPv6 (where available)." ), ) + + # convert_backend + parser_convert_backend = subparsers.add_parser( + "convert-backend", + parents=[parser_log], + help="Convert model to another backend.", + formatter_class=RawTextArgumentDefaultsHelpFormatter, + epilog=textwrap.dedent( + """\ + examples: + dp convert-backend model.pb model.pth + dp convert-backend model.pb model.dp + """ + ), + ) + parser_convert_backend.add_argument("INPUT", help="The input model file.") + parser_convert_backend.add_argument("OUTPUT", help="The output model file.") return parser @@ -767,6 +784,7 @@ def main(): "model-devi", "neighbor-stat", "gui", + "convert-backend", ): # common entrypoints from deepmd.entrypoints.main import main as deepmd_main diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index 73b2d76a6d..1e5f976baf 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -14,3 +14,7 @@ class BaseAtomicModel(BaseAtomicModel_): # export public methods that are not abstract get_nsel = torch.jit.export(BaseAtomicModel_.get_nsel) get_nnei = torch.jit.export(BaseAtomicModel_.get_nnei) + + @torch.jit.export + def get_model_def_script(self) -> str: + return self.model_def_script diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 926aea6d70..98bf6c0fde 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -50,6 +50,7 @@ class DPAtomicModel(torch.nn.Module, BaseAtomicModel): def __init__(self, descriptor, fitting, type_map: Optional[List[str]]): super().__init__() + self.model_def_script = "" ntypes = len(type_map) self.type_map = type_map self.ntypes = ntypes diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 16b06b2211..70afbcb0bc 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -288,6 +288,7 @@ def __init__( ): models = [dp_model, zbl_model] super().__init__(models, **kwargs) + self.model_def_script = "" self.dp_model = dp_model self.zbl_model = zbl_model diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index d8b830d1eb..cf5a70eb88 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -52,6 +52,7 @@ def __init__( self, tab_file: str, rcut: float, sel: Union[int, List[int]], **kwargs ): super().__init__() + self.model_def_script = "" self.tab_file = tab_file self.rcut = rcut self.tab = self._set_pairtab(tab_file, rcut) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 0c15cae46c..4a7469a804 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -87,6 +87,8 @@ def __init__( self.ffn = ffn self.ffn_embed_dim = ffn_embed_dim self.activation = activation + # TODO: To be fixed: precision should be given from inputs + self.prec = torch.float64 self.scaling_factor = scaling_factor self.head_num = head_num self.normalize = normalize diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index ad16741ee4..974c42ee41 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -10,6 +10,7 @@ """ import copy +import json from deepmd.pt.model.atomic_model import ( DPAtomicModel, @@ -98,7 +99,9 @@ def get_model(model_params): fitting_net["return_energy"] = True fitting = Fitting(**fitting_net) - return EnergyModel(descriptor, fitting, type_map=model_params["type_map"]) + model = EnergyModel(descriptor, fitting, type_map=model_params["type_map"]) + model.model_def_script = json.dumps(model_params) + return model __all__ = [ diff --git a/deepmd/pt/utils/serialization.py b/deepmd/pt/utils/serialization.py new file mode 100644 index 0000000000..91d1a3c76f --- /dev/null +++ b/deepmd/pt/utils/serialization.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json + +import torch + +from deepmd.pt.model.model import ( + get_model, +) +from deepmd.pt.model.model.ener_model import ( + EnergyModel, +) +from deepmd.pt.train.wrapper import ( + ModelWrapper, +) + + +def serialize_from_file(model_file: str) -> dict: + """Serialize the model file to a dictionary. + + Parameters + ---------- + model_file : str + The model file to be serialized. + + Returns + ------- + dict + The serialized model data. + """ + if model_file.endswith(".pth"): + saved_model = torch.jit.load(model_file, map_location="cpu") + model_def_script = json.loads(saved_model.model_def_script) + model = get_model(model_def_script) + model.load_state_dict(saved_model.state_dict()) + elif model_file.endswith(".pt"): + state_dict = torch.load(model_file, map_location="cpu") + if "model" in state_dict: + state_dict = state_dict["model"] + model_def_script = state_dict["_extra_state"]["model_params"] + model = get_model(model_def_script) + modelwrapper = ModelWrapper(model) + modelwrapper.load_state_dict(state_dict) + model = modelwrapper.model["Default"] + else: + raise ValueError("PyTorch backend only supports converting .pth or .pt file") + + model_dict = model.serialize() + data = { + "backend": "PyTorch", + "pt_version": torch.__version__, + "model": model_dict, + "model_def_script": model_def_script, + # TODO + "@variables": {}, + } + return data + + +def deserialize_to_file(model_file: str, data: dict) -> None: + """Deserialize the dictionary to a model file. + + Parameters + ---------- + model_file : str + The model file to be saved. + data : dict + The dictionary to be deserialized. + """ + if not model_file.endswith(".pth"): + raise ValueError("PyTorch backend only supports converting .pth file") + # TODO: read class type from data; see #3319 + model = EnergyModel.deserialize(data["model"]) + # JIT will happy in this way... + model.model_def_script = json.dumps(data["model_def_script"]) + model = torch.jit.script(model) + torch.jit.save(model, model_file) diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 6be63f907c..19bec5cec0 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -194,7 +194,7 @@ def __init__( ), "length of trainable should be that of n_neuron + 1" self.atom_ener = [] self.atom_ener_v = atom_ener - for at, ae in enumerate(atom_ener): + for at, ae in enumerate(atom_ener if atom_ener is not None else []): if ae is not None: self.atom_ener.append( tf.constant(ae, GLOBAL_TF_FLOAT_PRECISION, name="atom_%d_ener" % at) diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index aab0ab6578..73339a450f 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -581,7 +581,7 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict) -> dict: return cls.update_sel(global_jdata, local_jdata) @classmethod - def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": + def deserialize(cls, data: dict, suffix: str = "") -> "Model": """Deserialize the model. There is no suffix in a native DP model, but it is important @@ -592,17 +592,17 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Descriptor": data : dict The serialized data suffix : str, optional - Name suffix to identify this descriptor + Name suffix to identify this model Returns ------- - Descriptor - The deserialized descriptor + Model + The deserialized Model """ - if cls is Descriptor: - return Descriptor.get_class_by_type( - data.get("type", "standard") - ).deserialize(data) + if cls is Model: + return Model.get_class_by_type(data.get("type", "standard")).deserialize( + data + ) raise NotImplementedError("Not implemented in class %s" % cls.__name__) def serialize(self, suffix: str = "") -> dict: diff --git a/deepmd/tf/utils/serialization.py b/deepmd/tf/utils/serialization.py new file mode 100644 index 0000000000..7cf596f5bd --- /dev/null +++ b/deepmd/tf/utils/serialization.py @@ -0,0 +1,132 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import tempfile + +from deepmd.tf.entrypoints import ( + freeze, +) +from deepmd.tf.env import ( + GLOBAL_TF_FLOAT_PRECISION, + tf, +) +from deepmd.tf.model.model import ( + Model, +) +from deepmd.tf.utils.errors import ( + GraphWithoutTensorError, +) +from deepmd.tf.utils.graph import ( + get_tensor_by_name_from_graph, + load_graph_def, +) +from deepmd.tf.utils.sess import ( + run_sess, +) + + +def serialize_from_file(model_file: str) -> dict: + """Serialize the model file to a dictionary. + + Parameters + ---------- + model_file : str + The model file to be serialized. + + Returns + ------- + dict + The serialized model data. + """ + graph, graph_def = load_graph_def(model_file) + t_jdata = get_tensor_by_name_from_graph(graph, "train_attr/training_script") + jdata = json.loads(t_jdata) + model = Model(**jdata["model"]) + # important! must be called before serialize + model.init_variables(graph=graph, graph_def=graph_def) + model_dict = model.serialize() + data = { + "backend": "TensorFlow", + "tf_version": tf.__version__, + "model": model_dict, + "model_def_script": jdata["model"], + } + # neighbor stat information + try: + t_min_nbor_dist = get_tensor_by_name_from_graph( + graph, "train_attr/min_nbor_dist" + ) + except GraphWithoutTensorError as e: + pass + else: + data.setdefault("@variables", {}) + data["@variables"]["min_nbor_dist"] = t_min_nbor_dist + return data + + +def deserialize_to_file(model_file: str, data: dict) -> None: + """Deserialize the dictionary to a model file. + + Parameters + ---------- + model_file : str + The model file to be saved. + data : dict + The dictionary to be deserialized. + """ + model = Model.deserialize(data["model"]) + with tf.Graph().as_default() as graph, tf.Session(graph=graph) as sess: + place_holders = {} + for ii in ["coord", "box"]: + place_holders[ii] = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, [None], name="t_" + ii + ) + place_holders["type"] = tf.placeholder(tf.int32, [None], name="t_type") + place_holders["natoms_vec"] = tf.placeholder( + tf.int32, [model.get_ntypes() + 2], name="t_natoms" + ) + place_holders["default_mesh"] = tf.placeholder(tf.int32, [None], name="t_mesh") + inputs = {} + # fparam, aparam + if model.get_numb_fparam() > 0: + inputs["fparam"] = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, + [None, model.get_numb_fparam()], + name="t_fparam", + ) + if model.get_numb_aparam() > 0: + inputs["aparam"] = tf.placeholder( + GLOBAL_TF_FLOAT_PRECISION, + [None, model.get_numb_aparam()], + name="t_aparam", + ) + model.build( + place_holders["coord"], + place_holders["type"], + place_holders["natoms_vec"], + place_holders["box"], + place_holders["default_mesh"], + inputs, + reuse=False, + ) + init = tf.global_variables_initializer() + tf.constant( + json.dumps({"model": data["model_def_script"]}, separators=(",", ":")), + name="train_attr/training_script", + dtype=tf.string, + ) + if "min_nbor_dist" in data.get("@variables", {}): + tf.constant( + data["@variables"]["min_nbor_dist"], + name="train_attr/min_nbor_dist", + dtype=GLOBAL_TF_FLOAT_PRECISION, + ) + run_sess(sess, init) + saver = tf.train.Saver() + with tempfile.TemporaryDirectory() as nt: + saver.save( + sess, + os.path.join(nt, "model.ckpt"), + global_step=0, + ) + freeze(checkpoint_folder=nt, output=model_file, node_names=None) diff --git a/source/tests/common/dpmodel/test_network.py b/source/tests/common/dpmodel/test_network.py index bfed8da45a..047eee501c 100644 --- a/source/tests/common/dpmodel/test_network.py +++ b/source/tests/common/dpmodel/test_network.py @@ -285,7 +285,7 @@ def setUp(self) -> None: self.filename = "test_dp_dpmodel.dp" def test_save_load_model(self): - save_dp_model(self.filename, deepcopy(self.model_dict)) + save_dp_model(self.filename, {"model": deepcopy(self.model_dict)}) model = load_dp_model(self.filename) np.testing.assert_equal(model["model"], self.model_dict) assert "software" in model diff --git a/source/tests/consistent/io/__init__.py b/source/tests/consistent/io/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/source/tests/consistent/io/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py new file mode 100644 index 0000000000..b8fae40cda --- /dev/null +++ b/source/tests/consistent/io/test_io.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import unittest +from pathlib import ( + Path, +) + +import numpy as np + +from deepmd.backend.backend import ( + Backend, +) +from deepmd.dpmodel.model.model import ( + get_model, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) + +infer_path = Path(__file__).parent.parent.parent / "infer" + + +class IOTest: + data: dict + + def get_data_from_model(self, model_file: str) -> dict: + """Get data from a model file. + + Parameters + ---------- + model_file : str + The model file. + + Returns + ------- + dict + The data from the model file. + """ + inp_backend: Backend = Backend.detect_backend_by_model(model_file)() + inp_hook = inp_backend.serialize_hook + return inp_hook(model_file) + + def save_data_to_model(self, model_file: str, data: dict) -> None: + """Save data to a model file. + + Parameters + ---------- + model_file : str + The model file. + data : dict + The data to save. + """ + out_backend: Backend = Backend.detect_backend_by_model(model_file)() + out_hook = out_backend.deserialize_hook + out_hook(model_file, data) + + def test_data_equal(self): + prefix = "test_consistent_io_" + self.__class__.__name__.lower() + for backend_name in ("tensorflow", "pytorch", "dpmodel"): + with self.subTest(backend_name=backend_name): + backend = Backend.get_backend(backend_name)() + if not backend.is_available: + continue + reference_data = copy.deepcopy(self.data) + self.save_data_to_model(prefix + backend.suffixes[0], reference_data) + data = self.get_data_from_model(prefix + backend.suffixes[0]) + data = copy.deepcopy(data) + reference_data = copy.deepcopy(self.data) + # some keys are not expected to be not the same + for kk in [ + "backend", + "tf_version", + "pt_version", + "@variables", + # dpmodel only + "software", + "version", + "time", + ]: + data.pop(kk, None) + reference_data.pop(kk, None) + np.testing.assert_equal(data, reference_data) + + def test_deep_eval(self): + 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) + prefix = "test_consistent_io_" + self.__class__.__name__.lower() + rets = [] + for backend_name in ("tensorflow", "pytorch"): + backend = Backend.get_backend(backend_name)() + if not backend.is_available: + continue + reference_data = copy.deepcopy(self.data) + self.save_data_to_model(prefix + backend.suffixes[0], reference_data) + deep_eval = DeepPot(prefix + backend.suffixes[0]) + ret = deep_eval.eval( + self.coords, + self.box, + self.atype, + ) + rets.append(ret) + for ret in rets[1:]: + for vv1, vv2 in zip(rets[0], ret): + np.testing.assert_allclose(vv1, vv2, rtol=1e-12, atol=1e-12) + + +class TestDeepPot(unittest.TestCase, IOTest): + def setUp(self): + model_def_script = { + "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": { + "type": "ener", + "neuron": [ + 5, + 5, + ], + "resnet_dt": True, + "precision": "float64", + "atom_ener": [], + "seed": 1, + }, + } + model = get_model(copy.deepcopy(model_def_script)) + self.data = { + "model": model.serialize(), + "backend": "test", + "model_def_script": model_def_script, + }