From a4e0e9ee83eb4985334e2284a644742ce54f61bf Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 28 Aug 2023 19:34:05 -0400 Subject: [PATCH 1/3] support linear model Signed-off-by: Jinzhe Zeng --- README.md | 1 + deepmd/entrypoints/train.py | 2 + deepmd/infer/__init__.py | 49 +++++- deepmd/infer/deep_dipole.py | 5 + deepmd/infer/deep_dos.py | 4 + deepmd/infer/deep_eval.py | 11 +- deepmd/infer/deep_polar.py | 4 + deepmd/infer/deep_pot.py | 4 + deepmd/infer/deep_tensor.py | 9 +- deepmd/infer/deep_wfc.py | 5 + deepmd/model/frozen.py | 195 +++++++++++++++++++++++ deepmd/model/linear.py | 250 ++++++++++++++++++++++++++++++ deepmd/model/model.py | 14 +- deepmd/train/trainer.py | 8 + deepmd/utils/argcheck.py | 42 +++++ doc/model/index.md | 1 + doc/model/index.rst | 1 + doc/model/linear.md | 24 +++ examples/water/linear/input.json | 56 +++++++ source/tests/test_examples.py | 1 + source/tests/test_linear_model.py | 125 +++++++++++++++ 21 files changed, 799 insertions(+), 12 deletions(-) create mode 100644 deepmd/model/frozen.py create mode 100644 deepmd/model/linear.py create mode 100644 doc/model/linear.md create mode 100644 examples/water/linear/input.json create mode 100644 source/tests/test_linear_model.py diff --git a/README.md b/README.md index 76f5c9d3bb..9b9d0ff27d 100644 --- a/README.md +++ b/README.md @@ -112,6 +112,7 @@ A full [document](doc/train/train-input-auto.rst) on options in the training inp - [Train a Deep Potential model using `type embedding` approach](doc/model/train-se-e2-a-tebd.md) - [Deep potential long-range](doc/model/dplr.md) - [Deep Potential - Range Correction (DPRc)](doc/model/dprc.md) + - [Linear model](doc/model/linear.md) - [Training](doc/train/index.md) - [Training a model](doc/train/training.md) - [Advanced options](doc/train/training-advanced.md) diff --git a/deepmd/entrypoints/train.py b/deepmd/entrypoints/train.py index fa3a82bbdf..1a0d4b9c6d 100755 --- a/deepmd/entrypoints/train.py +++ b/deepmd/entrypoints/train.py @@ -520,6 +520,8 @@ def update_sel(jdata): rcut = get_rcut(jdata) get_min_nbor_dist(jdata, rcut) return jdata + elif jdata["model"].get("type") in ("linear_ener", "frozen"): + return jdata descrpt_data = jdata["model"]["descriptor"] if descrpt_data["type"] == "hybrid": for ii in range(len(descrpt_data["list"])): diff --git a/deepmd/infer/__init__.py b/deepmd/infer/__init__.py index 5055ca9cd9..14d75d0c44 100644 --- a/deepmd/infer/__init__.py +++ b/deepmd/infer/__init__.py @@ -5,6 +5,7 @@ Path, ) from typing import ( + Optional, Union, ) @@ -56,6 +57,7 @@ def DeepPotential( model_file: Union[str, Path], load_prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ) -> Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepDOS, DeepWFC]: """Factory function that will inialize appropriate potential read from `model_file`. @@ -67,6 +69,8 @@ def DeepPotential( The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Returns ------- @@ -81,23 +85,54 @@ def DeepPotential( mf = Path(model_file) model_type = DeepEval( - mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, ).model_type if model_type == "ener": - dp = DeepPot(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + dp = DeepPot( + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) elif model_type == "dos": - dp = DeepDOS(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + dp = DeepDOS( + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) elif model_type == "dipole": - dp = DeepDipole(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + dp = DeepDipole( + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) elif model_type == "polar": - dp = DeepPolar(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + dp = DeepPolar( + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) elif model_type == "global_polar": dp = DeepGlobalPolar( - mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, ) elif model_type == "wfc": - dp = DeepWFC(mf, load_prefix=load_prefix, default_tf_graph=default_tf_graph) + dp = DeepWFC( + mf, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, + ) else: raise RuntimeError(f"unknown model type {model_type}") diff --git a/deepmd/infer/deep_dipole.py b/deepmd/infer/deep_dipole.py index 0464a6c33f..6020118135 100644 --- a/deepmd/infer/deep_dipole.py +++ b/deepmd/infer/deep_dipole.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, + Optional, ) from deepmd.infer.deep_tensor import ( @@ -24,6 +25,8 @@ class DeepDipole(DeepTensor): The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Warnings -------- @@ -37,6 +40,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ) -> None: # use this in favor of dict update to move attribute from class to # instance namespace @@ -53,6 +57,7 @@ def __init__( model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph, + input_map=input_map, ) def get_dim_fparam(self) -> int: diff --git a/deepmd/infer/deep_dos.py b/deepmd/infer/deep_dos.py index 52ef056b1a..5f181bd336 100644 --- a/deepmd/infer/deep_dos.py +++ b/deepmd/infer/deep_dos.py @@ -46,6 +46,8 @@ class DeepDOS(DeepEval): auto_batch_size : bool or int or AutomaticBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Warnings -------- @@ -60,6 +62,7 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = True, + input_map: Optional[dict] = None, ) -> None: # add these tensors on top of what is defined by DeepTensor Class # use this in favor of dict update to move attribute from class to @@ -91,6 +94,7 @@ def __init__( load_prefix=load_prefix, default_tf_graph=default_tf_graph, auto_batch_size=auto_batch_size, + input_map=input_map, ) # load optional tensors diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index 899c8c9acf..3f5dede1ad 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -43,6 +43,8 @@ class DeepEval: auto_batch_size : bool or int or AutomaticBatchSize, default: False If True, automatic batch size will be used. If int, it will be used as the initial batch size. + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph """ load_prefix: str # set by subclass @@ -53,9 +55,13 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = False, + input_map: Optional[dict] = None, ): self.graph = self._load_graph( - model_file, prefix=load_prefix, default_tf_graph=default_tf_graph + model_file, + prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, ) self.load_prefix = load_prefix @@ -168,6 +174,7 @@ def _load_graph( frozen_graph_filename: "Path", prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ): # We load the protobuf file from the disk and parse it to retrieve the # unserialized graph_def @@ -178,7 +185,7 @@ def _load_graph( if default_tf_graph: tf.import_graph_def( graph_def, - input_map=None, + input_map=input_map, return_elements=None, name=prefix, producer_op_list=None, diff --git a/deepmd/infer/deep_polar.py b/deepmd/infer/deep_polar.py index 6ecbf8aae6..118f8c98a7 100644 --- a/deepmd/infer/deep_polar.py +++ b/deepmd/infer/deep_polar.py @@ -28,6 +28,8 @@ class DeepPolar(DeepTensor): The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Warnings -------- @@ -41,6 +43,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ) -> None: # use this in favor of dict update to move attribute from class to # instance namespace @@ -57,6 +60,7 @@ def __init__( model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph, + input_map=input_map, ) def get_dim_fparam(self) -> int: diff --git a/deepmd/infer/deep_pot.py b/deepmd/infer/deep_pot.py index b3e9be1e67..031c5de1bc 100644 --- a/deepmd/infer/deep_pot.py +++ b/deepmd/infer/deep_pot.py @@ -49,6 +49,8 @@ class DeepPot(DeepEval): auto_batch_size : bool or int or AutomaticBatchSize, default: True If True, automatic batch size will be used. If int, it will be used as the initial batch size. + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Examples -------- @@ -75,6 +77,7 @@ def __init__( load_prefix: str = "load", default_tf_graph: bool = False, auto_batch_size: Union[bool, int, AutoBatchSize] = True, + input_map: Optional[dict] = None, ) -> None: # add these tensors on top of what is defined by DeepTensor Class # use this in favor of dict update to move attribute from class to @@ -108,6 +111,7 @@ def __init__( load_prefix=load_prefix, default_tf_graph=default_tf_graph, auto_batch_size=auto_batch_size, + input_map=input_map, ) # load optional tensors diff --git a/deepmd/infer/deep_tensor.py b/deepmd/infer/deep_tensor.py index 30b6fcfea5..367a8ab5e7 100644 --- a/deepmd/infer/deep_tensor.py +++ b/deepmd/infer/deep_tensor.py @@ -35,6 +35,8 @@ class DeepTensor(DeepEval): The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph """ tensors = { @@ -58,10 +60,15 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ) -> None: """Constructor.""" DeepEval.__init__( - self, model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph + self, + model_file, + load_prefix=load_prefix, + default_tf_graph=default_tf_graph, + input_map=input_map, ) # check model type model_type = self.tensors["t_tensor"][2:-2] diff --git a/deepmd/infer/deep_wfc.py b/deepmd/infer/deep_wfc.py index 00b10bc543..ed682f642b 100644 --- a/deepmd/infer/deep_wfc.py +++ b/deepmd/infer/deep_wfc.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( TYPE_CHECKING, + Optional, ) from deepmd.infer.deep_tensor import ( @@ -24,6 +25,8 @@ class DeepWFC(DeepTensor): The prefix in the load computational graph default_tf_graph : bool If uses the default tf graph, otherwise build a new tf graph for evaluation + input_map : dict, optional + The input map for tf.import_graph_def. Only work with default tf graph Warnings -------- @@ -37,6 +40,7 @@ def __init__( model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False, + input_map: Optional[dict] = None, ) -> None: # use this in favor of dict update to move attribute from class to # instance namespace @@ -52,6 +56,7 @@ def __init__( model_file, load_prefix=load_prefix, default_tf_graph=default_tf_graph, + input_map=input_map, ) def get_dim_fparam(self) -> int: diff --git a/deepmd/model/frozen.py b/deepmd/model/frozen.py new file mode 100644 index 0000000000..689b07edd7 --- /dev/null +++ b/deepmd/model/frozen.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from enum import ( + Enum, +) +from typing import ( + Optional, + Union, +) + +from deepmd.env import ( + GLOBAL_TF_FLOAT_PRECISION, + MODEL_VERSION, + tf, +) +from deepmd.fit.fitting import ( + Fitting, +) +from deepmd.infer import ( + DeepPotential, +) +from deepmd.loss.loss import ( + Loss, +) + +from .model import ( + Model, +) + + +class FrozenModel(Model): + """Load model from a frozen model, which cannot be trained. + + Parameters + ---------- + model_file : str + The path to the frozen model + """ + + def __init__(self, model_file: str, **kwargs): + super().__init__(**kwargs) + self.model_file = model_file + self.model = DeepPotential(model_file) + self.model_type = self.model.model_type + + def build( + self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box: tf.Tensor, + mesh: tf.Tensor, + input_dict: dict, + frz_model: Optional[str] = None, + ckpt_meta: Optional[str] = None, + suffix: str = "", + reuse: Optional[Union[bool, Enum]] = None, + ) -> dict: + """Build the model. + + Parameters + ---------- + coord_ : tf.Tensor + The coordinates of atoms + atype_ : tf.Tensor + The atom types of atoms + natoms : tf.Tensor + The number of atoms + box : tf.Tensor + The box vectors + mesh : tf.Tensor + The mesh vectors + input_dict : dict + The input dict + frz_model : str, optional + The path to the frozen model + ckpt_meta : str, optional + The path to the checkpoint and meta file + suffix : str, optional + The suffix of the scope + reuse : bool or tf.AUTO_REUSE, optional + Whether to reuse the variables + + Returns + ------- + dict + The output dict + """ + # reset the model to import to the correct graph + extra_feed_dict = {} + if input_dict is not None: + if "fparam" in input_dict: + extra_feed_dict["fparam"] = input_dict["fparam"] + if "aparam" in input_dict: + extra_feed_dict["aparam"] = input_dict["aparam"] + input_map = self.get_feed_dict( + coord_, atype_, natoms, box, mesh, **extra_feed_dict + ) + self.model = DeepPotential( + self.model_file, + default_tf_graph=True, + load_prefix="load" + suffix, + input_map=input_map, + ) + + with tf.variable_scope("model_attr" + suffix, reuse=reuse): + t_tmap = tf.constant( + " ".join(self.get_type_map()), name="tmap", dtype=tf.string + ) + t_mt = tf.constant(self.model_type, name="model_type", dtype=tf.string) + t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string) + with tf.variable_scope("descrpt_attr" + suffix, reuse=reuse): + t_ntypes = tf.constant(self.get_ntypes(), name="ntypes", dtype=tf.int32) + t_rcut = tf.constant( + self.get_rcut(), name="rcut", dtype=GLOBAL_TF_FLOAT_PRECISION + ) + if self.model_type == "ener": + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): + t_dfparam = tf.constant( + self.model.get_dim_fparam(), name="dfparam", dtype=tf.int32 + ) + t_daparam = tf.constant( + self.model.get_dim_aparam(), name="daparam", dtype=tf.int32 + ) + return { + "energy": tf.identity(self.model.t_energy, name="o_energy" + suffix), + "force": tf.identity(self.model.t_force, name="o_force" + suffix), + "virial": tf.identity(self.model.t_virial, name="o_virial" + suffix), + "atom_ener": tf.identity( + self.model.t_ae, name="o_atom_energy" + suffix + ), + "atom_virial": tf.identity( + self.model.t_av, name="o_atom_virial" + suffix + ), + "coord": coord_, + "atype": atype_, + } + else: + raise NotImplementedError( + f"Model type {self.model_type} has not been implemented. " + "Contribution is welcome!" + ) + + def get_fitting(self) -> Union[Fitting, dict]: + """Get the fitting(s).""" + return {} + + def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + """Get the loss function(s).""" + # loss should be never used for a frozen model + return + + def get_rcut(self): + return self.model.get_rcut() + + def get_ntypes(self) -> int: + return self.model.get_ntypes() + + def data_stat(self, data): + pass + + def init_variables( + self, + graph: tf.Graph, + graph_def: tf.GraphDef, + model_type: str = "original_model", + suffix: str = "", + ) -> None: + """Init the embedding net variables with the given frozen model. + + Parameters + ---------- + graph : tf.Graph + The input frozen model graph + graph_def : tf.GraphDef + The input frozen model graph_def + model_type : str + the type of the model + suffix : str + suffix to name scope + """ + pass + + def enable_compression(self, suffix: str = "") -> None: + """Enable compression. + + Parameters + ---------- + suffix : str + suffix to name scope + """ + pass + + def get_type_map(self) -> list: + """Get the type map.""" + return self.model.get_type_map() diff --git a/deepmd/model/linear.py b/deepmd/model/linear.py new file mode 100644 index 0000000000..6399766662 --- /dev/null +++ b/deepmd/model/linear.py @@ -0,0 +1,250 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from enum import ( + Enum, +) +from functools import ( + lru_cache, +) +from typing import ( + List, + Optional, + Union, +) + +from deepmd.env import ( + GLOBAL_TF_FLOAT_PRECISION, + MODEL_VERSION, + tf, +) +from deepmd.fit.fitting import ( + Fitting, +) +from deepmd.loss.loss import ( + Loss, +) + +from .model import ( + Model, +) + + +class LinearModel(Model): + """Linear model make linear combinations of several existing models. + + Parameters + ---------- + models : list[dict] + A list of models to be combined. + weights : list[float] or str + If the type is list[float], a list of weights for each model. + If "mean", the weights are set to be 1 / len(models). + If "sum", the weights are set to be 1. + """ + + def __init__(self, models: List[dict], weights: List[float], **kwargs): + super().__init__(**kwargs) + self.models = [Model(**model) for model in models] + if isinstance(weights, list): + if len(weights) != len(models): + raise ValueError( + "The length of weights is not equal to the number of models" + ) + self.weights = weights + elif weights == "mean": + self.weights = [1 / len(models) for _ in range(len(models))] + elif weights == "sum": + self.weights = [1 for _ in range(len(models))] + # TODO: add more weights, for example, so-called committee models + else: + raise ValueError(f"Invalid weights {weights}") + + def get_fitting(self) -> Union[Fitting, dict]: + """Get the fitting(s).""" + return { + f"model{ii}": model.get_fitting() for ii, model in enumerate(self.models) + } + + def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: + """Get the loss function(s).""" + # the first model that is not None, or None if all models are None + for model in self.models: + loss = model.get_loss(loss, lr) + if loss is not None: + return loss + return None + + def get_rcut(self): + return max([model.get_rcut() for model in self.models]) + + @lru_cache(maxsize=1) + def get_ntypes(self) -> int: + # check if all models have the same ntypes + for model in self.models: + if model.get_ntypes() != self.models[0].get_ntypes(): + raise ValueError("Models have different ntypes") + return self.models[0].get_ntypes() + + def data_stat(self, data): + for model in self.models: + model.data_stat(data) + + def init_variables( + self, + graph: tf.Graph, + graph_def: tf.GraphDef, + model_type: str = "original_model", + suffix: str = "", + ) -> None: + """Init the embedding net variables with the given frozen model. + + Parameters + ---------- + graph : tf.Graph + The input frozen model graph + graph_def : tf.GraphDef + The input frozen model graph_def + model_type : str + the type of the model + suffix : str + suffix to name scope + """ + for ii, model in enumerate(self.models): + model.init_variables( + graph, graph_def, model_type, suffix=f"_model{ii}{suffix}" + ) + + def enable_compression(self, suffix: str = "") -> None: + """Enable compression. + + Parameters + ---------- + suffix : str + suffix to name scope + """ + for ii, model in enumerate(self.models): + model.enable_compression(suffix=f"_model{ii}{suffix}") + + def get_type_map(self) -> list: + """Get the type map.""" + return self.models[0].get_type_map() + + +class LinearEnergyModel(LinearModel): + """Linear energy model make linear combinations of several existing energy models.""" + + model_type = "ener" + + def build( + self, + coord_: tf.Tensor, + atype_: tf.Tensor, + natoms: tf.Tensor, + box: tf.Tensor, + mesh: tf.Tensor, + input_dict: dict, + frz_model: Optional[str] = None, + ckpt_meta: Optional[str] = None, + suffix: str = "", + reuse: Optional[Union[bool, Enum]] = None, + ) -> dict: + """Build the model. + + Parameters + ---------- + coord_ : tf.Tensor + The coordinates of atoms + atype_ : tf.Tensor + The atom types of atoms + natoms : tf.Tensor + The number of atoms + box : tf.Tensor + The box vectors + mesh : tf.Tensor + The mesh vectors + input_dict : dict + The input dict + frz_model : str, optional + The path to the frozen model + ckpt_meta : str, optional + The path to the checkpoint and meta file + suffix : str, optional + The suffix of the scope + reuse : bool or tf.AUTO_REUSE, optional + Whether to reuse the variables + + Returns + ------- + dict + The output dict + """ + with tf.variable_scope("model_attr" + suffix, reuse=reuse): + t_tmap = tf.constant( + " ".join(self.get_type_map()), name="tmap", dtype=tf.string + ) + t_mt = tf.constant(self.model_type, name="model_type", dtype=tf.string) + t_ver = tf.constant(MODEL_VERSION, name="model_version", dtype=tf.string) + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): + # non zero not supported + t_dfparam = tf.constant(0, name="dfparam", dtype=tf.int32) + t_daparam = tf.constant(0, name="daparam", dtype=tf.int32) + with tf.variable_scope("descrpt_attr" + suffix, reuse=reuse): + t_ntypes = tf.constant(self.get_ntypes(), name="ntypes", dtype=tf.int32) + t_rcut = tf.constant( + self.get_rcut(), name="rcut", dtype=GLOBAL_TF_FLOAT_PRECISION + ) + + subdicts = [] + for ii, model in enumerate(self.models): + subdict = model.build( + coord_, + atype_, + natoms, + box, + mesh, + input_dict, + frz_model=frz_model, + ckpt_meta=ckpt_meta, + suffix=f"_model{ii}{suffix}", + reuse=reuse, + ) + subdicts.append(subdict) + t_weight = tf.constant(self.weights, dtype=GLOBAL_TF_FLOAT_PRECISION) + + model_dict = {} + # energy shape is (n_batch,), other shapes are (n_batch, -1) + energy = tf.reduce_sum( + tf.stack([mm["energy"] for mm in subdicts], axis=0) * t_weight[:, None], + axis=0, + ) + force = tf.reduce_sum( + tf.stack([mm["force"] for mm in subdicts], axis=0) + * t_weight[:, None, None], + axis=0, + ) + virial = tf.reduce_sum( + tf.stack([mm["virial"] for mm in subdicts], axis=0) + * t_weight[:, None, None], + axis=0, + ) + atom_ener = tf.reduce_sum( + tf.stack([mm["atom_ener"] for mm in subdicts], axis=0) + * t_weight[:, None, None], + axis=0, + ) + atom_virial = tf.reduce_sum( + tf.stack([mm["atom_virial"] for mm in subdicts], axis=0) + * t_weight[:, None, None], + axis=0, + ) + + model_dict["energy"] = tf.identity(energy, name="o_energy" + suffix) + model_dict["force"] = tf.identity(force, name="o_force" + suffix) + model_dict["virial"] = tf.identity(virial, name="o_virial" + suffix) + model_dict["atom_ener"] = tf.identity(atom_ener, name="o_atom_energy" + suffix) + model_dict["atom_virial"] = tf.identity( + atom_virial, name="o_atom_virial" + suffix + ) + + model_dict["coord"] = coord_ + model_dict["atype"] = atype_ + return model_dict diff --git a/deepmd/model/model.py b/deepmd/model/model.py index 26d5a6fbb1..a06a3141c0 100644 --- a/deepmd/model/model.py +++ b/deepmd/model/model.py @@ -82,6 +82,12 @@ def __new__(cls, *args, **kwargs): if cls is Model: # init model # infer model type by fitting_type + from deepmd.model.frozen import ( + FrozenModel, + ) + from deepmd.model.linear import ( + LinearEnergyModel, + ) from deepmd.model.multi import ( MultiModel, ) @@ -96,6 +102,10 @@ def __new__(cls, *args, **kwargs): cls = MultiModel elif model_type == "pairwise_dprc": cls = PairwiseDPRc + elif model_type == "frozen": + cls = FrozenModel + elif model_type == "linear_ener": + cls = LinearEnergyModel else: raise ValueError(f"unknown model type: {model_type}") return cls.__new__(cls, *args, **kwargs) @@ -393,11 +403,11 @@ def get_numb_dos(self) -> Union[int, dict]: return 0 @abstractmethod - def get_fitting(self) -> Union[str, dict]: + def get_fitting(self) -> Union[Fitting, dict]: """Get the fitting(s).""" @abstractmethod - def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + def get_loss(self, loss: dict, lr) -> Optional[Union[Loss, dict]]: """Get the loss function(s).""" @abstractmethod diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index a6ac96dab4..b322336b39 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -325,6 +325,9 @@ def _build_lr(self): log.info("built lr") def _build_loss(self): + if self.stop_batch == 0: + # l2 is not used if stop_batch is zero + return None, None if not self.multi_task_mode: l2_l, l2_more = self.loss.build( self.learning_rate, @@ -449,6 +452,11 @@ def _build_optimizer(self, fitting_key=None): return optimizer def _build_training(self): + if self.stop_batch == 0: + # self.train_op is not used if stop_batch is zero + self.train_op = None + return + trainable_variables = tf.trainable_variables() if not self.multi_task_mode: diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index b67722bd89..b38f8c8063 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -799,6 +799,7 @@ def model_args(exclude_hybrid=False): hybrid_models.extend( [ pairwise_dprc(), + linear_ener_model_args(), ] ) return Argument( @@ -871,6 +872,7 @@ def model_args(exclude_hybrid=False): [ standard_model_args(), multi_model_args(), + frozen_model_args(), *hybrid_models, ], optional=True, @@ -945,6 +947,46 @@ def pairwise_dprc() -> Argument: return ca +def frozen_model_args() -> Argument: + doc_model_file = "Path to the frozen model file." + ca = Argument( + "frozen", + dict, + [ + Argument("model_file", str, optional=False, doc=doc_model_file), + ], + ) + return ca + + +def linear_ener_model_args() -> Argument: + doc_weights = ( + "If the type is list of float, a list of weights for each model. " + 'If "mean", the weights are set to be 1 / len(models). ' + 'If "sum", the weights are set to be 1.' + ) + models_args = model_args(exclude_hybrid=True) + models_args.name = "models" + models_args.fold_subdoc = True + models_args.set_dtype(list) + models_args.set_repeat(True) + models_args.doc = "The sub-models." + ca = Argument( + "linear_ener", + dict, + [ + models_args, + Argument( + "weights", + [list, str], + optional=False, + doc=doc_weights, + ), + ], + ) + return ca + + # --- Learning rate configurations: --- # def learning_rate_exp(): doc_start_lr = "The learning rate at the start of the training." diff --git a/doc/model/index.md b/doc/model/index.md index d649df1442..4ef508ec1b 100644 --- a/doc/model/index.md +++ b/doc/model/index.md @@ -16,3 +16,4 @@ - [Train a Deep Potential model using `type embedding` approach](train-se-e2-a-tebd.md) - [Deep potential long-range](dplr.md) - [Deep Potential - Range Correction (DPRc)](dprc.md) +- [Linear model](linear.md) diff --git a/doc/model/index.rst b/doc/model/index.rst index 6a01a3b015..6597ce1d21 100644 --- a/doc/model/index.rst +++ b/doc/model/index.rst @@ -19,3 +19,4 @@ Model train-se-a-mask dplr dprc + linear diff --git a/doc/model/linear.md b/doc/model/linear.md new file mode 100644 index 0000000000..b5e7c5c76a --- /dev/null +++ b/doc/model/linear.md @@ -0,0 +1,24 @@ +## Linear model + +One can linearly combine existing models with arbitrary coefficients: + +```json +"model": { + "type": "linear_ener", + "models": [ + { + "type": "frozen", + "model_file": "model0.pb" + }, + { + "type": "frozen", + "model_file": "model1.pb" + } + ], + "weights": [0.5, 0.5] +}, +``` + +{ref}`weights ` can be a list of floats, `mean`, or `sum`. + +To obtain the model, one needs to execute `dp train` to do a zero-step training with {ref}`numb_steps ` set to `0`, and then freeze the model with `dp freeze`. diff --git a/examples/water/linear/input.json b/examples/water/linear/input.json new file mode 100644 index 0000000000..e6d0e267f4 --- /dev/null +++ b/examples/water/linear/input.json @@ -0,0 +1,56 @@ +{ + "model": { + "type": "linear_ener", + "models": [ + { + "type": "frozen", + "model_file": "model0.pb" + }, + { + "type": "frozen", + "model_file": "model1.pb" + } + ], + "weights": "mean", + "_comment1": "that's all" + }, + + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-8, + "_comment2": "that's all" + }, + + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment3": " that's all" + }, + + "training": { + "training_data": { + "_comment4": "Currently there must be systems", + "_comment5": "TODO: support empty systems", + "systems": [ + "../data/data_0" + ], + "batch_size": "auto", + "_comment6": "that's all" + }, + "numb_steps": 0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 1000, + "_comment7": "that's all" + }, + + "_comment8": "that's all" +} diff --git a/source/tests/test_examples.py b/source/tests/test_examples.py index c9f71e4a81..d50ca5fee1 100644 --- a/source/tests/test_examples.py +++ b/source/tests/test_examples.py @@ -29,6 +29,7 @@ p_examples / "water" / "hybrid" / "input.json", p_examples / "water" / "dplr" / "train" / "dw.json", p_examples / "water" / "dplr" / "train" / "ener.json", + p_examples / "water" / "linear" / "input.json", p_examples / "nopbc" / "train" / "input.json", p_examples / "water_tensor" / "dipole" / "dipole_input.json", p_examples / "water_tensor" / "polar" / "polar_input.json", diff --git a/source/tests/test_linear_model.py b/source/tests/test_linear_model.py new file mode 100644 index 0000000000..13a2bc4850 --- /dev/null +++ b/source/tests/test_linear_model.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import os +import sys + +import numpy as np + +from deepmd.env import ( + GLOBAL_ENER_FLOAT_PRECISION, + GLOBAL_TF_FLOAT_PRECISION, + tf, +) +from deepmd.infer import ( + DeepPotential, +) +from deepmd.model.linear import ( + LinearEnergyModel, +) +from deepmd.utils.convert import ( + convert_pbtxt_to_pb, +) + +sys.path.insert(0, os.path.abspath(os.path.dirname(__file__))) +from common import ( + DataSystem, + del_data, + gen_data, + tests_path, +) + + +class TestLinearModel(tf.test.TestCase): + def setUp(self): + gen_data() + self.data_dir = "system" + with open(os.path.join(self.data_dir, "type_map.raw"), "w") as f: + f.write("O\nH") + self.pbtxts = [ + os.path.join(tests_path, "infer/deeppot.pbtxt"), + os.path.join(tests_path, "infer/deeppot-1.pbtxt"), + ] + self.graph_dirs = [pbtxt.replace("pbtxt", "pb") for pbtxt in self.pbtxts] + for pbtxt, pb in zip(self.pbtxts, self.graph_dirs): + convert_pbtxt_to_pb(pbtxt, pb) + self.graphs = [DeepPotential(pb) for pb in self.graph_dirs] + + def test_linear_ener_model(self): + numb_test = 1 + data = DataSystem([self.data_dir], "set", 1, 1, 6, run_opt=None) + test_data = data.get_test() + + model = LinearEnergyModel( + models=[ + { + "type": "frozen", + "model_file": model_file, + } + for model_file in self.graph_dirs + ], + weights="mean", + ) + + t_prop_c = tf.placeholder(tf.float32, [5], name="t_prop_c") + t_energy = tf.placeholder(GLOBAL_ENER_FLOAT_PRECISION, [None], name="t_energy") + t_coord = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name="i_coord") + t_type = tf.placeholder(tf.int32, [None], name="i_type") + t_natoms = tf.placeholder(tf.int32, [model.get_ntypes() + 2], 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") + is_training = tf.placeholder(tf.bool) + t_fparam = None + + model_pred = model.build( + t_coord, + t_type, + t_natoms, + t_box, + t_mesh, + t_fparam, + suffix="_linear_energy", + reuse=False, + ) + + energy = model_pred["energy"] + force = model_pred["force"] + virial = model_pred["virial"] + + feed_dict_test = { + t_prop_c: test_data["prop_c"], + t_energy: test_data["energy"][:numb_test], + t_coord: np.reshape(test_data["coord"][:numb_test, :], [-1]), + t_box: test_data["box"][:numb_test, :], + t_type: np.reshape(test_data["type"], [-1]), + t_natoms: test_data["natoms_vec"], + t_mesh: test_data["default_mesh"], + is_training: False, + } + sess = self.test_session().__enter__() + sess.run(tf.global_variables_initializer()) + [e, f, v] = sess.run([energy, force, virial], feed_dict=feed_dict_test) + e = np.reshape(e, [1, -1]) + f = np.reshape(f, [1, -1, 3]) + v = np.reshape(v, [1, 9]) + + es = [] + fs = [] + vs = [] + + for ii, graph in enumerate(self.graphs): + ei, fi, vi = graph.eval( + test_data["coord"][:numb_test, :], + test_data["box"][:numb_test, :], + np.reshape(test_data["type"], [-1]), + ) + es.append(ei) + fs.append(fi) + vs.append(vi) + + np.testing.assert_allclose(e, np.mean(es, axis=0), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(f, np.mean(fs, axis=0), rtol=1e-5, atol=1e-5) + np.testing.assert_allclose(v, np.mean(vs, axis=0), rtol=1e-5, atol=1e-5) + + def tearDown(self): + for pb in self.graph_dirs: + os.remove(pb) + del_data() From fe9858eb12fc7632c6aada8ad66b02b4ed413e4b Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 28 Aug 2023 19:59:09 -0400 Subject: [PATCH 2/3] mv fparam out of ener; it seems all DPs have fparam Signed-off-by: Jinzhe Zeng --- deepmd/model/frozen.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/deepmd/model/frozen.py b/deepmd/model/frozen.py index 689b07edd7..c1b7d0286e 100644 --- a/deepmd/model/frozen.py +++ b/deepmd/model/frozen.py @@ -113,14 +113,14 @@ def build( t_rcut = tf.constant( self.get_rcut(), name="rcut", dtype=GLOBAL_TF_FLOAT_PRECISION ) + with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): + t_dfparam = tf.constant( + self.model.get_dim_fparam(), name="dfparam", dtype=tf.int32 + ) + t_daparam = tf.constant( + self.model.get_dim_aparam(), name="daparam", dtype=tf.int32 + ) if self.model_type == "ener": - with tf.variable_scope("fitting_attr" + suffix, reuse=reuse): - t_dfparam = tf.constant( - self.model.get_dim_fparam(), name="dfparam", dtype=tf.int32 - ) - t_daparam = tf.constant( - self.model.get_dim_aparam(), name="daparam", dtype=tf.int32 - ) return { "energy": tf.identity(self.model.t_energy, name="o_energy" + suffix), "force": tf.identity(self.model.t_force, name="o_force" + suffix), From fe9b7b016276adb34136b1a4bdb234cb4dfb0cd2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 2 Sep 2023 22:50:17 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- doc/getting-started/quick_start.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/getting-started/quick_start.ipynb b/doc/getting-started/quick_start.ipynb index e743b5cf5c..31209ae381 100644 --- a/doc/getting-started/quick_start.ipynb +++ b/doc/getting-started/quick_start.ipynb @@ -1,5 +1,5 @@ { - "cells": [ + "cells": [ { "attachments": {}, "cell_type": "markdown",