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..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: @@ -452,10 +453,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 1649921270..8185fd1591 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.sw_rmin = sw_rmin self.sw_rmax = sw_rmax 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