From 0d732ade343f976196f3a81c811e2dbb11bcb3e5 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Thu, 1 Jun 2023 22:59:20 -0400 Subject: [PATCH] refactor: uncouple Descriptor and Fitting from Trainer (#2549) Signed-off-by: Jinzhe Zeng --- deepmd/descriptor/descriptor.py | 5 + deepmd/descriptor/hybrid.py | 9 +- deepmd/descriptor/loc_frame.py | 7 +- deepmd/descriptor/se_a.py | 1 + deepmd/descriptor/se_a_ebd.py | 1 + deepmd/descriptor/se_a_ef.py | 1 + deepmd/descriptor/se_a_mask.py | 1 + deepmd/descriptor/se_atten.py | 6 + deepmd/descriptor/se_r.py | 1 + deepmd/descriptor/se_t.py | 1 + deepmd/fit/dipole.py | 30 ++ deepmd/fit/dos.py | 26 ++ deepmd/fit/ener.py | 35 +++ deepmd/fit/fitting.py | 23 ++ deepmd/fit/polar.py | 41 +++ deepmd/loss/dos.py | 1 + deepmd/loss/ener.py | 1 + deepmd/model/dos.py | 55 ++-- deepmd/model/ener.py | 114 ++++--- deepmd/model/model.py | 323 ++++++++++++++++++++ deepmd/model/multi.py | 190 ++++++++++-- deepmd/model/tensor.py | 45 ++- deepmd/train/trainer.py | 359 +++-------------------- deepmd/utils/argcheck.py | 78 +++-- deepmd/utils/type_embed.py | 1 + source/tests/.gitignore | 13 + source/tests/test_init_frz_model_spin.py | 5 + source/tests/test_model_spin.py | 7 + 28 files changed, 931 insertions(+), 449 deletions(-) create mode 100644 source/tests/.gitignore diff --git a/deepmd/descriptor/descriptor.py b/deepmd/descriptor/descriptor.py index 3542918cb3..90eda03d35 100644 --- a/deepmd/descriptor/descriptor.py +++ b/deepmd/descriptor/descriptor.py @@ -522,3 +522,8 @@ def build_type_exclude_mask( # same as inputs_i, (nsamples * natoms, ndescrpt) mask = tf.reshape(mask, [-1, ndescrpt]) return mask + + @property + def explicit_ntypes(self) -> bool: + """Explicit ntypes with type embedding.""" + return False diff --git a/deepmd/descriptor/hybrid.py b/deepmd/descriptor/hybrid.py index e1ddf5f2e9..834af56226 100644 --- a/deepmd/descriptor/hybrid.py +++ b/deepmd/descriptor/hybrid.py @@ -35,7 +35,9 @@ class DescrptHybrid(Descriptor): Build a descriptor from the concatenation of the list of descriptors. """ - def __init__(self, list: list, multi_task: bool = False, spin: Spin = None) -> None: + def __init__( + self, list: list, multi_task: bool = False, spin: Spin = None, **kwargs + ) -> None: """Constructor.""" # warning: list is conflict with built-in list descrpt_list = list @@ -387,3 +389,8 @@ def pass_tensors_from_frz_model( n_tensors = len(ii.get_tensor_names()) ii.pass_tensors_from_frz_model(*tensors[jj : jj + n_tensors]) jj += n_tensors + + @property + def explicit_ntypes(self) -> bool: + """Explicit ntypes with type embedding.""" + return any(ii.explicit_ntypes for ii in self.descrpt_list) diff --git a/deepmd/descriptor/loc_frame.py b/deepmd/descriptor/loc_frame.py index dd99f7fb57..a962641021 100644 --- a/deepmd/descriptor/loc_frame.py +++ b/deepmd/descriptor/loc_frame.py @@ -54,7 +54,12 @@ class DescrptLocFrame(Descriptor): """ def __init__( - self, rcut: float, sel_a: List[int], sel_r: List[int], axis_rule: List[int] + self, + rcut: float, + sel_a: List[int], + sel_r: List[int], + axis_rule: List[int], + **kwargs, ) -> None: """Constructor.""" self.sel_a = sel_a diff --git a/deepmd/descriptor/se_a.py b/deepmd/descriptor/se_a.py index 641210f0d1..69252a8fe6 100644 --- a/deepmd/descriptor/se_a.py +++ b/deepmd/descriptor/se_a.py @@ -164,6 +164,7 @@ def __init__( uniform_seed: bool = False, multi_task: bool = False, spin: Optional[Spin] = None, + **kwargs, ) -> None: """Constructor.""" if rcut < rcut_smth: diff --git a/deepmd/descriptor/se_a_ebd.py b/deepmd/descriptor/se_a_ebd.py index e0684403e1..f19548b240 100644 --- a/deepmd/descriptor/se_a_ebd.py +++ b/deepmd/descriptor/se_a_ebd.py @@ -87,6 +87,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", exclude_types: List[List[int]] = [], + **kwargs, ) -> None: """Constructor.""" DescrptSeA.__init__( diff --git a/deepmd/descriptor/se_a_ef.py b/deepmd/descriptor/se_a_ef.py index 0dbad403f4..f9f091af66 100644 --- a/deepmd/descriptor/se_a_ef.py +++ b/deepmd/descriptor/se_a_ef.py @@ -82,6 +82,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed=False, + **kwargs, ) -> None: """Constructor.""" self.descrpt_para = DescrptSeAEfLower( diff --git a/deepmd/descriptor/se_a_mask.py b/deepmd/descriptor/se_a_mask.py index 93256e3910..de77747da8 100644 --- a/deepmd/descriptor/se_a_mask.py +++ b/deepmd/descriptor/se_a_mask.py @@ -126,6 +126,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + **kwargs, ) -> None: """Constructor.""" self.sel_a = sel diff --git a/deepmd/descriptor/se_atten.py b/deepmd/descriptor/se_atten.py index b113e5c994..d5a2e565c3 100644 --- a/deepmd/descriptor/se_atten.py +++ b/deepmd/descriptor/se_atten.py @@ -110,6 +110,7 @@ def __init__( attn_dotr: bool = True, attn_mask: bool = False, multi_task: bool = False, + **kwargs, ) -> None: if not set_davg_zero: warnings.warn( @@ -1065,3 +1066,8 @@ def build_type_exclude_mask( # same as inputs_i, (nsamples * natoms, ndescrpt) mask = tf.reshape(mask, [-1, ndescrpt]) return mask + + @property + def explicit_ntypes(self) -> bool: + """Explicit ntypes with type embedding.""" + return True diff --git a/deepmd/descriptor/se_r.py b/deepmd/descriptor/se_r.py index 7ef5f47347..74fc66c351 100644 --- a/deepmd/descriptor/se_r.py +++ b/deepmd/descriptor/se_r.py @@ -97,6 +97,7 @@ def __init__( uniform_seed: bool = False, multi_task: bool = False, spin: Optional[Spin] = None, + **kwargs, ) -> None: """Constructor.""" if rcut < rcut_smth: diff --git a/deepmd/descriptor/se_t.py b/deepmd/descriptor/se_t.py index 7615275a82..e12561f028 100644 --- a/deepmd/descriptor/se_t.py +++ b/deepmd/descriptor/se_t.py @@ -90,6 +90,7 @@ def __init__( precision: str = "default", uniform_seed: bool = False, multi_task: bool = False, + **kwargs, ) -> None: """Constructor.""" if rcut < rcut_smth: diff --git a/deepmd/fit/dipole.py b/deepmd/fit/dipole.py index 9f497584ae..a814242bd1 100644 --- a/deepmd/fit/dipole.py +++ b/deepmd/fit/dipole.py @@ -16,6 +16,12 @@ from deepmd.fit.fitting import ( Fitting, ) +from deepmd.loss.loss import ( + Loss, +) +from deepmd.loss.tensor import ( + TensorLoss, +) from deepmd.utils.graph import ( get_fitting_net_variables_from_graph_def, ) @@ -60,6 +66,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + **kwargs, ) -> None: """Constructor.""" self.ntypes = descrpt.get_ntypes() @@ -296,3 +303,26 @@ def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """ self.mixed_prec = mixed_prec self.fitting_precision = get_precision(mixed_prec["output_prec"]) + + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function. + + Parameters + ---------- + loss : dict + the loss dict + lr : LearningRateExp + the learning rate + + Returns + ------- + Loss + the loss function + """ + return TensorLoss( + loss, + model=self, + tensor_name="dipole", + tensor_size=3, + label_name="dipole", + ) diff --git a/deepmd/fit/dos.py b/deepmd/fit/dos.py index 4bb1838af1..1dbfc2c0c0 100644 --- a/deepmd/fit/dos.py +++ b/deepmd/fit/dos.py @@ -19,6 +19,12 @@ from deepmd.fit.fitting import ( Fitting, ) +from deepmd.loss.dos import ( + DOSLoss, +) +from deepmd.loss.loss import ( + Loss, +) from deepmd.nvnmd.fit.ener import ( one_layer_nvnmd, ) @@ -98,6 +104,7 @@ def __init__( uniform_seed: bool = False, layer_name: Optional[List[Optional[str]]] = None, use_aparam_as_mask: bool = False, + **kwargs, ) -> None: """Constructor.""" # model param @@ -609,3 +616,22 @@ def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """ self.mixed_prec = mixed_prec self.fitting_precision = get_precision(mixed_prec["output_prec"]) + + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function. + + Parameters + ---------- + loss : dict + the loss dict + lr : LearningRateExp + the learning rate + + Returns + ------- + Loss + the loss function + """ + return DOSLoss( + **loss, starter_learning_rate=lr.start_lr(), numb_dos=self.get_numb_dos() + ) diff --git a/deepmd/fit/ener.py b/deepmd/fit/ener.py index f482173495..f7ed408c15 100644 --- a/deepmd/fit/ener.py +++ b/deepmd/fit/ener.py @@ -23,6 +23,14 @@ from deepmd.infer import ( DeepPotential, ) +from deepmd.loss.ener import ( + EnerDipoleLoss, + EnerSpinLoss, + EnerStdLoss, +) +from deepmd.loss.loss import ( + Loss, +) from deepmd.nvnmd.fit.ener import ( one_layer_nvnmd, ) @@ -137,6 +145,7 @@ def __init__( layer_name: Optional[List[Optional[str]]] = None, use_aparam_as_mask: bool = False, spin: Optional[Spin] = None, + **kwargs, ) -> None: """Constructor.""" # model param @@ -870,3 +879,29 @@ def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: """ self.mixed_prec = mixed_prec self.fitting_precision = get_precision(mixed_prec["output_prec"]) + + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function. + + Parameters + ---------- + loss : dict + The loss function parameters. + lr : LearningRateExp + The learning rate. + + Returns + ------- + Loss + The loss function. + """ + _loss_type = loss.pop("type", "ener") + loss["starter_learning_rate"] = lr.start_lr() + if _loss_type == "ener": + return EnerStdLoss(**loss) + elif _loss_type == "ener_dipole": + return EnerDipoleLoss(**loss) + elif _loss_type == "ener_spin": + return EnerSpinLoss(**loss, use_spin=self.spin.use_spin) + else: + raise RuntimeError("unknown loss type") diff --git a/deepmd/fit/fitting.py b/deepmd/fit/fitting.py index b0a9efe3be..803132452a 100644 --- a/deepmd/fit/fitting.py +++ b/deepmd/fit/fitting.py @@ -1,3 +1,6 @@ +from abc import ( + abstractmethod, +) from typing import ( Callable, ) @@ -5,6 +8,9 @@ from deepmd.env import ( tf, ) +from deepmd.loss.loss import ( + Loss, +) from deepmd.utils import ( Plugin, PluginVariant, @@ -78,3 +84,20 @@ def init_variables( "Fitting %s doesn't support initialization from the given variables!" % type(self).__name__ ) + + @abstractmethod + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function. + + Parameters + ---------- + loss : dict + the loss dict + lr : LearningRateExp + the learning rate + + Returns + ------- + Loss + the loss function + """ diff --git a/deepmd/fit/polar.py b/deepmd/fit/polar.py index 42b8d155f5..fa64e573c8 100644 --- a/deepmd/fit/polar.py +++ b/deepmd/fit/polar.py @@ -20,6 +20,12 @@ from deepmd.fit.fitting import ( Fitting, ) +from deepmd.loss.loss import ( + Loss, +) +from deepmd.loss.tensor import ( + TensorLoss, +) from deepmd.utils.graph import ( get_fitting_net_variables_from_graph_def, ) @@ -74,6 +80,7 @@ def __init__( activation_function: str = "tanh", precision: str = "default", uniform_seed: bool = False, + **kwargs, ) -> None: """Constructor.""" self.ntypes = descrpt.get_ntypes() @@ -490,6 +497,16 @@ def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: self.mixed_prec = mixed_prec self.fitting_precision = get_precision(mixed_prec["output_prec"]) + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function.""" + return TensorLoss( + loss, + model=self, + tensor_name="polar", + tensor_size=9, + label_name="polarizability", + ) + class GlobalPolarFittingSeA: r"""Fit the system polarizability with descriptor se_a. @@ -632,3 +649,27 @@ def enable_mixed_precision(self, mixed_prec: Optional[dict] = None) -> None: The mixed precision setting used in the embedding net """ self.polar_fitting.enable_mixed_precision(mixed_prec) + + def get_loss(self, loss: dict, lr) -> Loss: + """Get the loss function. + + Parameters + ---------- + loss : dict + the loss dict + lr : LearningRateExp + the learning rate + + Returns + ------- + Loss + the loss function + """ + return TensorLoss( + loss, + model=self, + tensor_name="global_polar", + tensor_size=9, + atomic=False, + label_name="polarizability", + ) diff --git a/deepmd/loss/dos.py b/deepmd/loss/dos.py index 5a7e67e29c..8e50738fd4 100644 --- a/deepmd/loss/dos.py +++ b/deepmd/loss/dos.py @@ -34,6 +34,7 @@ def __init__( limit_pref_acdf: float = 0.0, protect_value: float = 1e-8, log_fit: bool = False, + **kwargs, ) -> None: self.starter_learning_rate = starter_learning_rate self.numb_dos = numb_dos diff --git a/deepmd/loss/ener.py b/deepmd/loss/ener.py index 07c97b09bc..faae465930 100644 --- a/deepmd/loss/ener.py +++ b/deepmd/loss/ener.py @@ -45,6 +45,7 @@ def __init__( limit_pref_pf: float = 0.0, relative_f: Optional[float] = None, enable_atom_ener_coeff: bool = False, + **kwargs, ) -> None: self.starter_learning_rate = starter_learning_rate self.start_pref_e = start_pref_e diff --git a/deepmd/model/dos.py b/deepmd/model/dos.py index 0da709409f..6d90b7bd28 100644 --- a/deepmd/model/dos.py +++ b/deepmd/model/dos.py @@ -1,6 +1,7 @@ from typing import ( List, Optional, + Union, ) from deepmd.env import ( @@ -8,9 +9,12 @@ global_cvt_2_ener_float, tf, ) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) from .model import ( - Model, + StandardModel, ) from .model_stat import ( make_stat_input, @@ -18,15 +22,17 @@ ) -class DOSModel(Model): +class DOSModel(StandardModel): """DOS model. Parameters ---------- - descrpt + descriptor Descriptor - fitting + fitting_net Fitting net + type_embedding + Type embedding net type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. @@ -40,31 +46,28 @@ class DOSModel(Model): def __init__( self, - descrpt, - fitting, - typeebd=None, - type_map: List[str] = None, + descriptor: dict, + fitting_net: dict, + type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, + type_map: Optional[List[str]] = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, + **kwargs, ) -> None: """Constructor.""" - # descriptor - self.descrpt = descrpt - self.rcut = self.descrpt.get_rcut() - self.ntypes = self.descrpt.get_ntypes() + super().__init__( + descriptor=descriptor, + fitting_net=fitting_net, + type_embedding=type_embedding, + type_map=type_map, + data_stat_nbatch=data_stat_nbatch, + data_stat_protect=data_stat_protect, + **kwargs, + ) # fitting - self.fitting = fitting self.numb_dos = self.fitting.get_numb_dos() self.numb_fparam = self.fitting.get_numb_fparam() - # type embedding - self.typeebd = typeebd - # other inputs - if type_map is None: - self.type_map = [] - else: - self.type_map = type_map - self.data_stat_nbatch = data_stat_nbatch - self.data_stat_protect = data_stat_protect + self.numb_aparam = self.fitting.get_numb_aparam() def get_numb_dos(self): return self.numb_dos @@ -78,6 +81,14 @@ def get_ntypes(self): def get_type_map(self): return self.type_map + def get_numb_fparam(self) -> int: + """Get the number of frame parameters.""" + return self.numb_fparam + + def get_numb_aparam(self) -> int: + """Get the number of atomic parameters.""" + return self.numb_fparam + def data_stat(self, data): all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) diff --git a/deepmd/model/ener.py b/deepmd/model/ener.py index f9387c67fc..2491091dac 100644 --- a/deepmd/model/ener.py +++ b/deepmd/model/ener.py @@ -1,6 +1,7 @@ from typing import ( List, Optional, + Union, ) import numpy as np @@ -11,15 +12,18 @@ op_module, tf, ) -from deepmd.utils.pair_tab import ( - PairTab, +from deepmd.utils.data_system import ( + DeepmdDataSystem, ) from deepmd.utils.spin import ( Spin, ) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) from .model import ( - Model, + StandardModel, ) from .model_stat import ( make_stat_input, @@ -27,15 +31,17 @@ ) -class EnerModel(Model): +class EnerModel(StandardModel): """Energy model. Parameters ---------- - descrpt + descriptor Descriptor - fitting + fitting_net Fitting net + type_embedding + Type embedding net type_map Mapping atom type to the name (str) of the type. For example `type_map[1]` gives the name of the type 1. @@ -51,15 +57,19 @@ class EnerModel(Model): The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. sw_rmin The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + spin + spin + data_stat_nsample + The number of training samples in a system to compute and change the energy bias. """ model_type = "ener" def __init__( self, - descrpt, - fitting, - typeebd=None, + descriptor: dict, + fitting_net: dict, + type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, type_map: Optional[List[str]] = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, @@ -68,34 +78,27 @@ def __init__( sw_rmin: Optional[float] = None, sw_rmax: Optional[float] = None, spin: Optional[Spin] = None, + data_bias_nsample: int = 10, + **kwargs, ) -> None: """Constructor.""" - # descriptor - self.descrpt = descrpt - self.rcut = self.descrpt.get_rcut() - self.ntypes = self.descrpt.get_ntypes() - # fitting - self.fitting = fitting + super().__init__( + descriptor=descriptor, + fitting_net=fitting_net, + type_embedding=type_embedding, + type_map=type_map, + data_stat_nbatch=data_stat_nbatch, + data_bias_nsample=data_bias_nsample, + data_stat_protect=data_stat_protect, + use_srtab=use_srtab, + smin_alpha=smin_alpha, + sw_rmin=sw_rmin, + sw_rmax=sw_rmax, + spin=spin, + **kwargs, + ) self.numb_fparam = self.fitting.get_numb_fparam() - # type embedding - self.typeebd = typeebd - # spin - self.spin = spin - # other inputs - if type_map is None: - self.type_map = [] - else: - self.type_map = type_map - self.data_stat_nbatch = data_stat_nbatch - self.data_stat_protect = data_stat_protect - self.srtab_name = use_srtab - if self.srtab_name is not None: - self.srtab = PairTab(self.srtab_name) - self.smin_alpha = smin_alpha - self.sw_rmin = sw_rmin - self.sw_rmax = sw_rmax - else: - self.srtab = None + self.numb_aparam = self.fitting.get_numb_aparam() def get_rcut(self): return self.rcut @@ -106,6 +109,14 @@ def get_ntypes(self): def get_type_map(self): return self.type_map + def get_numb_fparam(self) -> int: + """Get the number of frame parameters.""" + return self.numb_fparam + + def get_numb_aparam(self) -> int: + """Get the number of atomic parameters.""" + return self.numb_fparam + def data_stat(self, data): all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys=False) m_all_stat = merge_sys_stat(all_stat) @@ -454,3 +465,38 @@ def natoms_not_match(self, force, natoms, atype): ghost_force = tf.concat([ghost_force_real, ghost_force_mag], axis=1) force = tf.concat([loc_force, ghost_force], axis=1) return force + + def change_energy_bias( + self, + data: DeepmdDataSystem, + frozen_model: str, + origin_type_map: list, + full_type_map: str, + bias_shift: str = "delta", + ) -> None: + """Change the energy bias according to the input data and the pretrained model. + + Parameters + ---------- + data : DeepmdDataSystem + The training data. + frozen_model : str + The path file of frozen model. + origin_type_map : list + The original type_map in dataset, they are targets to change the energy bias. + full_type_map : str + The full type_map in pretrained model + bias_shift : str + The mode for changing energy bias : ['delta', 'statistic'] + 'delta' : perform predictions on energies of target dataset, + and do least sqaure on the errors to obtain the target shift as bias. + 'statistic' : directly use the statistic energy bias in the target dataset. + """ + self.fitting.change_energy_bias( + data, + frozen_model, + origin_type_map, + full_type_map, + bias_shift, + self.data_bias_nsample, + ) diff --git a/deepmd/model/model.py b/deepmd/model/model.py index 8e6ffad910..9a1f72368b 100644 --- a/deepmd/model/model.py +++ b/deepmd/model/model.py @@ -11,16 +11,132 @@ Union, ) +from deepmd.descriptor.descriptor import ( + Descriptor, +) from deepmd.env import ( GLOBAL_TF_FLOAT_PRECISION, tf, ) +from deepmd.fit.fitting import ( + Fitting, +) +from deepmd.loss.loss import ( + Loss, +) +from deepmd.utils.argcheck import ( + type_embedding_args, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) from deepmd.utils.graph import ( load_graph_def, ) +from deepmd.utils.pair_tab import ( + PairTab, +) +from deepmd.utils.spin import ( + Spin, +) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) class Model(ABC): + """Abstract base model. + + Parameters + ---------- + type_embedding + Type embedding net + type_map + Mapping atom type to the name (str) of the type. + For example `type_map[1]` gives the name of the type 1. + data_stat_nbatch + Number of frames used for data statistic + data_bias_nsample + The number of training samples in a system to compute and change the energy bias. + data_stat_protect + Protect parameter for atomic energy regression + use_srtab + The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly. + smin_alpha + The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided. + sw_rmin + The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + sw_rmin + The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided. + spin + spin + compress + Compression information for internal use + """ + + def __new__(cls, *args, **kwargs): + if cls is Model: + # init model + # infer model type by fitting_type + from deepmd.model.multi import ( + MultiModel, + ) + + model_type = kwargs.get("type", "standard") + if model_type == "standard": + cls = StandardModel + elif model_type == "multi": + cls = MultiModel + else: + raise ValueError(f"unknown model type: {model_type}") + return cls.__new__(cls, *args, **kwargs) + return super().__new__(cls) + + def __init__( + self, + type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, + type_map: Optional[List[str]] = None, + data_stat_nbatch: int = 10, + data_bias_nsample: int = 10, + data_stat_protect: float = 1e-2, + use_srtab: Optional[str] = None, + smin_alpha: Optional[float] = None, + sw_rmin: Optional[float] = None, + sw_rmax: Optional[float] = None, + spin: Optional[Spin] = None, + compress: Optional[dict] = None, + **kwargs, + ) -> None: + super().__init__() + # spin + if isinstance(spin, Spin): + self.spin = spin + elif spin is not None: + self.spin = Spin(**spin) + else: + self.spin = None + self.compress = compress + # other inputs + if type_map is None: + self.type_map = [] + else: + self.type_map = type_map + self.data_stat_nbatch = data_stat_nbatch + self.data_bias_nsample = data_bias_nsample + self.data_stat_protect = data_stat_protect + self.srtab_name = use_srtab + if self.srtab_name is not None: + self.srtab = PairTab(self.srtab_name) + self.smin_alpha = smin_alpha + self.sw_rmin = sw_rmin + self.sw_rmax = sw_rmax + else: + self.srtab = None + + def get_type_map(self) -> list: + """Get the type map.""" + return self.type_map + @abstractmethod def build( self, @@ -187,3 +303,210 @@ def _import_graph_def_from_ckpt_meta( return tf.import_graph_def( sub_graph_def, input_map=feed_dict, return_elements=return_elements, name="" ) + + def enable_mixed_precision(self, mixed_prec: dict): + """Enable mixed precision for the model. + + Parameters + ---------- + mixed_prec : dict + The mixed precision config + """ + raise RuntimeError("Not supported") + + def change_energy_bias( + self, + data: DeepmdDataSystem, + frozen_model: str, + origin_type_map: list, + full_type_map: str, + bias_shift: str = "delta", + ) -> None: + """Change the energy bias according to the input data and the pretrained model. + + Parameters + ---------- + data : DeepmdDataSystem + The training data. + frozen_model : str + The path file of frozen model. + origin_type_map : list + The original type_map in dataset, they are targets to change the energy bias. + full_type_map : str + The full type_map in pretrained model + bias_shift : str + The mode for changing energy bias : ['delta', 'statistic'] + 'delta' : perform predictions on energies of target dataset, + and do least sqaure on the errors to obtain the target shift as bias. + 'statistic' : directly use the statistic energy bias in the target dataset. + """ + raise RuntimeError("Not supported") + + def enable_compression(self): + """Enable compression.""" + raise RuntimeError("Not supported") + + def get_numb_fparam(self) -> Union[int, dict]: + """Get the number of frame parameters.""" + return 0 + + def get_numb_aparam(self) -> Union[int, dict]: + """Get the number of atomic parameters.""" + return 0 + + def get_numb_dos(self) -> Union[int, dict]: + """Get the number of gridpoints in energy space.""" + return 0 + + @abstractmethod + def get_fitting(self) -> Union[str, dict]: + """Get the fitting(s).""" + + @abstractmethod + def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + """Get the loss function(s).""" + + @abstractmethod + def get_rcut(self) -> float: + """Get cutoff radius of the model.""" + + @abstractmethod + def get_ntypes(self) -> int: + """Get the number of types.""" + + @abstractmethod + def data_stat(self, data: dict): + """Data staticis.""" + + +class StandardModel(Model): + """Standard model, which must contain a descriptor and a fitting. + + Parameters + ---------- + descriptor : Union[dict, Descriptor] + The descriptor + fitting_net : Union[dict, Fitting] + The fitting network + type_embedding : dict, optional + The type embedding + type_map : list of dict, optional + The type map + """ + + def __new__(cls, *args, **kwargs): + from .dos import ( + DOSModel, + ) + from .ener import ( + EnerModel, + ) + from .tensor import ( + DipoleModel, + PolarModel, + ) + + if cls is StandardModel: + fitting_type = kwargs["fitting_net"]["type"] + # init model + # infer model type by fitting_type + if fitting_type == "ener": + cls = EnerModel + elif fitting_type == "dos": + cls = DOSModel + elif fitting_type == "dipole": + cls = DipoleModel + elif fitting_type == "polar": + cls = PolarModel + else: + raise RuntimeError("get unknown fitting type when building model") + return cls.__new__(cls) + return super().__new__(cls) + + def __init__( + self, + descriptor: Union[dict, Descriptor], + fitting_net: Union[dict, Fitting], + type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, + type_map: Optional[List[str]] = None, + **kwargs, + ) -> None: + super().__init__( + descriptor=descriptor, fitting=fitting_net, type_map=type_map, **kwargs + ) + if isinstance(descriptor, Descriptor): + self.descrpt = descriptor + else: + self.descrpt = Descriptor( + **descriptor, ntypes=len(type_map), spin=self.spin + ) + + if isinstance(fitting_net, Fitting): + self.fitting = fitting_net + else: + self.fitting = Fitting(**fitting_net, descrpt=self.descrpt, spin=self.spin) + self.rcut = self.descrpt.get_rcut() + self.ntypes = self.descrpt.get_ntypes() + + # type embedding + if type_embedding is not None and isinstance(type_embedding, TypeEmbedNet): + self.typeebd = type_embedding + elif type_embedding is not None: + self.typeebd = TypeEmbedNet( + **type_embedding, + padding=self.descrpt.explicit_ntypes, + ) + elif self.descrpt.explicit_ntypes: + default_args = type_embedding_args() + default_args_dict = {i.name: i.default for i in default_args} + default_args_dict["activation_function"] = None + self.typeebd = TypeEmbedNet( + **default_args_dict, + padding=True, + ) + else: + self.typeebd = None + + def enable_mixed_precision(self, mixed_prec: dict): + """Enable mixed precision for the model. + + Parameters + ---------- + mixed_prec : dict + The mixed precision config + """ + self.descrpt.enable_mixed_precision(mixed_prec) + self.fitting.enable_mixed_precision(mixed_prec) + + def enable_compression(self): + """Enable compression.""" + graph, graph_def = load_graph_def(self.compress["model_file"]) + self.descrpt.enable_compression( + self.compress["min_nbor_dist"], + graph, + graph_def, + self.compress["table_config"][0], + self.compress["table_config"][1], + self.compress["table_config"][2], + self.compress["table_config"][3], + ) + # for fparam or aparam settings in 'ener' type fitting net + self.fitting.init_variables(graph, graph_def) + if self.typeebd is not None: + self.typeebd.init_variables(graph, graph_def) + + def get_fitting(self) -> Union[Fitting, dict]: + """Get the fitting(s).""" + return self.fitting + + def get_loss(self, loss: dict, lr) -> Union[Loss, dict]: + """Get the loss function(s).""" + return self.fitting.get_loss(loss, lr) + + def get_rcut(self) -> float: + """Get cutoff radius of the model.""" + return self.rcut + + def get_ntypes(self) -> int: + """Get the number of types.""" + return self.ntypes diff --git a/deepmd/model/multi.py b/deepmd/model/multi.py index dcf7566f36..13af572c51 100644 --- a/deepmd/model/multi.py +++ b/deepmd/model/multi.py @@ -1,23 +1,49 @@ import json from typing import ( + Dict, List, Optional, ) import numpy as np +from deepmd.descriptor.descriptor import ( + Descriptor, +) from deepmd.env import ( MODEL_VERSION, global_cvt_2_ener_float, op_module, tf, ) +from deepmd.fit import ( + DipoleFittingSeA, + DOSFitting, + EnerFitting, + GlobalPolarFittingSeA, + PolarFittingSeA, +) +from deepmd.fit.fitting import ( + Fitting, +) +from deepmd.loss.loss import ( + Loss, +) +from deepmd.utils.argcheck import ( + type_embedding_args, +) from deepmd.utils.graph import ( get_tensor_by_name_from_graph, ) from deepmd.utils.pair_tab import ( PairTab, ) +from deepmd.utils.spin import ( + Spin, +) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) from .model import ( Model, @@ -33,13 +59,13 @@ class MultiModel(Model): Parameters ---------- - descrpt + descriptor Descriptor - fitting_dict + fitting_net_dict Dictionary of fitting nets fitting_type_dict - Dictionary of types of fitting nets - typeebd + deprecated argument + type_embedding Type embedding net type_map Mapping atom type to the name (str) of the type. @@ -62,10 +88,10 @@ class MultiModel(Model): def __init__( self, - descrpt, - fitting_dict, - fitting_type_dict, - typeebd=None, + descriptor: dict, + fitting_net_dict: dict, + fitting_type_dict: Optional[dict] = None, # deprecated + type_embedding=None, type_map: Optional[List[str]] = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, @@ -73,22 +99,72 @@ def __init__( smin_alpha: Optional[float] = None, sw_rmin: Optional[float] = None, sw_rmax: Optional[float] = None, + **kwargs, ) -> None: """Constructor.""" + super().__init__( + descriptor=descriptor, + fitting_net_dict=fitting_net_dict, + type_embedding=type_embedding, + type_map=type_map, + data_stat_nbatch=data_stat_nbatch, + data_stat_protect=data_stat_protect, + use_srtab=use_srtab, + smin_alpha=smin_alpha, + sw_rmin=sw_rmin, + sw_rmax=sw_rmax, + ) + if self.spin is not None and not isinstance(self.spin, Spin): + self.spin = Spin(**self.spin) + if isinstance(descriptor, Descriptor): + self.descrpt = descriptor + else: + self.descrpt = Descriptor( + **descriptor, + ntypes=len(type_map), + multi_task=True, + spin=self.spin, + ) + + fitting_dict = {} + for item in fitting_net_dict: + item_fitting_param = fitting_net_dict[item] + if isinstance(item_fitting_param, Fitting): + fitting_dict[item] = item_fitting_param + else: + fitting_dict[item] = Fitting( + **item_fitting_param, descrpt=self.descrpt, spin=self.spin + ) + + # type embedding + if type_embedding is not None and isinstance(type_embedding, TypeEmbedNet): + self.typeebd = type_embedding + elif type_embedding is not None: + self.typeebd = TypeEmbedNet( + **type_embedding, + padding=self.descrpt.explicit_ntypes, + ) + elif self.descrpt.explicit_ntypes: + default_args = type_embedding_args() + default_args_dict = {i.name: i.default for i in default_args} + default_args_dict["activation_function"] = None + self.typeebd = TypeEmbedNet( + **default_args_dict, + padding=True, + ) + else: + self.typeebd = None + # descriptor - self.descrpt = descrpt self.rcut = self.descrpt.get_rcut() self.ntypes = self.descrpt.get_ntypes() # fitting self.fitting_dict = fitting_dict - self.fitting_type_dict = fitting_type_dict self.numb_fparam_dict = { item: self.fitting_dict[item].get_numb_fparam() for item in self.fitting_dict - if self.fitting_type_dict[item] == "ener" + if isinstance(self.fitting_dict[item], EnerFitting) } - # type embedding - self.typeebd = typeebd # other inputs if type_map is None: self.type_map = [] @@ -195,11 +271,10 @@ def build( natomsel = {} nout = {} for fitting_key in self.fitting_dict: - if self.fitting_type_dict[fitting_key] in [ - "dipole", - "polar", - "global_polar", - ]: + if isinstance( + self.fitting_dict[fitting_key], + (DipoleFittingSeA, PolarFittingSeA, GlobalPolarFittingSeA), + ): sel_type[fitting_key] = self.fitting_dict[ fitting_key ].get_sel_type() @@ -299,7 +374,7 @@ def build( self.atom_ener = {} model_dict = {} for fitting_key in self.fitting_dict: - if self.fitting_type_dict[fitting_key] == "ener": + if isinstance(self.fitting_dict[fitting_key], EnerFitting): atom_ener = self.fitting_dict[fitting_key].build( dout, natoms, @@ -384,12 +459,15 @@ def build( model_dict[fitting_key]["atom_virial"] = atom_virial model_dict[fitting_key]["coord"] = coord model_dict[fitting_key]["atype"] = atype - elif self.fitting_type_dict[fitting_key] in [ - "dipole", - "polar", - "global_polar", - ]: - tensor_name = self.fitting_type_dict[fitting_key] + elif isinstance( + self.fitting_dict[fitting_key], + (DipoleFittingSeA, PolarFittingSeA, GlobalPolarFittingSeA), + ): + tensor_name = { + DipoleFittingSeA: "dipole", + PolarFittingSeA: "polar", + GlobalPolarFittingSeA: "global_polar", + }[type(self.fitting_dict[fitting_key])] output = self.fitting_dict[fitting_key].build( dout, rot_mat, @@ -464,7 +542,6 @@ def build( model_dict[fitting_key]["force"] = force model_dict[fitting_key]["virial"] = virial model_dict[fitting_key]["atom_virial"] = atom_virial - return model_dict def init_variables( @@ -506,3 +583,64 @@ def init_variables( tf.constant("original_model", name="model_type", dtype=tf.string) if self.typeebd is not None: self.typeebd.init_variables(graph, graph_def, suffix=suffix) + + def enable_mixed_precision(self, mixed_prec: dict): + """Enable mixed precision for the model. + + Parameters + ---------- + mixed_prec : dict + The mixed precision config + """ + self.descrpt.enable_mixed_precision(mixed_prec) + for fitting_key in self.fitting_dict: + self.fitting_dict[fitting_key].enable_mixed_precision(self.mixed_prec) + + def get_numb_fparam(self) -> dict: + """Get the number of frame parameters.""" + numb_fparam_dict = {} + for fitting_key in self.fitting_dict: + if isinstance(self.fitting_dict[fitting_key], (EnerFitting, DOSFitting)): + numb_fparam_dict[fitting_key] = self.fitting_dict[ + fitting_key + ].get_numb_fparam() + else: + numb_fparam_dict[fitting_key] = 0 + return numb_fparam_dict + + def get_numb_aparam(self) -> dict: + """Get the number of atomic parameters.""" + numb_aparam_dict = {} + for fitting_key in self.fitting_dict: + if isinstance(self.fitting_dict[fitting_key], (EnerFitting, DOSFitting)): + numb_aparam_dict[fitting_key] = self.fitting_dict[ + fitting_key + ].get_numb_aparam() + else: + numb_aparam_dict[fitting_key] = 0 + return numb_aparam_dict + + def get_numb_dos(self) -> dict: + """Get the number of gridpoints in energy space.""" + numb_dos_dict = {} + for fitting_key in self.fitting_dict: + if isinstance(self.fitting_dict[fitting_key], DOSFitting): + numb_dos_dict[fitting_key] = self.fitting_dict[ + fitting_key + ].get_numb_dos() + else: + numb_dos_dict[fitting_key] = 0 + return numb_dos_dict + + def get_fitting(self) -> dict: + """Get the fitting(s).""" + return self.fitting_dict.copy() + + def get_loss(self, loss: dict, lr: dict) -> Dict[str, Loss]: + loss_dict = {} + for fitting_key in self.fitting_dict: + loss_param = loss.get(fitting_key, {}) + loss_dict[fitting_key] = self.fitting_dict[fitting_key].get_loss( + loss_param, lr[fitting_key] + ) + return loss_dict diff --git a/deepmd/model/tensor.py b/deepmd/model/tensor.py index eb32fa88cf..d97d52f454 100644 --- a/deepmd/model/tensor.py +++ b/deepmd/model/tensor.py @@ -1,15 +1,19 @@ from typing import ( List, Optional, + Union, ) from deepmd.env import ( MODEL_VERSION, tf, ) +from deepmd.utils.type_embed import ( + TypeEmbedNet, +) from .model import ( - Model, + StandardModel, ) from .model_stat import ( make_stat_input, @@ -17,18 +21,18 @@ ) -class TensorModel(Model): +class TensorModel(StandardModel): """Tensor model. Parameters ---------- tensor_name Name of the tensor. - descrpt + descriptor Descriptor - fitting + fitting_net Fitting net - typeebd + type_embedding Type embedding net type_map Mapping atom type to the name (str) of the type. @@ -42,30 +46,25 @@ class TensorModel(Model): def __init__( self, tensor_name: str, - descrpt, - fitting, - typeebd=None, + descriptor: dict, + fitting_net: dict, + type_embedding: Optional[Union[dict, TypeEmbedNet]] = None, type_map: Optional[List[str]] = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, + **kwargs, ) -> None: """Constructor.""" + super().__init__( + descriptor=descriptor, + fitting_net=fitting_net, + type_embedding=type_embedding, + type_map=type_map, + data_stat_nbatch=data_stat_nbatch, + data_stat_protect=data_stat_protect, + **kwargs, + ) self.model_type = tensor_name - # descriptor - self.descrpt = descrpt - self.rcut = self.descrpt.get_rcut() - self.ntypes = self.descrpt.get_ntypes() - # fitting - self.fitting = fitting - # type embedding - self.typeebd = typeebd - # other params - if type_map is None: - self.type_map = [] - else: - self.type_map = type_map - self.data_stat_nbatch = data_stat_nbatch - self.data_stat_protect = data_stat_protect def get_rcut(self): return self.rcut diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py index 580d434533..a191d267aa 100644 --- a/deepmd/train/trainer.py +++ b/deepmd/train/trainer.py @@ -26,9 +26,6 @@ get_precision, j_must_have, ) -from deepmd.descriptor.descriptor import ( - Descriptor, -) from deepmd.env import ( GLOBAL_ENER_FLOAT_PRECISION, GLOBAL_TF_FLOAT_PRECISION, @@ -37,27 +34,16 @@ tf, tfv2, ) -from deepmd.fit import ( - Fitting, -) -from deepmd.loss import ( - DOSLoss, - EnerDipoleLoss, - EnerSpinLoss, - EnerStdLoss, - TensorLoss, +from deepmd.fit.ener import ( + EnerFitting, ) from deepmd.model import ( - DipoleModel, - DOSModel, - EnerModel, MultiModel, - PolarModel, ) -from deepmd.utils import random as dp_random -from deepmd.utils.argcheck import ( - type_embedding_args, +from deepmd.model.model import ( + Model, ) +from deepmd.utils import random as dp_random from deepmd.utils.data_system import ( DeepmdDataSystem, ) @@ -75,12 +61,6 @@ from deepmd.utils.sess import ( run_sess, ) -from deepmd.utils.spin import ( - Spin, -) -from deepmd.utils.type_embed import ( - TypeEmbedNet, -) log = logging.getLogger(__name__) @@ -108,27 +88,8 @@ def __init__(self, jdata, run_opt, is_compress=False): def _init_param(self, jdata): # model config model_param = j_must_have(jdata, "model") - self.multi_task_mode = "fitting_net_dict" in model_param - descrpt_param = j_must_have(model_param, "descriptor") - fitting_param = ( - j_must_have(model_param, "fitting_net") - if not self.multi_task_mode - else j_must_have(model_param, "fitting_net_dict") - ) - typeebd_param = model_param.get("type_embedding", None) - spin_param = model_param.get("spin", None) - self.model_param = model_param - self.descrpt_param = descrpt_param - - # spin - if spin_param is not None: - self.spin = Spin( - use_spin=spin_param["use_spin"], - virtual_len=spin_param["virtual_len"], - spin_norm=spin_param["spin_norm"], - ) - else: - self.spin = None + if "fitting_key" in model_param: + model_param["type"] = "multi" # nvnmd self.nvnmd_param = jdata.get("nvnmd", {}) @@ -138,149 +99,10 @@ def _init_param(self, jdata): nvnmd_cfg.disp_message() nvnmd_cfg.save() - # descriptor - try: - descrpt_type = descrpt_param["type"] - self.descrpt_type = descrpt_type - except KeyError: - raise KeyError("the type of descriptor should be set by `type`") - - explicit_ntypes_descrpt = ["se_atten"] - hybrid_with_tebd = False - if descrpt_param["type"] in explicit_ntypes_descrpt: - descrpt_param["ntypes"] = len(model_param["type_map"]) - elif descrpt_param["type"] == "hybrid": - for descrpt_item in descrpt_param["list"]: - if descrpt_item["type"] in explicit_ntypes_descrpt: - descrpt_item["ntypes"] = len(model_param["type_map"]) - hybrid_with_tebd = True - if self.multi_task_mode: - descrpt_param["multi_task"] = True - if descrpt_param["type"] in ["se_e2_a", "se_a", "se_e2_r", "se_r", "hybrid"]: - descrpt_param["spin"] = self.spin - self.descrpt = Descriptor(**descrpt_param) - - # fitting net - if not self.multi_task_mode: - fitting_type = fitting_param.get("type", "ener") - self.fitting_type = fitting_type - fitting_param["descrpt"] = self.descrpt - if fitting_type == "ener": - fitting_param["spin"] = self.spin - self.fitting = Fitting(**fitting_param) - else: - self.fitting_dict = {} - self.fitting_type_dict = {} - self.nfitting = len(fitting_param) - for item in fitting_param: - item_fitting_param = fitting_param[item] - item_fitting_type = item_fitting_param.get("type", "ener") - self.fitting_type_dict[item] = item_fitting_type - item_fitting_param["descrpt"] = self.descrpt - if item_fitting_type == "ener": - item_fitting_param["spin"] = self.spin - self.fitting_dict[item] = Fitting(**item_fitting_param) - - # type embedding - padding = False - if descrpt_type == "se_atten" or hybrid_with_tebd: - padding = True - if typeebd_param is not None: - self.typeebd = TypeEmbedNet( - neuron=typeebd_param["neuron"], - resnet_dt=typeebd_param["resnet_dt"], - activation_function=typeebd_param["activation_function"], - precision=typeebd_param["precision"], - trainable=typeebd_param["trainable"], - seed=typeebd_param["seed"], - padding=padding, - ) - elif descrpt_type == "se_atten" or hybrid_with_tebd: - default_args = type_embedding_args() - default_args_dict = {i.name: i.default for i in default_args} - self.typeebd = TypeEmbedNet( - neuron=default_args_dict["neuron"], - resnet_dt=default_args_dict["resnet_dt"], - activation_function=None, - precision=default_args_dict["precision"], - trainable=default_args_dict["trainable"], - seed=default_args_dict["seed"], - padding=padding, - ) - else: - self.typeebd = None - # init model - # infer model type by fitting_type - if not self.multi_task_mode: - if self.fitting_type == "ener": - self.model = EnerModel( - self.descrpt, - self.fitting, - self.typeebd, - model_param.get("type_map"), - model_param.get("data_stat_nbatch", 10), - model_param.get("data_stat_protect", 1e-2), - model_param.get("use_srtab"), - model_param.get("smin_alpha"), - model_param.get("sw_rmin"), - model_param.get("sw_rmax"), - self.spin, - ) - # elif fitting_type == 'wfc': - # self.model = WFCModel(model_param, self.descrpt, self.fitting) - elif self.fitting_type == "dos": - self.model = DOSModel( - self.descrpt, - self.fitting, - self.typeebd, - model_param.get("type_map"), - model_param.get("data_stat_nbatch", 10), - model_param.get("data_stat_protect", 1e-2), - ) - - elif self.fitting_type == "dipole": - self.model = DipoleModel( - self.descrpt, - self.fitting, - self.typeebd, - model_param.get("type_map"), - model_param.get("data_stat_nbatch", 10), - model_param.get("data_stat_protect", 1e-2), - ) - elif self.fitting_type == "polar": - self.model = PolarModel( - self.descrpt, - self.fitting, - self.typeebd, - model_param.get("type_map"), - model_param.get("data_stat_nbatch", 10), - model_param.get("data_stat_protect", 1e-2), - ) - # elif self.fitting_type == 'global_polar': - # self.model = GlobalPolarModel( - # self.descrpt, - # self.fitting, - # model_param.get('type_map'), - # model_param.get('data_stat_nbatch', 10), - # model_param.get('data_stat_protect', 1e-2) - # ) - else: - raise RuntimeError("get unknown fitting type when building model") - else: # multi-task mode - self.model = MultiModel( - self.descrpt, - self.fitting_dict, - self.fitting_type_dict, - self.typeebd, - model_param.get("type_map"), - model_param.get("data_stat_nbatch", 10), - model_param.get("data_stat_protect", 1e-2), - model_param.get("use_srtab"), - model_param.get("smin_alpha"), - model_param.get("sw_rmin"), - model_param.get("sw_rmax"), - ) + self.model = Model(**model_param) + self.multi_task_mode = isinstance(self.model, MultiModel) + self.fitting = self.model.get_fitting() def get_lr_and_coef(lr_param): scale_by_worker = lr_param.get("scale_by_worker", "linear") @@ -307,86 +129,20 @@ def get_lr_and_coef(lr_param): self.lr_dict = {} self.scale_lr_coef_dict = {} lr_param_dict = jdata.get("learning_rate_dict", {}) - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: lr_param = lr_param_dict.get(fitting_key, {}) ( self.lr_dict[fitting_key], self.scale_lr_coef_dict[fitting_key], ) = get_lr_and_coef(lr_param) - # loss # infer loss type by fitting_type - def loss_init(_loss_param, _fitting_type, _fitting, _lr): - _loss_type = _loss_param.get("type", "ener") - if _fitting_type == "ener": - _loss_param.pop("type", None) - _loss_param["starter_learning_rate"] = _lr.start_lr() - if _loss_type == "ener": - loss = EnerStdLoss(**_loss_param) - elif _loss_type == "ener_dipole": - loss = EnerDipoleLoss(**_loss_param) - elif _loss_type == "ener_spin": - loss = EnerSpinLoss(**_loss_param, use_spin=self.spin.use_spin) - else: - raise RuntimeError("unknown loss type") - elif _fitting_type == "dos": - _loss_param.pop("type", None) - _loss_param["starter_learning_rate"] = _lr.start_lr() - _loss_param["numb_dos"] = self.fitting.get_numb_dos() - loss = DOSLoss(**_loss_param) - elif _fitting_type == "wfc": - loss = TensorLoss( - _loss_param, - model=_fitting, - tensor_name="wfc", - tensor_size=_fitting.get_out_size(), - label_name="wfc", - ) - elif _fitting_type == "dipole": - loss = TensorLoss( - _loss_param, - model=_fitting, - tensor_name="dipole", - tensor_size=3, - label_name="dipole", - ) - elif _fitting_type == "polar": - loss = TensorLoss( - _loss_param, - model=_fitting, - tensor_name="polar", - tensor_size=9, - label_name="polarizability", - ) - elif _fitting_type == "global_polar": - loss = TensorLoss( - _loss_param, - model=_fitting, - tensor_name="global_polar", - tensor_size=9, - atomic=False, - label_name="polarizability", - ) - else: - raise RuntimeError( - "get unknown fitting type when building loss function" - ) - return loss - if not self.multi_task_mode: loss_param = jdata.get("loss", {}) - self.loss = loss_init(loss_param, self.fitting_type, self.fitting, self.lr) + self.loss = self.model.get_loss(loss_param, self.lr) else: - self.loss_dict = {} - loss_param_dict = jdata.get("loss_dict", {}) - for fitting_key in self.fitting_type_dict: - loss_param = loss_param_dict.get(fitting_key, {}) - self.loss_dict[fitting_key] = loss_init( - loss_param, - self.fitting_type_dict[fitting_key], - self.fitting_dict[fitting_key], - self.lr_dict[fitting_key], - ) + loss_param = jdata.get("loss_dict", {}) + self.loss_dict = self.model.get_loss(loss_param, self.lr_dict) # training tr_data = jdata["training"] @@ -394,7 +150,7 @@ def loss_init(_loss_param, _fitting_type, _fitting, _lr): if self.multi_task_mode: self.fitting_key_list = [] self.fitting_prob = [] - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: self.fitting_key_list.append(fitting_key) # multi-task mode must have self.fitting_weight self.fitting_prob.append(self.fitting_weight[fitting_key]) @@ -426,30 +182,15 @@ def loss_init(_loss_param, _fitting_type, _fitting, _lr): # self.auto_prob_style = tr_data['auto_prob'] self.useBN = False if not self.multi_task_mode: - if ( - self.fitting_type == "ener" or self.fitting_type == "dos" - ) and self.fitting.get_numb_fparam() > 0: - self.numb_fparam = self.fitting.get_numb_fparam() - else: - self.numb_fparam = 0 + self.numb_fparam = self.model.get_numb_fparam() if tr_data.get("validation_data", None) is not None: self.valid_numb_batch = tr_data["validation_data"].get("numb_btch", 1) else: self.valid_numb_batch = 1 else: - self.numb_fparam_dict = {} + self.numb_fparam_dict = self.model.get_numb_fparam() self.valid_numb_batch_dict = {} - for fitting_key in self.fitting_type_dict: - if ( - self.fitting_type_dict[fitting_key] == "ener" - and self.fitting_dict[fitting_key].get_numb_fparam() > 0 - ): - self.numb_fparam_dict[fitting_key] = self.fitting_dict[ - fitting_key - ].get_numb_fparam() - else: - self.numb_fparam_dict[fitting_key] = 0 data_dict = tr_data.get("data_dict", None) for systems in data_dict: if data_dict[systems].get("validation_data", None) is not None: @@ -470,12 +211,9 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): if not self.multi_task_mode: if not self.is_compress and data.mixed_type: - assert self.descrpt_type in [ - "se_atten" - ], "Data in mixed_type format must use attention descriptor!" - assert self.fitting_type in [ - "ener" - ], "Data in mixed_type format must use ener fitting!" + assert isinstance( + self.fitting, EnerFitting + ), "Data in mixed_type format must use ener fitting!" if self.numb_fparam > 0: log.info("training with %d frame parameter(s)" % self.numb_fparam) @@ -489,16 +227,11 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): for fitting_key in data: self.valid_fitting_key.append(fitting_key) if data[fitting_key].mixed_type: - assert self.descrpt_type in ["se_atten"], ( - "Data for fitting net {} in mixed_type format " - "must use attention descriptor!".format(fitting_key) - ) - assert self.fitting_type_dict[fitting_key] in [ - "ener" - ], "Data for fitting net {} in mixed_type format must use ener fitting!".format( + assert isinstance( + self.fitting[fitting_key], EnerFitting + ), "Data for fitting net {} in mixed_type format must use ener fitting!".format( fitting_key ) - if self.numb_fparam_dict[fitting_key] > 0: log.info( "fitting net %s training with %d frame parameter(s)" @@ -561,20 +294,7 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): # TODO: this is a simple fix but we should have a clear # architecture to call neighbor stat else: - graph, graph_def = load_graph_def( - self.model_param["compress"]["model_file"] - ) - self.descrpt.enable_compression( - self.model_param["compress"]["min_nbor_dist"], - graph, - graph_def, - self.model_param["compress"]["table_config"][0], - self.model_param["compress"]["table_config"][1], - self.model_param["compress"]["table_config"][2], - self.model_param["compress"]["table_config"][3], - ) - # for fparam or aparam settings in 'ener' type fitting net - self.fitting.init_variables(graph, graph_def) + self.model.enable_compression() if self.is_compress or self.model_type == "compressed_model": tf.constant("compressed_model", name="model_type", dtype=tf.string) @@ -582,14 +302,7 @@ def build(self, data=None, stop_batch=0, origin_type_map=None, suffix=""): tf.constant("original_model", name="model_type", dtype=tf.string) if self.mixed_prec is not None: - self.descrpt.enable_mixed_precision(self.mixed_prec) - if not self.multi_task_mode: - self.fitting.enable_mixed_precision(self.mixed_prec) - else: - for fitting_key in self.fitting_dict: - self.fitting_dict[fitting_key].enable_mixed_precision( - self.mixed_prec - ) + self.model.enable_mixed_precision(self.mixed_prec) self._build_lr() self._build_network(data, suffix) @@ -602,7 +315,8 @@ def _build_lr(self): self.learning_rate = self.lr.build(self.global_step, self.stop_batch) else: self.learning_rate_dict = {} - for fitting_key in self.fitting_type_dict: + + for fitting_key in self.fitting: self.learning_rate_dict[fitting_key] = self.lr_dict[fitting_key].build( self.global_step, self.stop_batch ) @@ -623,7 +337,7 @@ def _build_loss(self): l2_l = tf.cast(l2_l, get_precision(self.mixed_prec["output_prec"])) else: l2_l, l2_more = {}, {} - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: lr = self.learning_rate_dict[fitting_key] model = self.model_pred[fitting_key] loss_dict = self.loss_dict[fitting_key] @@ -748,7 +462,7 @@ def _build_training(self): self.train_op = tf.group(*train_ops) else: self.train_op = {} - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: optimizer = self._build_optimizer(fitting_key=fitting_key) apply_op = optimizer.minimize( loss=self.l2_l[fitting_key], @@ -841,7 +555,7 @@ def train(self, train_data=None, valid_data=None): ) ) else: - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: log.info( "%s: start training at lr %.2e (== %.2e), decay_step %d, decay_rate %f, final lr will be %.2e" % ( @@ -903,7 +617,7 @@ def train(self, train_data=None, valid_data=None): else: datasetloader = {} data_op = {} - for fitting_key in self.fitting_type_dict: + for fitting_key in self.fitting: datasetloader[fitting_key] = DatasetLoader(train_data[fitting_key]) data_op[fitting_key] = datasetloader[fitting_key].build() @@ -915,7 +629,8 @@ def train(self, train_data=None, valid_data=None): batch_train_op = self.train_op else: fitting_idx = dp_random.choice( - np.arange(self.nfitting), p=np.array(self.fitting_prob) + np.arange(len(self.fitting_key_list)), + p=np.array(self.fitting_prob), ) fitting_key = self.fitting_key_list[fitting_idx] train_batch = train_data[fitting_key].get_batch() @@ -931,7 +646,7 @@ def train(self, train_data=None, valid_data=None): next_train_batch_op = data_op else: fitting_idx = dp_random.choice( - np.arange(self.nfitting), p=np.array(self.fitting_prob) + np.arange(len(self.fitting_key_list)), p=np.array(self.fitting_prob) ) next_fitting_key = self.fitting_key_list[fitting_idx] next_datasetloader = datasetloader[next_fitting_key] @@ -1418,16 +1133,12 @@ def _change_energy_bias( self, data, frozen_model, origin_type_map, bias_shift="delta" ): full_type_map = data.get_type_map() - assert ( - self.fitting_type == "ener" - ), "energy bias changing only supports 'ener' fitting net!" - self.model.fitting.change_energy_bias( + self.model.change_energy_bias( data, frozen_model, origin_type_map, full_type_map, bias_shift=bias_shift, - ntest=self.model_param.get("data_bias_nsample", 10), ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 1fe00b4f43..5748ad9bf7 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -744,9 +744,6 @@ def model_args(): doc_data_stat_protect = "Protect parameter for atomic energy regression." doc_data_bias_nsample = "The number of training samples in a system to compute and change the energy bias." doc_type_embedding = "The type embedding." - doc_descrpt = "The descriptor of atomic environment." - doc_fitting = "The fitting of physical properties." - doc_fitting_net_dict = "The dictionary of multiple fitting nets in multi-task mode. Each fitting_net_dict[fitting_key] is the single definition of fitting of physical properties with user-defined name `fitting_key`." doc_modifier = "The modifier of model output." doc_use_srtab = "The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly." doc_smin_alpha = "The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided." @@ -754,8 +751,7 @@ def model_args(): doc_sw_rmax = "The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided." doc_compress_config = "Model compression configurations" doc_spin = "The settings for systems with spin." - - ca = Argument( + return Argument( "model", dict, [ @@ -793,18 +789,6 @@ def model_args(): optional=True, doc=doc_type_embedding, ), - Argument( - "descriptor", dict, [], [descrpt_variant_type_args()], doc=doc_descrpt - ), - Argument( - "fitting_net", - dict, - [], - [fitting_variant_type_args()], - optional=True, - doc=doc_fitting, - ), - Argument("fitting_net_dict", dict, optional=True, doc=doc_fitting_net_dict), Argument( "modifier", dict, @@ -820,11 +804,68 @@ def model_args(): [model_compression_type_args()], optional=True, doc=doc_compress_config, + fold_subdoc=True, ), Argument("spin", dict, spin_args(), [], optional=True, doc=doc_spin), ], + [ + Variant( + "type", + [ + standard_model_args(), + multi_model_args(), + ], + optional=True, + default_tag="standard", + ), + ], + ) + + +def standard_model_args() -> Argument: + doc_descrpt = "The descriptor of atomic environment." + doc_fitting = "The fitting of physical properties." + + ca = Argument( + "standard", + dict, + [ + Argument( + "descriptor", dict, [], [descrpt_variant_type_args()], doc=doc_descrpt + ), + Argument( + "fitting_net", + dict, + [], + [fitting_variant_type_args()], + doc=doc_fitting, + ), + ], + doc="Stardard model, which contains a descriptor and a fitting.", + ) + return ca + + +def multi_model_args() -> Argument: + doc_descrpt = "The descriptor of atomic environment. See model[standard]/descriptor for details." + doc_fitting_net_dict = "The dictionary of multiple fitting nets in multi-task mode. Each fitting_net_dict[fitting_key] is the single definition of fitting of physical properties with user-defined name `fitting_key`." + + ca = Argument( + "multi", + dict, + [ + Argument( + "descriptor", + dict, + [], + [descrpt_variant_type_args()], + doc=doc_descrpt, + fold_subdoc=True, + ), + Argument("fitting_net_dict", dict, doc=doc_fitting_net_dict), + ], + doc="Multiple-task model.", ) - # print(ca.gen_doc()) return ca @@ -1574,6 +1615,7 @@ def normalize_multi_task(data): assert ( "type_map" in data["model"] ), "In multi-task mode, 'model/type_map' must be defined! " + data["model"]["type"] = "multi" data["model"]["fitting_net_dict"] = normalize_fitting_net_dict( data["model"]["fitting_net_dict"] ) diff --git a/deepmd/utils/type_embed.py b/deepmd/utils/type_embed.py index 7a3e0925b8..a2c747c248 100644 --- a/deepmd/utils/type_embed.py +++ b/deepmd/utils/type_embed.py @@ -93,6 +93,7 @@ def __init__( seed: Optional[int] = None, uniform_seed: bool = False, padding: bool = False, + **kwargs, ) -> None: """Constructor.""" self.neuron = neuron diff --git a/source/tests/.gitignore b/source/tests/.gitignore new file mode 100644 index 0000000000..ad1f2dffce --- /dev/null +++ b/source/tests/.gitignore @@ -0,0 +1,13 @@ +# all of the follwing are generated by tests +/e.out +/f.out +/frozen_model.pbtxt +/model.ckpt-0.data-00000-of-00001 +/model.ckpt-0.index +/model.ckpt-1.data-00000-of-00001 +/model.ckpt-1.index +/v.out +/nvnmd/out/map.npy +/nvnmd/out/train_cnn.json +/nvnmd/out/train_qnn.json +/nvnmd/out/weight.npy diff --git a/source/tests/test_init_frz_model_spin.py b/source/tests/test_init_frz_model_spin.py index 18fcb71773..a8223acd93 100644 --- a/source/tests/test_init_frz_model_spin.py +++ b/source/tests/test_init_frz_model_spin.py @@ -52,6 +52,9 @@ def _init_models(): jdata = j_loader("test_model_spin.json") jdata["training"]["save_ckpt"] = ckpt + jdata["training"]["training_data"]["systems"] = [str(tests_path / "model_spin/")] + jdata["training"]["validation_data"]["systems"] = [str(tests_path / "model_spin/")] + del jdata["training"]["set_prefix"] with open(INPUT, "w") as fp: json.dump(jdata, fp, indent=4) @@ -65,6 +68,8 @@ def _init_models(): jdata = j_loader("test_model_spin.json") jdata["training"]["save_ckpt"] = ckpt + jdata["training"]["training_data"]["systems"] = [str(tests_path / "model_spin/")] + jdata["training"]["validation_data"]["systems"] = [str(tests_path / "model_spin/")] del jdata["training"]["set_prefix"] jdata["loss"]["type"] = "ener_spin" jdata = update_deepmd_input(jdata, warning=True, dump="input_v2_compat.json") diff --git a/source/tests/test_model_spin.py b/source/tests/test_model_spin.py index 0bd2856552..06fa6602f2 100644 --- a/source/tests/test_model_spin.py +++ b/source/tests/test_model_spin.py @@ -6,6 +6,7 @@ del_data, gen_data, j_loader, + tests_path, ) from deepmd.common import ( @@ -51,6 +52,12 @@ def test_model_spin(self): test_size = j_must_have(jdata["training"]["validation_data"], "numb_btch") stop_batch = j_must_have(jdata["training"], "numb_steps") rcut = j_must_have(jdata["model"]["descriptor"], "rcut") + jdata["training"]["training_data"]["systems"] = [ + str(tests_path / "model_spin/") + ] + jdata["training"]["validation_data"]["systems"] = [ + str(tests_path / "model_spin/") + ] data = DataSystem(systems, set_pfx, batch_size, test_size, rcut, run_opt=None) test_data = data.get_test()