add linear models

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)

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"])):

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}")
 

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:

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

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,

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:

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

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]

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: