From 0eb7847cc88a15dca25f7c0e4511b037eb587f78 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Mon, 11 Mar 2024 19:58:26 -0400 Subject: [PATCH] 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