From 0eb7847cc88a15dca25f7c0e4511b037eb587f78 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 11 Mar 2024 19:58:26 -0400 Subject: [PATCH 1/6] fix: remove model_def_script from AtomicModel After #3438, `model_def_script` is no more saved in AtomicModel. Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/atomic_model/make_base_atomic_model.py | 4 ---- deepmd/dpmodel/model/base_model.py | 7 +++++++ deepmd/dpmodel/model/make_model.py | 4 ---- deepmd/pt/model/atomic_model/base_atomic_model.py | 3 --- deepmd/pt/model/atomic_model/dp_atomic_model.py | 1 - deepmd/pt/model/atomic_model/linear_atomic_model.py | 1 - deepmd/pt/model/atomic_model/pairtab_atomic_model.py | 1 - deepmd/pt/model/model/make_model.py | 5 ----- deepmd/pt/model/model/model.py | 6 ++++++ 9 files changed, 13 insertions(+), 19 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index e3d6d8bcd1..dfbb4e435a 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -178,10 +178,6 @@ def do_grad_(self, var_name: str, base: str) -> bool: return self.fitting_output_def()[var_name].c_differentiable return self.fitting_output_def()[var_name].r_differentiable - def get_model_def_script(self) -> str: - # TODO: implement this method; saved to model - raise NotImplementedError - setattr(BAM, fwd_method_name, BAM.fwd) delattr(BAM, "fwd") diff --git a/deepmd/dpmodel/model/base_model.py b/deepmd/dpmodel/model/base_model.py index 95c448442e..5169d1b5fe 100644 --- a/deepmd/dpmodel/model/base_model.py +++ b/deepmd/dpmodel/model/base_model.py @@ -172,3 +172,10 @@ class BaseModel(make_base_model()): deepmd.dpmodel.model.base_model.BaseBaseModel Backend-independent BaseModel class. """ + + def __init__(self) -> None: + self.model_def_script = "" + + def get_model_def_script(self) -> str: + """Get the model definition script.""" + return self.model_def_script diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 6022fd3e73..c14f098a1d 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -452,10 +452,6 @@ def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.atomic_model.get_nnei() - def get_model_def_script(self) -> str: - """Get the model definition script.""" - return self.atomic_model.get_model_def_script() - def get_sel(self) -> List[int]: """Returns the number of selected atoms for each type.""" return self.atomic_model.get_sel() diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py index d045220b6e..d3a1cfb459 100644 --- a/deepmd/pt/model/atomic_model/base_atomic_model.py +++ b/deepmd/pt/model/atomic_model/base_atomic_model.py @@ -58,9 +58,6 @@ def reinit_pair_exclude( else: self.pair_excl = PairExcludeMask(self.get_ntypes(), self.pair_exclude_types) - def get_model_def_script(self) -> str: - return self.model_def_script - def atomic_output_def(self) -> FittingOutputDef: old_def = self.fitting_output_def() if self.atom_excl is None: diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index ec08850524..e764be6998 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -55,7 +55,6 @@ def __init__( **kwargs, ): torch.nn.Module.__init__(self) - 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 3fb3ee90dd..fa4d875e2c 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -360,7 +360,6 @@ def __init__( ): models = [dp_model, zbl_model] super().__init__(models, type_map, **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 db0a2efa4a..897a7cd9d4 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -78,7 +78,6 @@ def __init__( **kwargs, ): torch.nn.Module.__init__(self) - 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/model/make_model.py b/deepmd/pt/model/model/make_model.py index 0a5f286040..167ad81923 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -469,11 +469,6 @@ def get_nnei(self) -> int: """Returns the total number of selected neighboring atoms in the cut-off radius.""" return self.atomic_model.get_nnei() - @torch.jit.export - def get_model_def_script(self) -> str: - """Get the model definition script.""" - return self.atomic_model.get_model_def_script() - def atomic_output_def(self) -> FittingOutputDef: """Get the output def of the atomic model.""" return self.atomic_model.atomic_output_def() diff --git a/deepmd/pt/model/model/model.py b/deepmd/pt/model/model/model.py index 3d4618449a..a62050b2d1 100644 --- a/deepmd/pt/model/model/model.py +++ b/deepmd/pt/model/model/model.py @@ -17,6 +17,7 @@ class BaseModel(torch.nn.Module, make_base_model()): def __init__(self, *args, **kwargs): """Construct a basic model for different tasks.""" torch.nn.Module.__init__(self) + self.model_def_script = "" def compute_or_load_stat( self, @@ -39,3 +40,8 @@ def compute_or_load_stat( The path to the statistics files. """ raise NotImplementedError + + @torch.jit.export + def get_model_def_script(self) -> str: + """Get the model definition script.""" + return self.model_def_script From 72fbcfd97afe558222b676f9b89e1633cf83767e Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 11 Mar 2024 20:00:55 -0400 Subject: [PATCH 2/6] init for dpmodel make_model Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/model/make_model.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index c14f098a1d..68889ad331 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -73,6 +73,7 @@ def __init__( atomic_model_: Optional[T_AtomicModel] = None, **kwargs, ): + BaseModel.__init__(self) if atomic_model_ is not None: self.atomic_model: T_AtomicModel = atomic_model_ else: From 5b718836af87f426aeeecd2e67417c2a4195425a Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 11 Mar 2024 21:32:41 -0400 Subject: [PATCH 3/6] frozen_model Signed-off-by: Jinzhe Zeng --- deepmd/dpmodel/utils/network.py | 4 + deepmd/pt/model/model/__init__.py | 4 + deepmd/pt/model/model/frozen.py | 174 +++++++++++++++++++ deepmd/tf/fit/ener.py | 2 +- deepmd/tf/model/frozen.py | 35 +++- deepmd/tf/model/model.py | 3 +- source/tests/consistent/model/test_frozen.py | 158 +++++++++++++++++ source/tests/infer/deeppot.dp | Bin 0 -> 43424 bytes 8 files changed, 377 insertions(+), 3 deletions(-) create mode 100644 deepmd/pt/model/model/frozen.py create mode 100644 source/tests/consistent/model/test_frozen.py create mode 100644 source/tests/infer/deeppot.dp diff --git a/deepmd/dpmodel/utils/network.py b/deepmd/dpmodel/utils/network.py index 6206367b1b..1747e25527 100644 --- a/deepmd/dpmodel/utils/network.py +++ b/deepmd/dpmodel/utils/network.py @@ -230,6 +230,10 @@ def deserialize(cls, data: dict) -> "NativeLayer": variables.get("b", None), variables.get("idt", None), ) + if obj.b is not None: + obj.b = obj.b.ravel() + if obj.idt is not None: + obj.idt = obj.idt.ravel() obj.check_shape_consistency() return obj diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 3098dc7677..f93ec88bde 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -37,6 +37,9 @@ from .ener_model import ( EnergyModel, ) +from .frozen import ( + FrozenModel, +) from .make_hessian_model import ( make_hessian_model, ) @@ -173,6 +176,7 @@ def get_model(model_params): "get_model", "DPModel", "EnergyModel", + "FrozenModel", "SpinModel", "SpinEnergyModel", "DPZBLModel", diff --git a/deepmd/pt/model/model/frozen.py b/deepmd/pt/model/model/frozen.py new file mode 100644 index 0000000000..e3dcd389bb --- /dev/null +++ b/deepmd/pt/model/model/frozen.py @@ -0,0 +1,174 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import tempfile +from typing import ( + Dict, + List, + Optional, +) + +import torch + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, +) +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) +from deepmd.pt.model.model.model import ( + BaseModel, +) + + +@BaseModel.register("frozen") +class FrozenModel(BaseModel): + """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 + if model_file.endswith(".pth"): + self.model = torch.jit.load(model_file) + else: + # try to convert from other formats + with tempfile.NamedTemporaryFile(suffix=".pth") as f: + convert_backend(INPUT=model_file, OUTPUT=f.name) + self.model = torch.jit.load(f.name) + + @torch.jit.export + def fitting_output_def(self) -> FittingOutputDef: + """Get the output def of developer implemented atomic models.""" + return self.model.fitting_output_def() + + @torch.jit.export + def get_rcut(self) -> float: + """Get the cut-off radius.""" + return self.model.get_rcut() + + @torch.jit.export + def get_type_map(self) -> List[str]: + """Get the type map.""" + return self.model.get_type_map() + + @torch.jit.export + def get_sel(self) -> List[int]: + """Returns the number of selected atoms for each type.""" + return self.model.get_sel() + + @torch.jit.export + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this atomic model.""" + return self.model.get_dim_fparam() + + @torch.jit.export + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this atomic model.""" + return self.model.get_dim_aparam() + + @torch.jit.export + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ + return self.model.get_sel_type() + + @torch.jit.export + def is_aparam_nall(self) -> bool: + """Check whether the shape of atomic parameters is (nframes, nall, ndim). + + If False, the shape is (nframes, nloc, ndim). + """ + return self.model.is_aparam_nall() + + @torch.jit.export + def mixed_types(self) -> bool: + """If true, the model + 1. assumes total number of atoms aligned across frames; + 2. uses a neighbor list that does not distinguish different atomic types. + + If false, the model + 1. assumes total number of atoms of each atom type aligned across frames; + 2. uses a neighbor list that distinguishes different atomic types. + + """ + return self.model.mixed_types() + + @torch.jit.export + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + return self.model.forward( + coord, + atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + + @torch.jit.export + def get_model_def_script(self) -> str: + """Get the model definition script.""" + # try to use the original script instead of "frozen model" + # Note: this cannot change the script of the parent model + # it may still try to load hard-coded filename, which might + # be a problem + return self.model.get_model_def_script() + + def serialize(self) -> dict: + from deepmd.pt.model.model import ( + get_model, + ) + + # try to recover the original model + model_def_script = json.loads(self.get_model_def_script()) + model = get_model(model_def_script) + model.load_state_dict(self.model.state_dict()) + return model.serialize() + + @classmethod + def deserialize(cls, data: dict): + raise RuntimeError("Should not touch here.") + + @torch.jit.export + def get_nnei(self) -> int: + """Returns the total number of selected neighboring atoms in the cut-off radius.""" + return self.model.get_nnei() + + @torch.jit.export + def get_nsel(self) -> int: + """Returns the total number of selected neighboring atoms in the cut-off radius.""" + return self.model.get_nsel() + + @classmethod + def update_sel(cls, global_jdata: dict, local_jdata: dict): + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + global_jdata : dict + The global data, containing the training section + local_jdata : dict + The local data refer to the current class + """ + return local_jdata + + @torch.jit.export + def model_output_type(self) -> str: + """Get the output type for the model.""" + return self.model.model_output_type() diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 780ae76c96..883ba7d2b8 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -213,7 +213,7 @@ def __init__( else: self.atom_ener.append(None) self.useBN = False - self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64) + self.bias_atom_e = np.zeros((self.ntypes, 1), dtype=np.float64) # data requirement if self.numb_fparam > 0: add_data_requirement( diff --git a/deepmd/tf/model/frozen.py b/deepmd/tf/model/frozen.py index 1933690ca7..86676bfe0b 100644 --- a/deepmd/tf/model/frozen.py +++ b/deepmd/tf/model/frozen.py @@ -1,4 +1,7 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +import tempfile from enum import ( Enum, ) @@ -7,6 +10,9 @@ Union, ) +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) from deepmd.infer.deep_pot import ( DeepPot, ) @@ -24,6 +30,10 @@ from deepmd.tf.loss.loss import ( Loss, ) +from deepmd.tf.utils.graph import ( + get_tensor_by_name_from_graph, + load_graph_def, +) from .model import ( Model, @@ -43,7 +53,14 @@ class FrozenModel(Model): def __init__(self, model_file: str, **kwargs): super().__init__(**kwargs) self.model_file = model_file - self.model = DeepPotential(model_file) + if not model_file.endswith(".pb"): + # try to convert from other formats + with tempfile.NamedTemporaryFile( + suffix=".pb", dir=os.curdir, delete=False + ) as f: + convert_backend(INPUT=model_file, OUTPUT=f.name) + self.model_file = f.name + self.model = DeepPotential(self.model_file) if isinstance(self.model, DeepPot): self.model_type = "ener" else: @@ -228,3 +245,19 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): """ # we don't know how to compress it, so no neighbor statistics here return local_jdata + + def serialize(self, suffix: str = "") -> dict: + # try to recover the original model + # the current graph contains a prefix "load", + # so it cannot used to recover the original model + graph, graph_def = load_graph_def(self.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) + return model.serialize() + + @classmethod + def deserialize(cls, data: dict, suffix: str = ""): + raise RuntimeError("Should not touch here.") diff --git a/deepmd/tf/model/model.py b/deepmd/tf/model/model.py index ca660f8e95..a0e234a547 100644 --- a/deepmd/tf/model/model.py +++ b/deepmd/tf/model/model.py @@ -566,7 +566,8 @@ def deserialize(cls, data: dict, suffix: str = "") -> "Model": """ if cls is Model: return Model.get_class_by_type(data.get("type", "standard")).deserialize( - data + data, + suffix=suffix, ) raise NotImplementedError("Not implemented in class %s" % cls.__name__) diff --git a/source/tests/consistent/model/test_frozen.py b/source/tests/consistent/model/test_frozen.py new file mode 100644 index 0000000000..c6265500b4 --- /dev/null +++ b/source/tests/consistent/model/test_frozen.py @@ -0,0 +1,158 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import unittest +from typing import ( + Any, + Tuple, +) + +import numpy as np + +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) + +from ..common import ( + INSTALLED_PT, + INSTALLED_TF, + CommonTest, + parameterized, +) +from .common import ( + ModelTest, +) + +if INSTALLED_PT: + from deepmd.pt.model.model import BaseModel as FrozenModelPT + +else: + FrozenModelPT = None +if INSTALLED_TF: + from deepmd.tf.model.model import Model as FrozenModelTF +else: + FrozenModelTF = None +from pathlib import ( + Path, +) + +from deepmd.entrypoints.convert_backend import ( + convert_backend, +) +from deepmd.utils.argcheck import ( + model_args, +) + +original_model = str(Path(__file__).parent.parent.parent / "infer" / "deeppot.dp") +pt_model = "deeppot.pth" +tf_model = "deeppot.pb" +dp_model = original_model + + +def setUpModule(): + convert_backend( + INPUT=dp_model, + OUTPUT=tf_model, + ) + convert_backend( + INPUT=dp_model, + OUTPUT=pt_model, + ) + + +@parameterized((pt_model, tf_model, dp_model)) +class TestFrozen(CommonTest, ModelTest, unittest.TestCase): + @property + def data(self) -> dict: + (model_file,) = self.param + if not INSTALLED_PT and model_file.endswith(".pth"): + raise unittest.SkipTest("PyTorch is not installed") + if not INSTALLED_TF and model_file.endswith(".pb"): + raise unittest.SkipTest("TensorFlow is not installed") + return { + "type": "frozen", + "model_file": model_file, + } + + tf_class = FrozenModelTF + dp_class = None + pt_class = FrozenModelPT + args = model_args() + + def skip_dp(self): + return True + + 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) + + # 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, + 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, ...]: + # shape not matched. ravel... + if backend is self.RefBackend.DP: + return (ret["energy_redu"].ravel(), ret["energy"].ravel()) + elif backend is self.RefBackend.PT: + return (ret["energy"].ravel(), ret["atom_energy"].ravel()) + elif backend is self.RefBackend.TF: + return (ret[0].ravel(), ret[1].ravel()) + raise ValueError(f"Unknown backend: {backend}") diff --git a/source/tests/infer/deeppot.dp b/source/tests/infer/deeppot.dp new file mode 100644 index 0000000000000000000000000000000000000000..2f7d9e3f6f8175bcf013137b2989e4407c197e68 GIT binary patch literal 43424 zcmeHw30zF;|Nj&bqGU_B63J3&rlP`qib}2!iL}U66QjMREG;TaLZnbasI(|8N>eJI zX)o=iMD|G8u1n}r`k!geaWy617yowV-d@+V)(!|fv<6J`Z=hQNc#r$&8JKZcMnrUZ@+1O zo?VIHs^bRm=i1e`@F4sPPXnL*jiA5Z&m?d>)VOTt=i)IE_xpXgR75@%!r}+}VB<9! z^TI7ocD8JKJRe`z6>P0?Vi0FpL_7)`Ln74&vN?%y&iRDZZ(`j$a**bhJpv6QiebK@ z@r5nJ%aM?&rS?m)8;}7BNtwnOXpXK*8q*>dR41h#lGmw)r#{YuYSoz_k!Z78MWF@7 z?$jH%NHq_#Mobp&c%F;oR`~3AGOY;&TU5@h9$$_=X7O1urqzMXrLY@ghCPExjjJf4 zQk9UJOPbR0DjA)w^bO0OmsiBff?uy)64`r@28lpnIplkr7Blk|1oXM*VQ`Us(S(Wvca z{GQO7Ky>?I$ElZ+jd0;j!B}_C8hBYGFgr0j7p`qd%qf^v20~}1j}sc20~2X+wdrHa zQR)7>&B>GM!27;npdp6O@8fFJR?Mmdm9k9_9@~CGXWuMjSSsFt?XBu%hCX$O>QnBz z>G(%@7O`}|u=X7EAu~tmG=DLi@GaHx!1Py)axb15djnpteSV5;nt>*XY(GDVnhO@O zfl@n`*TU(c>zfxYWP(b9Yr8(CKfl3eUAZw69#)E_HH2lN_P3o0d6@q6NfvK2l#9{e)mACuw@bn6b@qY* zd2#5)VA+;)nEv*Zv2F!2LB}0d11eTgNS{8n4DD#xG%y&`p9oXc-SJ*ET+%lnQi{qoU`7r&BS|;lj?3#~4HoFeJeJ2SeR7PnQWBQY- zcWm9$lnD}i0v@5c(Wo|O+B(hpKs39prlWXjBMe%UU$-06Uy0R&hGR~-aOI)&U1@R| z*i}!vDS_!P%j%j#g;*te!f;WUHL(s(=p;^Ch2axYZ@(!(x)PjoTLqdi{cTxn5~d_^ z1AJ1frKm++Ni`L7Ppzydxw{-q&5PNWjsrDW#n2-+SjrtNd%UH0@ zPgk&dMiN>TC}G2IX@@f8&dxa?-UybO=$ZDx8W5I>Q#KCE25Q~Lqx=&qU{8?sYPab* zFn-Ih($u(cDBR#h1JL6&!^fQpPxzX~O)0n9F zrTVSgJJV5S9sSk(-AuHK-#U>I$V9R;t#%$}CZL4Q6?e2qX`oTC=Cn43|9tO*`bm`u zP<`fQ`$G)>R;d+&k16S}XZ*~IF1s_)Ae)JAmt**=j&~*dWB8|sYQ>zx@XtHrtF`#| zIJ9gOBS{&dq1AB)|KF0reeZ>5!_lofB|1k{ zfN^u7!$;{HcpWvd(c?lm6ft)#5|XZggc+$??il_gXqf9ZF9k%L<#yIU8IlPW-s|+B z7K#eO)g3VWw+jXgldC9$=WTLID=_?9KPrSg#_;#{icvnbBm<2(Ry#jPH3OyX=dbsk z!bB@RG;H;kPDhTHUMX34F;OfZ>#*YiCJH(~dHw*)%ZNOuHNb5=3l=lr^rmYyaHZWT zMP4Kwf>_xP>hhb=w1w+Od6(yapY)A2CHm`F!!@P`T*Bh);1Sb@Hj384@Z`g5UimzQ zt#dEaOEEtiJZ4hc63ovc3#Og)8hHbXMkbUVHoJ^!BpyA{naP6M{rAP!-N)?mL+Azd zoHTg-=fuA+6g8o;xd(nrt;zvTHtm0(->SQa`hlpkxase&yEfu=mjoARKez7s79Io+ zxbWDM5a%H1@5jTFn1X;9pZV9ugAkTmc!K@`Jce)pP{VN`nxBgYAw_>_zh%#a{)p!* zZu*UwFVKOU9+5)Z;Ih&abhe?K1ge@Xu)_P4Zg2NxOHx2GH6!xEx> z^KHh5?~m%6kBDxsM)%G4tB?M#afh%g!p{i$`|V1C^S2C=1Q$dEvRvK&5lPsUAimI7 z8hzhb4Cm8c(l_4>oX>zT*srh?fA>&he1)N(cYn*?2)iQc5^nnY?M)50H+dnTh6C!$Y}CFh1-LE$oBHYxti5KiBnN>;D0KCU@PC z67gssahtKkN$+hi-fbY(`#8}Xjh_^W6Ft4RX?S=2gx<&dOYe2hTrx&weBlFFx?%11~=C;sY-}^o{^8KJwxt zFFx|(WAD*{S3mLMBQHMo=lIz3J+GhjUK496C&v8T_nLa*F{4W?;`{m%cpxOcX3kCj zzHFL3c;-xp>d1P(;|{p9!}k~W&G)O1 z_Oz>i%g-u>I9T*IKl``vAPw&ic>G)T>mV{kKe{m+C3r2kdiBla9xZR2hY!y_c>ckQ z54`xu%MZN#$Q=iG^#QLw;ME7b`hYtw5Hv47@Ztk6KJelL5f8ZA{vIEOzw<2}ngD2< zu2j3mndBZVZ=8n@&pvqm!HW;P_{hr-y!^-=2YB@XuRh?_2fX@#J1!73FFx?%11~=C z;sX&6xZD07A3W0ip67Z;qm?5~yhnsZ!^X1)#?RKW!0=N(!3@h zdc!sbZdAC0Pn1YPixwQC>dvo%7j^nFISaGUiB-exGWfE<`e3c#A@6fYuSMZg?%+Jo zl~PJ_v^$3M&_nN9lS&jLKHaS7It$sd|jO@P{(9fDHV^B``7_#?--n&1BI{5ExH|h9v!eU)AI^Vgw)Vkw>n+o!Qq42 zO`pyPq&fNcl@;j+(Qw8P&o_%|fjZx)ey?vn3^LRCAfTNL9%CtvvW0OH< zb6T~%egM2IEIK*9B@fOuucgci1aK2LddMK=Jd8CqUD`aP5Y4H%Qgz~O3^M-k_|7B~ zH!$iXZ%LX2;5+p1ea*`05VhAcMW(0*ZJFUpQ?4rlIVx@R_Ir^)>R^yOHUpaWbp85@ zA*t|GZw37#tppTA?CbgH{xJQh|JW0g!_kbkhr^TvijmW{jUx`bN5ei5b>aKQ7a(Ld zt+a)|2DM(?c6X-9NzjR?iP*8V9$v_hudC~)p|hFeVpLPlqNzs}CW!{cLXD7(=-fOH z)L9($dh85ObWLUPs=d~Qs5;jBFVlC&(aYvh`yWbO1EZ?C$4P0i=)F$BE|bPo@H;r} zF{Q#Ctv{aRy*4}r9&0mK-PEfTfK8p?)YZ2Ag zC_~eP~>YuCHCrg4Wm1z0!x|!twbO z0a8Rg8k;pFFmy)-@{EwWa-;4%`eflX&P|{e4P;Cpm2NCY4oP7XtV&q`6GZ>E?8t_S z_W0^Meo?SeWc)4NH~vUiz-vtd^%N}HzP$CaY(B`lT1we1`W>~6A*n7-&qq({EQFNo zOHdAT)s34c^8mdyb1QO70ot>(A%`RifM%F^Sfo824O~*`VbLCqrsqxM3y*V0)tYm* z`;N>4rnX>WR(39W)#0ugI7z$k9k&x!mER^7K2h# zLDw&L=hSi*yf;g3J1`~_O_coRj-SH6AjS(WK@?maI>8)6112Yciq+lur{E|q1V6j{EWyx11XogO_;|8xYj)gLx4IOq-ue=TX8 z?O%;<7IYZK-9C)&b`&)Yd|rt<=T&%#)Cn-gE=ojRV|2uwlZ&AxUv5{Pa3xx5 z;DyE{M4(~ImEYewRRuxQr;0B*oB?|UJK*E~95iK_$uRHt*OA9)gR1rq;mB;W_{AG? zNl5Qtlh=UAY9#9R_Y6Qslu2(zjultHrWf=8MS3xse(r#0WqmQ|wyk`;wx|$|k4X-+xS0ZleC4l> z#$>|8#d@r~K0kh7WWg^eE$CB$a%h0-60b|4(Yhm;Uz7kVx{`D8K3WZNsz;Orl zpnKT-iZMeCf?w3ZxgiV2%Zg^8M!w-Kj?b>3kPJCzDJdqBZwXrMlaPw4HqOdDvX_bM zB+k3sJ;FrkH$zXWm~3Kw?fuoB?>+xl@6*fU@BRK-@6-3hPV-W zS<_7Au>%wQOPb6@Q@CgaE?SX`R^p=Pa?$g+Xk{*1r8`aLl8?+KADK%&GM9X0F8Rn@ z@{zgZBXh||=8})hB_D-LJ_?t76fXHFT=G%4!X+PtOFjyhdz}X4A;HJMHpJ)%Rsd4h>+6VFd zI&Itlbe(arzS0=rd=lHg!hm;g#`)xT^v$;o=Tmd-o9|a2?T9&1k(U{bN=~+S%xjB= zyXgr3ljxk(tV)t?3w2U7%58u(}eRsrlQ;xPXP6 zkJWvASCarbtYwD=*XM!Mgfp7zEjQ8LF)=BL{LwIZn`G^YElkuRlW8dEauul>4GSiV zWT24;2R{u8Igc!>u1DTWNCd5#I|FX1#>0H618*irMWBYwBXXWr9zex4LVK&ReI}`a z2hy$l@?rKF0lLkDOYnY*VcnG|fHr=zQo!P5SUWN7-Q>am@Vd4@XjyX}9CKW!_%sAy zk^HhvyJkngv<~3`e~J~NlOE5`n+C<8-HZbRA|yP3RJ&kZ=>&k8Swb7b)Y9QX$(&+C z>FelT`!*~1RiV33Wp0Gft49g zXHPT=N1cB~#?60OgpM5@DUF!XaOiPU?rx(CU~<{bw~bVTW;aO}EcOqABE64G-rLnf zjQ7&3)262(li}u%irvp5FSPXilE7FP(nvlRtnGmw(Eaz_9DWczJe^rF%e4><6&QbcXtGb~ET*!cT-#4X@ly^AS|T@Y?U8zW>zlhxS6I>bhZRs zr|(l!>qv!=S+h#FPp(2|ep|G@L!c0eMBaLGJ2DszbgQ3JU6PF4wx5khXR^`b&SJs8 zj~7Cy`ZM1LRyE)(tmU=yQ5ab8Q;d%WheFu~tF)-Asjz9oBNJvl3uxnCU2s*%M5Ekd zH2qIq0*_j^KVm**q2girTNPHEfInm`--;e8LMBgcy%x`8qWRm-w?!D8f+6t=la8#) zMnMzjPSEv_02sPGBh4=oM4V)DHz-7av(W`v>&7U!6QikAtA{-on+DAcF^K|m*}B^) zCzx<4UBJ)#Z5;AdTsG0rryAToma~@Dq`-+CgI---RtjTXC6cypI0b86NlqeK`RM7| z9Zw?suOX-XMyf|hYZ;)gb z?|K=?$M5r%)Y5=(fsw7c`;I_ye&O@e$P?yMf0G{RTa99xCz=L7IgAdxh)*2%v=Z^l z`+Ry`mWyIMtDBtHC!*4ocH1A_D27dugAXPOV*5<5S(M6}2&5Lj@P21#6-<(14LW=& z15^%(H0%q`K@9$B!+l$>qjZhsSrs|q$mrlGLn}%Wy5j0$Js_eQo#Ut6)l_4m>wBE; zD2`1=?_?Yl_O%t2fYvp$c51|<#MI+re_RTL%(%IGuKQ$zzym|+3rDhG zMXSO2)-id==k?(Dfu-fhY-hrw`6Wqczg^)Q%3U`kW4thYVbnepK563L@2^&X(i_Lz zC}uGVda?Av_NHP8-a-qLQ7%Fuhhl}MKTLtZkPm;}x|<1nL&VYYgaq`vMQfc-2NNRY zhZ-b4DuX9lX9vVIC4gq*jl9*^K9jMA#%%F9>2P@5QaL$lCK|qJ*1Sd7KGXSk7mUQQ zeWp<^mch@keJ0(eXErACap-kww$&SKpXrtkq-m}RK@Pv&TC=L40llap`E5}tL4x~t z?XPHK0ru;^eEYC{ri{!r7S~Oh(4(E;U}l{Tn*!fU8tUdE*8)vLb!?wW%rUHS@WTRB zB^Uj!ye<|QyV|zA*_93k_B8pVODy2;T<*m8ArY-oOZTz0wnsq+9t+otHNrw~5lI(p zpQ(Mk*9qGT*^n-|{ZGru6;O@T3m;45fO}}cQlZ3f81;6Bm>agwB)uWeC>h&lQdXvr z7GwKNzk2@uQTck-mEG%XqMjt$@2?}9zw(^Xzv^rPhp)BYkB^!oN00K{ ziQ#-{NBZWQf%6pw^v(Cbj(*F}@MV7fJIeRlkpzPylOUI9zg8maiAQ%o!{Pso2RHrw z_{_lZ5#hq;=la=_#Vd8#Hc&K%BNkx~VL0t7>11cg*h+PzNvcX(&}jBH7P3}!hNO(7 z3(e7qZfA>iE6B;q$y3R47Bm-%^6X7C3uOhexw5jH1x22sD6c@4$G91E8} zSu#aJo~%kyQB|5JCr=^EE3>RsIsqs(uKl9Q|X%tf`_q2}Y&J0P_xpML{lI&AcCmY6QNmY3{ zB^gO4wxC9eb7fR0*sW}7&W>!@j3_dSGRjy7bsOEul;~B&E~h!!(io-|*mRauYbP3Z zXGbdCmi&D`4Af@aDd#PBe}rl#4tnA$tqF=*zD zE>&TAVvB=KPIcL2x`XCu_xT=FjKzh@W;3;Pw#E4+F$_0jE46o|nbUD?TUy&u8FLk} zUYd;=&BB6iy9vA77kXGhV{El^v|4OuZH?cTO9QqRbQ?)kOcl1cs1!CTT^Dr;k16l- z)EF>b1K92GGxpNHxSJ*+1b$m?z0hq9G5_B{8E@ zF*c3@ot-c}VJ_jspkaoHtGQ1VaKs_5?_Cr%;%$j8g>1!eoI~_{q+r9 zxwloQ^s`lTC4X*1CvycI6L-wiyONa+6^lf!T^Y33WQwtM*F8ItafLK8J3?4EG1!H~ zm+_CRAcWKYMEZ&$FBk+h^U7+C)5mcF#`zRwAtX4^BCV`0_QmfBU)u!lEo$8PyW zHp4QGX;(>Mk1a<5OK&#xZ8QraGvnc&&1Ab;2qyC#%mH?;mAg>hD>;U>UzFlYK7n@Erd`r3J3@ zZx$A~9eks#U@I3>V$T|bP-LGF_Rw6E z-u+vlfB3gTZ~m>&!@t>OWcNVPl{z_94Ocm+z_wjZL1o%iL~^PXyoZkMoOG3toL-`S zVHca-db;spp&zfHde$=bRJx;Sj~WNNplc%ws}?OyIa^TdVzg_2r|&%~HVd{QthT~; zqq^@haqmHWu|4&<)+9EdzNkWoaLu(P)g=%2ev=c8h83JxvDs(-{W7dzi=;YRGcer0 zs6;u@rw1~sHJqtG%ZWpV*Yn>hsQ-Vgs(Vr1t+>WkVnJn4 yF(acgOl{4m40CL|kc<`f|McQJcU?t226NXymp$7$QElwmyMbhRcX!Ud?f(PmE`N^z literal 0 HcmV?d00001 From 752a06d4e352c8608e4ea9d91f27b003d8aeb0b4 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 11 Mar 2024 21:34:27 -0400 Subject: [PATCH 4/6] docs Signed-off-by: Jinzhe Zeng --- deepmd/utils/argcheck.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index e822e18d50..57f1145d55 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1461,7 +1461,6 @@ def frozen_model_args() -> Argument: [ Argument("model_file", str, optional=False, doc=doc_model_file), ], - doc=doc_only_tf_supported, ) return ca From da48b244192897e6566f77a7202c54d7c522c526 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 12 Mar 2024 03:51:42 -0400 Subject: [PATCH 5/6] delete test files Signed-off-by: Jinzhe Zeng --- source/tests/consistent/model/test_frozen.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/source/tests/consistent/model/test_frozen.py b/source/tests/consistent/model/test_frozen.py index c6265500b4..a60a6abb3f 100644 --- a/source/tests/consistent/model/test_frozen.py +++ b/source/tests/consistent/model/test_frozen.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import os import unittest from typing import ( Any, @@ -42,8 +43,8 @@ ) original_model = str(Path(__file__).parent.parent.parent / "infer" / "deeppot.dp") -pt_model = "deeppot.pth" -tf_model = "deeppot.pb" +pt_model = "deeppot_for_consistent_frozen.pth" +tf_model = "deeppot_for_consistent_frozen.pb" dp_model = original_model @@ -58,6 +59,14 @@ def setUpModule(): ) +def tearDownModule(): + for model_file in (pt_model, tf_model): + try: + os.remove(model_file) + except FileNotFoundError: + pass + + @parameterized((pt_model, tf_model, dp_model)) class TestFrozen(CommonTest, ModelTest, unittest.TestCase): @property From a002a380d1bb1e76579c06a821caf1e3ae20a6c1 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Tue, 12 Mar 2024 04:11:04 -0400 Subject: [PATCH 6/6] reshape bias_atom_e to ntypes,1 Signed-off-by: Jinzhe Zeng --- deepmd/tf/fit/ener.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 883ba7d2b8..f8f5c3b346 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -213,7 +213,7 @@ def __init__( else: self.atom_ener.append(None) self.useBN = False - self.bias_atom_e = np.zeros((self.ntypes, 1), dtype=np.float64) + self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64) # data requirement if self.numb_fparam > 0: add_data_requirement( @@ -868,7 +868,7 @@ def deserialize(cls, data: dict, suffix: str = ""): data["nets"], suffix=suffix, ) - fitting.bias_atom_e = data["@variables"]["bias_atom_e"] + fitting.bias_atom_e = data["@variables"]["bias_atom_e"].ravel() if fitting.numb_fparam > 0: fitting.fparam_avg = data["@variables"]["fparam_avg"] fitting.fparam_inv_std = data["@variables"]["fparam_inv_std"] @@ -922,7 +922,7 @@ def serialize(self, suffix: str = "") -> dict: suffix=suffix, ), "@variables": { - "bias_atom_e": self.bias_atom_e, + "bias_atom_e": self.bias_atom_e.reshape(-1, 1), "fparam_avg": self.fparam_avg, "fparam_inv_std": self.fparam_inv_std, "aparam_avg": self.aparam_avg,