From 99369d9b2c054b216375a0739827b6c1d5d172ff Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Wed, 27 Nov 2024 09:01:47 +0800 Subject: [PATCH] rename to case embd --- .../dpmodel/atomic_model/dp_atomic_model.py | 6 +-- .../atomic_model/linear_atomic_model.py | 12 ++--- .../atomic_model/make_base_atomic_model.py | 4 +- .../atomic_model/pairtab_atomic_model.py | 4 +- deepmd/dpmodel/fitting/dipole_fitting.py | 4 +- deepmd/dpmodel/fitting/dos_fitting.py | 4 +- deepmd/dpmodel/fitting/ener_fitting.py | 4 +- deepmd/dpmodel/fitting/general_fitting.py | 40 +++++++-------- deepmd/dpmodel/fitting/invar_fitting.py | 4 +- .../dpmodel/fitting/polarizability_fitting.py | 4 +- deepmd/dpmodel/fitting/property_fitting.py | 4 +- deepmd/dpmodel/model/make_model.py | 4 +- .../pt/model/atomic_model/dp_atomic_model.py | 6 +-- .../model/atomic_model/linear_atomic_model.py | 12 ++--- .../atomic_model/pairtab_atomic_model.py | 4 +- deepmd/pt/model/model/make_model.py | 4 +- deepmd/pt/model/task/dipole.py | 8 +-- deepmd/pt/model/task/dos.py | 4 +- deepmd/pt/model/task/ener.py | 4 +- deepmd/pt/model/task/fitting.py | 50 +++++++++---------- deepmd/pt/model/task/invar_fitting.py | 8 +-- deepmd/pt/model/task/polarizability.py | 8 +-- deepmd/pt/model/task/property.py | 8 +-- deepmd/pt/train/training.py | 26 +++++----- deepmd/tf/descriptor/se_a_ebd.py | 2 - deepmd/tf/fit/dipole.py | 14 +++--- deepmd/tf/fit/dos.py | 16 +++--- deepmd/tf/fit/ener.py | 16 +++--- deepmd/tf/fit/polar.py | 14 +++--- deepmd/utils/argcheck.py | 30 +++++------ doc/train/multi-task-training.md | 4 +- .../pytorch_example/input_torch_sharefit.json | 2 +- .../pt/model/water/multitask_sharefit.json | 2 +- source/tests/pt/test_multitask.py | 2 +- .../universal/dpmodel/fitting/test_fitting.py | 20 ++++---- 35 files changed, 179 insertions(+), 179 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/dp_atomic_model.py b/deepmd/dpmodel/atomic_model/dp_atomic_model.py index c93de33f62..2fa072cc78 100644 --- a/deepmd/dpmodel/atomic_model/dp_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/dp_atomic_model.py @@ -65,12 +65,12 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.descriptor.get_sel() - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ - self.fitting.set_caseid(case_idx) + self.fitting.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 0d2d32470b..8108292bd2 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -134,13 +134,13 @@ def get_model_rcuts(self) -> list[float]: def get_sel(self) -> list[int]: return [max([model.get_nsel() for model in self.models])] - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ for model in self.models: - model.set_caseid(case_idx) + model.set_case_embd(case_idx) def get_model_nsels(self) -> list[int]: """Get the processed sels for each individual models. Not distinguishing types.""" @@ -436,13 +436,13 @@ def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": data.pop("type", None) return super().deserialize(data) - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ # only set case_idx for dpmodel - self.models[0].set_caseid(case_idx) + self.models[0].set_case_embd(case_idx) def _compute_weight( self, diff --git a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py index f9b3c98d78..01caa7cd64 100644 --- a/deepmd/dpmodel/atomic_model/make_base_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/make_base_atomic_model.py @@ -69,9 +69,9 @@ def get_sel(self) -> list[int]: pass @abstractmethod - def set_caseid(self, case_idx: int) -> None: + def set_case_embd(self, case_idx: int) -> None: """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ pass diff --git a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py index 8e96579fe3..ecd51452d1 100644 --- a/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/pairtab_atomic_model.py @@ -120,9 +120,9 @@ def get_type_map(self) -> list[str]: def get_sel(self) -> list[int]: return [self.sel] - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ raise NotImplementedError( diff --git a/deepmd/dpmodel/fitting/dipole_fitting.py b/deepmd/dpmodel/fitting/dipole_fitting.py index 42c8e5929e..fcaea43338 100644 --- a/deepmd/dpmodel/fitting/dipole_fitting.py +++ b/deepmd/dpmodel/fitting/dipole_fitting.py @@ -95,7 +95,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[list[bool]] = None, @@ -131,7 +131,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, diff --git a/deepmd/dpmodel/fitting/dos_fitting.py b/deepmd/dpmodel/fitting/dos_fitting.py index 6fa3de2b24..2f6df77eac 100644 --- a/deepmd/dpmodel/fitting/dos_fitting.py +++ b/deepmd/dpmodel/fitting/dos_fitting.py @@ -36,7 +36,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, bias_dos: Optional[np.ndarray] = None, rcond: Optional[float] = None, trainable: Union[bool, list[bool]] = True, @@ -61,7 +61,7 @@ def __init__( bias_atom=bias_dos, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, trainable=trainable, activation_function=activation_function, diff --git a/deepmd/dpmodel/fitting/ener_fitting.py b/deepmd/dpmodel/fitting/ener_fitting.py index 235427954d..6435b6468f 100644 --- a/deepmd/dpmodel/fitting/ener_fitting.py +++ b/deepmd/dpmodel/fitting/ener_fitting.py @@ -32,7 +32,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[list[bool]] = None, @@ -56,7 +56,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, diff --git a/deepmd/dpmodel/fitting/general_fitting.py b/deepmd/dpmodel/fitting/general_fitting.py index 2757276c3a..c05d84c4a1 100644 --- a/deepmd/dpmodel/fitting/general_fitting.py +++ b/deepmd/dpmodel/fitting/general_fitting.py @@ -105,7 +105,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, bias_atom_e: Optional[np.ndarray] = None, rcond: Optional[float] = None, tot_ener_zero: bool = False, @@ -128,7 +128,7 @@ def __init__( self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid + self.dim_case_embd = dim_case_embd self.rcond = rcond self.tot_ener_zero = tot_ener_zero self.trainable = trainable @@ -173,16 +173,16 @@ def __init__( self.aparam_inv_std = np.ones(self.numb_aparam, dtype=self.prec) else: self.aparam_avg, self.aparam_inv_std = None, None - if self.numb_caseid > 0: - self.caseid = np.zeros(self.numb_caseid, dtype=self.prec) + if self.dim_case_embd > 0: + self.case_embd = np.zeros(self.dim_case_embd, dtype=self.prec) else: - self.caseid = None + self.case_embd = None # init networks in_dim = ( self.dim_descrpt + self.numb_fparam + (0 if self.use_aparam_as_mask else self.numb_aparam) - + self.numb_caseid + + self.dim_case_embd ) self.nets = NetworkCollection( 1 if not self.mixed_types else 0, @@ -229,12 +229,12 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this fitting net by the given case_idx, + Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ - self.caseid = np.eye(self.numb_caseid, dtype=self.prec)[case_idx] + self.case_embd = np.eye(self.dim_case_embd, dtype=self.prec)[case_idx] def change_type_map( self, type_map: list[str], model_with_new_type_stat=None @@ -269,8 +269,8 @@ def __setitem__(self, key, value) -> None: self.aparam_avg = value elif key in ["aparam_inv_std"]: self.aparam_inv_std = value - elif key in ["caseid"]: - self.caseid = value + elif key in ["case_embd"]: + self.case_embd = value elif key in ["scale"]: self.scale = value else: @@ -287,8 +287,8 @@ def __getitem__(self, key): return self.aparam_avg elif key in ["aparam_inv_std"]: return self.aparam_inv_std - elif key in ["caseid"]: - return self.caseid + elif key in ["case_embd"]: + return self.case_embd elif key in ["scale"]: return self.scale else: @@ -313,7 +313,7 @@ def serialize(self) -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "rcond": self.rcond, "activation_function": self.activation_function, "precision": self.precision, @@ -322,7 +322,7 @@ def serialize(self) -> dict: "nets": self.nets.serialize(), "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), - "caseid": to_numpy_array(self.caseid), + "case_embd": to_numpy_array(self.case_embd), "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), "aparam_avg": to_numpy_array(self.aparam_avg), @@ -443,16 +443,16 @@ def _call_common( axis=-1, ) - if self.numb_caseid > 0: - assert self.caseid is not None - caseid = xp.tile(xp.reshape(self.caseid, [1, 1, -1]), [nf, nloc, 1]) + if self.dim_case_embd > 0: + assert self.case_embd is not None + case_embd = xp.tile(xp.reshape(self.case_embd, [1, 1, -1]), [nf, nloc, 1]) xx = xp.concat( - [xx, caseid], + [xx, case_embd], axis=-1, ) if xx_zeros is not None: xx_zeros = xp.concat( - [xx_zeros, caseid], + [xx_zeros, case_embd], axis=-1, ) diff --git a/deepmd/dpmodel/fitting/invar_fitting.py b/deepmd/dpmodel/fitting/invar_fitting.py index b1d2630788..b5d3a02d86 100644 --- a/deepmd/dpmodel/fitting/invar_fitting.py +++ b/deepmd/dpmodel/fitting/invar_fitting.py @@ -123,7 +123,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, bias_atom: Optional[np.ndarray] = None, rcond: Optional[float] = None, tot_ener_zero: bool = False, @@ -156,7 +156,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, bias_atom_e=bias_atom, tot_ener_zero=tot_ener_zero, diff --git a/deepmd/dpmodel/fitting/polarizability_fitting.py b/deepmd/dpmodel/fitting/polarizability_fitting.py index c75884dd56..0db6a23377 100644 --- a/deepmd/dpmodel/fitting/polarizability_fitting.py +++ b/deepmd/dpmodel/fitting/polarizability_fitting.py @@ -101,7 +101,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[list[bool]] = None, @@ -151,7 +151,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, tot_ener_zero=tot_ener_zero, trainable=trainable, diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index af4ba18420..8b903af00e 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -78,7 +78,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -100,7 +100,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, rcond=rcond, trainable=trainable, activation_function=activation_function, diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index 63bdec9a93..783e012f98 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -552,8 +552,8 @@ def serialize(self) -> dict: def deserialize(cls, data) -> "CM": return cls(atomic_model_=T_AtomicModel.deserialize(data)) - def set_caseid(self, case_idx: int): - self.atomic_model.set_caseid(case_idx) + def set_case_embd(self, case_idx: int): + self.atomic_model.set_case_embd(case_idx) def get_dim_fparam(self) -> int: """Get the number (dimension) of frame parameters of this atomic model.""" diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index a8dece8236..c988d63213 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -93,12 +93,12 @@ def get_sel(self) -> list[int]: """Get the neighbor selection.""" return self.sel - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ - self.fitting_net.set_caseid(case_idx) + self.fitting_net.set_case_embd(case_idx) def mixed_types(self) -> bool: """If true, the model diff --git a/deepmd/pt/model/atomic_model/linear_atomic_model.py b/deepmd/pt/model/atomic_model/linear_atomic_model.py index 38ed3f2713..36c636ddfb 100644 --- a/deepmd/pt/model/atomic_model/linear_atomic_model.py +++ b/deepmd/pt/model/atomic_model/linear_atomic_model.py @@ -158,13 +158,13 @@ def get_model_rcuts(self) -> list[float]: def get_sel(self) -> list[int]: return [max([model.get_nsel() for model in self.models])] - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ for model in self.models: - model.set_caseid(case_idx) + model.set_case_embd(case_idx) def get_model_nsels(self) -> list[int]: """Get the processed sels for each individual models. Not distinguishing types.""" @@ -569,13 +569,13 @@ def serialize(self) -> dict: ) return dd - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ # only set case_idx for dpmodel - self.models[0].set_caseid(case_idx) + self.models[0].set_case_embd(case_idx) @classmethod def deserialize(cls, data) -> "DPZBLLinearEnergyAtomicModel": diff --git a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py index 1a04d370a0..cb5e48400b 100644 --- a/deepmd/pt/model/atomic_model/pairtab_atomic_model.py +++ b/deepmd/pt/model/atomic_model/pairtab_atomic_model.py @@ -141,9 +141,9 @@ def get_type_map(self) -> list[str]: def get_sel(self) -> list[int]: return [self.sel] - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this atomic model by the given case_idx, + Set the case embedding of this atomic model by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ raise NotImplementedError( diff --git a/deepmd/pt/model/model/make_model.py b/deepmd/pt/model/model/make_model.py index f857fc014b..472eae5329 100644 --- a/deepmd/pt/model/model/make_model.py +++ b/deepmd/pt/model/model/make_model.py @@ -514,8 +514,8 @@ def serialize(self) -> dict: def deserialize(cls, data) -> "CM": return cls(atomic_model_=T_AtomicModel.deserialize(data)) - def set_caseid(self, case_idx: int): - self.atomic_model.set_caseid(case_idx) + def set_case_embd(self, case_idx: int): + self.atomic_model.set_case_embd(case_idx) @torch.jit.export def get_dim_fparam(self) -> int: diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 74fc8ee066..65b64220ae 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -51,8 +51,8 @@ class DipoleFittingNet(GeneralFitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. - numb_caseid : int - Number of case specific identification. + dim_case_embd : int + Dimension of case specific embedding. activation_function : str Activation function. precision : str @@ -83,7 +83,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -106,7 +106,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/model/task/dos.py b/deepmd/pt/model/task/dos.py index fa8cfd110f..568ef81c92 100644 --- a/deepmd/pt/model/task/dos.py +++ b/deepmd/pt/model/task/dos.py @@ -47,7 +47,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, rcond: Optional[float] = None, bias_dos: Optional[torch.Tensor] = None, trainable: Union[bool, list[bool]] = True, @@ -74,7 +74,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 2f8ac9b9b0..07351b33f6 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -50,7 +50,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -68,7 +68,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index a68f4e0f64..2486ab576f 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -64,7 +64,7 @@ def share_params(self, base_class, shared_level, resume=False) -> None: self.__class__ == base_class.__class__ ), "Only fitting nets of the same type can share params!" if shared_level == 0: - # only not share the bias_atom_e and the caseid + # only not share the bias_atom_e and the case_embd # the following will successfully link all the params except buffers, which need manually link. for item in self._modules: self._modules[item] = base_class._modules[item] @@ -95,8 +95,8 @@ class GeneralFitting(Fitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. - numb_caseid : int - Number of case specific identification. + dim_case_embd : int + Dimension of case specific embedding. activation_function : str Activation function. precision : str @@ -134,7 +134,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -156,7 +156,7 @@ def __init__( self.resnet_dt = resnet_dt self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid + self.dim_case_embd = dim_case_embd self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] @@ -208,20 +208,20 @@ def __init__( else: self.aparam_avg, self.aparam_inv_std = None, None - if self.numb_caseid > 0: + if self.dim_case_embd > 0: self.register_buffer( - "caseid", - torch.zeros(self.numb_caseid, dtype=self.prec, device=device), - # torch.eye(self.numb_caseid, dtype=self.prec, device=device)[0], + "case_embd", + torch.zeros(self.dim_case_embd, dtype=self.prec, device=device), + # torch.eye(self.dim_case_embd, dtype=self.prec, device=device)[0], ) else: - self.caseid = None + self.case_embd = None in_dim = ( self.dim_descrpt + self.numb_fparam + (0 if self.use_aparam_as_mask else self.numb_aparam) - + self.numb_caseid + + self.dim_case_embd ) self.filter_layers = NetworkCollection( @@ -289,7 +289,7 @@ def serialize(self) -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "activation_function": self.activation_function, "precision": self.precision, "mixed_types": self.mixed_types, @@ -298,7 +298,7 @@ def serialize(self) -> dict: "exclude_types": self.exclude_types, "@variables": { "bias_atom_e": to_numpy_array(self.bias_atom_e), - "caseid": to_numpy_array(self.caseid), + "case_embd": to_numpy_array(self.case_embd), "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), "aparam_avg": to_numpy_array(self.aparam_avg), @@ -358,12 +358,12 @@ def get_type_map(self) -> list[str]: """Get the name to each type of atoms.""" return self.type_map - def set_caseid(self, case_idx: int): + def set_case_embd(self, case_idx: int): """ - Set the case identification of this fitting net by the given case_idx, + Set the case embedding of this fitting net by the given case_idx, typically concatenated with the output of the descriptor and fed into the fitting net. """ - self.caseid = torch.eye(self.numb_caseid, dtype=self.prec, device=device)[ + self.case_embd = torch.eye(self.dim_case_embd, dtype=self.prec, device=device)[ case_idx ] @@ -379,8 +379,8 @@ def __setitem__(self, key, value) -> None: self.aparam_avg = value elif key in ["aparam_inv_std"]: self.aparam_inv_std = value - elif key in ["caseid"]: - self.caseid = value + elif key in ["case_embd"]: + self.case_embd = value elif key in ["scale"]: self.scale = value else: @@ -397,8 +397,8 @@ def __getitem__(self, key): return self.aparam_avg elif key in ["aparam_inv_std"]: return self.aparam_inv_std - elif key in ["caseid"]: - return self.caseid + elif key in ["case_embd"]: + return self.case_embd elif key in ["scale"]: return self.scale else: @@ -497,16 +497,16 @@ def _forward_common( dim=-1, ) - if self.numb_caseid > 0: - assert self.caseid is not None - caseid = torch.tile(self.caseid.reshape([1, 1, -1]), [nf, nloc, 1]) + if self.dim_case_embd > 0: + assert self.case_embd is not None + case_embd = torch.tile(self.case_embd.reshape([1, 1, -1]), [nf, nloc, 1]) xx = torch.cat( - [xx, caseid], + [xx, case_embd], dim=-1, ) if xx_zeros is not None: xx_zeros = torch.cat( - [xx_zeros, caseid], + [xx_zeros, case_embd], dim=-1, ) diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 801ef6f7b9..b1599eac60 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -56,8 +56,8 @@ class InvarFitting(GeneralFitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. - numb_caseid : int - Number of case specific identification. + dim_case_embd : int + Dimension of case specific embedding. activation_function : str Activation function. precision : str @@ -93,7 +93,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -116,7 +116,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 66ae1aa180..d9a421d635 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -53,8 +53,8 @@ class PolarFittingNet(GeneralFitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. - numb_caseid : int - Number of case specific identification. + dim_case_embd : int + Dimension of case specific embedding. activation_function : str Activation function. precision : str @@ -87,7 +87,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -131,7 +131,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index b1c00c4004..dec0f1447b 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -60,8 +60,8 @@ class PropertyFittingNet(InvarFitting): Number of frame parameters. numb_aparam : int Number of atomic parameters. - numb_caseid : int - Number of case specific identification. + dim_case_embd : int + Dimension of case specific embedding. activation_function : str Activation function. precision : str @@ -85,7 +85,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, @@ -105,7 +105,7 @@ def __init__( resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, - numb_caseid=numb_caseid, + dim_case_embd=dim_case_embd, activation_function=activation_function, precision=precision, mixed_types=mixed_types, diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 19df530b9e..562f93875f 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -1275,37 +1275,39 @@ def get_model_for_wrapper(_model_params, resuming=False): else: _model = {} model_keys = list(_model_params["model_dict"]) - do_caseid, caseid_index = get_caseid_config(_model_params) + do_case_embd, case_embd_index = get_case_embd_config(_model_params) for _model_key in model_keys: _model[_model_key] = get_single_model( _model_params["model_dict"][_model_key], ) - if do_caseid and not resuming: - # only set caseid when from scratch multitask training - _model[_model_key].set_caseid(caseid_index[_model_key]) + if do_case_embd and not resuming: + # only set case_embd when from scratch multitask training + _model[_model_key].set_case_embd(case_embd_index[_model_key]) return _model -def get_caseid_config(_model_params): +def get_case_embd_config(_model_params): assert ( "model_dict" in _model_params ), "Only support setting data config for multi-task model!" model_keys = list(_model_params["model_dict"]) sorted_model_keys = sorted(model_keys) - numb_caseid_list = [ + numb_case_embd_list = [ _model_params["model_dict"][model_key] .get("fitting_net", {}) - .get("numb_caseid", 0) + .get("dim_case_embd", 0) for model_key in sorted_model_keys ] - if not all(item == numb_caseid_list[0] for item in numb_caseid_list): + if not all(item == numb_case_embd_list[0] for item in numb_case_embd_list): raise ValueError( - f"All models must have the same dimension of data identification, while the settings are: {numb_caseid_list}" + f"All models must have the same dimension of data identification, while the settings are: {numb_case_embd_list}" ) - if numb_caseid_list[0] == 0: + if numb_case_embd_list[0] == 0: return False, {} - caseid_index = {model_key: idx for idx, model_key in enumerate(sorted_model_keys)} - return True, caseid_index + case_embd_index = { + model_key: idx for idx, model_key in enumerate(sorted_model_keys) + } + return True, case_embd_index def model_change_out_bias( diff --git a/deepmd/tf/descriptor/se_a_ebd.py b/deepmd/tf/descriptor/se_a_ebd.py index 1e35fd001d..ae76308e69 100644 --- a/deepmd/tf/descriptor/se_a_ebd.py +++ b/deepmd/tf/descriptor/se_a_ebd.py @@ -58,8 +58,6 @@ class DescrptSeAEbd(DescrptSeA): Number of hidden layers for the type embedding net (skip connected). numb_aparam Number of atomic parameters. If >0 it will be embedded with atom types. - numb_caseid - Number of case specific identification. set_davg_zero Set the shift of embedding net input to zero. activation_function diff --git a/deepmd/tf/fit/dipole.py b/deepmd/tf/fit/dipole.py index 1c31f1bfce..4428d06536 100644 --- a/deepmd/tf/fit/dipole.py +++ b/deepmd/tf/fit/dipole.py @@ -58,8 +58,8 @@ class DipoleFittingSeA(Fitting): Number of frame parameters numb_aparam Number of atomic parameters - numb_caseid - Number of case specific identification. + dim_case_embd + Dimension of case specific embedding. sel_type : list[int] The atom types selected to have an atomic dipole prediction. If is None, all atoms are selected. seed : int @@ -86,7 +86,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, sel_type: Optional[list[int]] = None, seed: Optional[int] = None, activation_function: str = "tanh", @@ -122,13 +122,13 @@ def __init__( self.type_map = type_map self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid + self.dim_case_embd = dim_case_embd if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") - if numb_caseid > 0: - raise ValueError("numb_caseid is not supported in TensorFlow.") + if dim_case_embd > 0: + raise ValueError("dim_case_embd is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None @@ -401,7 +401,7 @@ def serialize(self, suffix: str) -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "exclude_types": [], diff --git a/deepmd/tf/fit/dos.py b/deepmd/tf/fit/dos.py index e3ea5f971a..1da0e55a92 100644 --- a/deepmd/tf/fit/dos.py +++ b/deepmd/tf/fit/dos.py @@ -74,8 +74,8 @@ class DOSFitting(Fitting): Number of frame parameter numb_aparam Number of atomic parameter - numb_caseid - Number of case specific identification. + dim_case_embd + Dimension of case specific embedding. ! numb_dos (added) Number of gridpoints on which the DOS is evaluated (NEDOS in VASP) rcond @@ -113,7 +113,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, numb_dos: int = 300, rcond: Optional[float] = None, trainable: Optional[list[bool]] = None, @@ -135,9 +135,9 @@ def __init__( self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid - if numb_caseid > 0: - raise ValueError("numb_caseid is not supported in TensorFlow.") + self.dim_case_embd = dim_case_embd + if dim_case_embd > 0: + raise ValueError("dim_case_embd is not supported in TensorFlow.") self.numb_dos = numb_dos @@ -715,7 +715,7 @@ def serialize(self, suffix: str = "") -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "rcond": self.rcond, "trainable": self.trainable, "activation_function": self.activation_function, @@ -738,7 +738,7 @@ def serialize(self, suffix: str = "") -> dict: "fparam_inv_std": self.fparam_inv_std, "aparam_avg": self.aparam_avg, "aparam_inv_std": self.aparam_inv_std, - "caseid": None, + "case_embd": None, }, "type_map": self.type_map, } diff --git a/deepmd/tf/fit/ener.py b/deepmd/tf/fit/ener.py index 2745b19eb5..068d3d8e35 100644 --- a/deepmd/tf/fit/ener.py +++ b/deepmd/tf/fit/ener.py @@ -117,8 +117,8 @@ class EnerFitting(Fitting): Number of frame parameter numb_aparam Number of atomic parameter - numb_caseid - Number of case specific identification. + dim_case_embd + Dimension of case specific embedding. rcond The condition number for the regression of atomic energy. tot_ener_zero @@ -158,7 +158,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, rcond: Optional[float] = None, tot_ener_zero: bool = False, trainable: Optional[list[bool]] = None, @@ -193,9 +193,9 @@ def __init__( # .add("trainable", [list, bool], default = True) self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid - if numb_caseid > 0: - raise ValueError("numb_caseid is not supported in TensorFlow.") + self.dim_case_embd = dim_case_embd + if dim_case_embd > 0: + raise ValueError("dim_case_embd is not supported in TensorFlow.") self.n_neuron = neuron self.resnet_dt = resnet_dt self.rcond = rcond @@ -920,7 +920,7 @@ def serialize(self, suffix: str = "") -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "rcond": self.rcond, "tot_ener_zero": self.tot_ener_zero, "trainable": self.trainable, @@ -952,7 +952,7 @@ def serialize(self, suffix: str = "") -> dict: "fparam_inv_std": self.fparam_inv_std, "aparam_avg": self.aparam_avg, "aparam_inv_std": self.aparam_inv_std, - "caseid": None, + "case_embd": None, }, "type_map": self.type_map, } diff --git a/deepmd/tf/fit/polar.py b/deepmd/tf/fit/polar.py index e42acd6739..14dd6ee092 100644 --- a/deepmd/tf/fit/polar.py +++ b/deepmd/tf/fit/polar.py @@ -63,8 +63,8 @@ class PolarFittingSeA(Fitting): Number of frame parameters numb_aparam Number of atomic parameters - numb_caseid - Number of case specific identification. + dim_case_embd + Dimension of case specific embedding. sel_type : list[int] The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected. fit_diag : bool @@ -97,7 +97,7 @@ def __init__( resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, - numb_caseid: int = 0, + dim_case_embd: int = 0, sel_type: Optional[list[int]] = None, fit_diag: bool = True, scale: Optional[list[float]] = None, @@ -165,13 +165,13 @@ def __init__( self.type_map = type_map self.numb_fparam = numb_fparam self.numb_aparam = numb_aparam - self.numb_caseid = numb_caseid + self.dim_case_embd = dim_case_embd if numb_fparam > 0: raise ValueError("numb_fparam is not supported in the dipole fitting") if numb_aparam > 0: raise ValueError("numb_aparam is not supported in the dipole fitting") - if numb_caseid > 0: - raise ValueError("numb_caseid is not supported in TensorFlow.") + if dim_case_embd > 0: + raise ValueError("dim_case_embd is not supported in TensorFlow.") self.fparam_avg = None self.fparam_std = None self.fparam_inv_std = None @@ -594,7 +594,7 @@ def serialize(self, suffix: str) -> dict: "resnet_dt": self.resnet_dt, "numb_fparam": self.numb_fparam, "numb_aparam": self.numb_aparam, - "numb_caseid": self.numb_caseid, + "dim_case_embd": self.dim_case_embd, "activation_function": self.activation_function_name, "precision": self.fitting_precision.name, "exclude_types": [], diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index c6efb52c5e..e0877b0e8f 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1433,7 +1433,7 @@ def descrpt_variant_type_args(exclude_hybrid: bool = False) -> Variant: def fitting_ener(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." - doc_numb_caseid = "The dimension of the case identification embedding. When training or fine-tuning a multitask model with case identification embeddings, this number should be set to the number of model branches." + doc_numb_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." @@ -1461,11 +1461,11 @@ def fitting_ener(): Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), Argument( - "numb_caseid", + "dim_case_embd", int, optional=True, default=0, - doc=doc_only_pt_supported + doc_numb_caseid, + doc=doc_only_pt_supported + doc_numb_case_embd, ), Argument( "neuron", @@ -1517,7 +1517,7 @@ def fitting_ener(): def fitting_dos(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." - doc_numb_caseid = "The dimension of the case identification embedding. When training or fine-tuning a multitask model with case identification embeddings, this number should be set to the number of model branches." + doc_numb_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." @@ -1535,11 +1535,11 @@ def fitting_dos(): Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), Argument( - "numb_caseid", + "dim_case_embd", int, optional=True, default=0, - doc=doc_only_pt_supported + doc_numb_caseid, + doc=doc_only_pt_supported + doc_numb_case_embd, ), Argument( "neuron", list[int], optional=True, default=[120, 120, 120], doc=doc_neuron @@ -1572,7 +1572,7 @@ def fitting_dos(): def fitting_property(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." - doc_numb_caseid = "The dimension of the case identification embedding. When training or fine-tuning a multitask model with case identification embeddings, this number should be set to the number of model branches." + doc_numb_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built" doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' @@ -1585,11 +1585,11 @@ def fitting_property(): Argument("numb_fparam", int, optional=True, default=0, doc=doc_numb_fparam), Argument("numb_aparam", int, optional=True, default=0, doc=doc_numb_aparam), Argument( - "numb_caseid", + "dim_case_embd", int, optional=True, default=0, - doc=doc_only_pt_supported + doc_numb_caseid, + doc=doc_only_pt_supported + doc_numb_case_embd, ), Argument( "neuron", @@ -1621,7 +1621,7 @@ def fitting_property(): def fitting_polar(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." - doc_numb_caseid = "The dimension of the case identification embedding. When training or fine-tuning a multitask model with case identification embeddings, this number should be set to the number of model branches." + doc_numb_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' @@ -1651,11 +1651,11 @@ def fitting_polar(): doc=doc_only_pt_supported + doc_numb_aparam, ), Argument( - "numb_caseid", + "dim_case_embd", int, optional=True, default=0, - doc=doc_only_pt_supported + doc_numb_caseid, + doc=doc_only_pt_supported + doc_numb_case_embd, ), Argument( "neuron", @@ -1699,7 +1699,7 @@ def fitting_polar(): def fitting_dipole(): doc_numb_fparam = "The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams." doc_numb_aparam = "The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams." - doc_numb_caseid = "The dimension of the case identification embedding. When training or fine-tuning a multitask model with case identification embeddings, this number should be set to the number of model branches." + doc_numb_case_embd = "The dimension of the case embedding embedding. When training or fine-tuning a multitask model with case embedding embeddings, this number should be set to the number of model branches." doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' @@ -1722,11 +1722,11 @@ def fitting_dipole(): doc=doc_only_pt_supported + doc_numb_aparam, ), Argument( - "numb_caseid", + "dim_case_embd", int, optional=True, default=0, - doc=doc_only_pt_supported + doc_numb_caseid, + doc=doc_only_pt_supported + doc_numb_case_embd, ), Argument( "neuron", diff --git a/doc/train/multi-task-training.md b/doc/train/multi-task-training.md index d33c33ae31..2194fd4701 100644 --- a/doc/train/multi-task-training.md +++ b/doc/train/multi-task-training.md @@ -60,12 +60,12 @@ Specifically, there are several parts that need to be modified: - Level 0: Shares all parameters (default) - Level 1: Shares type embedding only - Not all descriptors support all levels (e.g., se_a only supports level 0) -- - For fitting nets, we only support the default `shared_level`=0, where all parameters will be shared except for `bias_atom_e` and `caseid`. +- - For fitting nets, we only support the default `shared_level`=0, where all parameters will be shared except for `bias_atom_e` and `case_embd`. - - To conduct multitask training, there are two typical approaches: 1. **Descriptor sharing only**: Share the descriptor with `shared_level`=0. See [here](../../examples/water_multi_task/pytorch_example/input_torch.json) for an example. 2. **Descriptor and fitting network sharing with data identification**: - Share the descriptor and the fitting network with `shared_level`=0. - - {ref}`numb_caseid ` must be set to the number of model branches, which will distinguish different data tasks using a one-hot embedding. + - {ref}`dim_case_embd ` must be set to the number of model branches, which will distinguish different data tasks using a one-hot embedding. - See [here](../../examples/water_multi_task/pytorch_example/input_torch_sharefit.json) for an example. - - The parts that are exclusive to each model can be written following the previous definition. diff --git a/examples/water_multi_task/pytorch_example/input_torch_sharefit.json b/examples/water_multi_task/pytorch_example/input_torch_sharefit.json index 134a15a361..2fc23007c6 100644 --- a/examples/water_multi_task/pytorch_example/input_torch_sharefit.json +++ b/examples/water_multi_task/pytorch_example/input_torch_sharefit.json @@ -64,7 +64,7 @@ ], "resnet_dt": true, "seed": 1, - "numb_caseid": 2, + "dim_case_embd": 2, "_comment": " that's all" }, "_comment": "that's all" diff --git a/source/tests/pt/model/water/multitask_sharefit.json b/source/tests/pt/model/water/multitask_sharefit.json index d61a586f54..246b5992f7 100644 --- a/source/tests/pt/model/water/multitask_sharefit.json +++ b/source/tests/pt/model/water/multitask_sharefit.json @@ -26,7 +26,7 @@ "_comment": " that's all" }, "my_fitting": { - "numb_caseid": 2, + "dim_case_embd": 2, "neuron": [ 240, 240, diff --git a/source/tests/pt/test_multitask.py b/source/tests/pt/test_multitask.py index 7816beba4a..62964abad3 100644 --- a/source/tests/pt/test_multitask.py +++ b/source/tests/pt/test_multitask.py @@ -74,7 +74,7 @@ def test_multitask_train(self) -> None: self.share_fitting and "model_1.atomic_model.fitting_net" in state_key and "fitting_net.bias_atom_e" not in state_key - and "fitting_net.caseid" not in state_key + and "fitting_net.case_embd" not in state_key ): torch.testing.assert_close( multi_state_dict[state_key], diff --git a/source/tests/universal/dpmodel/fitting/test_fitting.py b/source/tests/universal/dpmodel/fitting/test_fitting.py index 6d173317fb..db199c02a3 100644 --- a/source/tests/universal/dpmodel/fitting/test_fitting.py +++ b/source/tests/universal/dpmodel/fitting/test_fitting.py @@ -39,7 +39,7 @@ def FittingParamEnergy( exclude_types=[], precision="float64", embedding_width=None, - numb_param=0, # test numb_fparam, numb_aparam and numb_caseid together + numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): input_dict = { "ntypes": ntypes, @@ -51,7 +51,7 @@ def FittingParamEnergy( "precision": precision, "numb_fparam": numb_param, "numb_aparam": numb_param, - "numb_caseid": numb_param, + "dim_case_embd": numb_param, } return input_dict @@ -78,7 +78,7 @@ def FittingParamDos( exclude_types=[], precision="float64", embedding_width=None, - numb_param=0, # test numb_fparam, numb_aparam and numb_caseid together + numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): input_dict = { "ntypes": ntypes, @@ -90,7 +90,7 @@ def FittingParamDos( "precision": precision, "numb_fparam": numb_param, "numb_aparam": numb_param, - "numb_caseid": numb_param, + "dim_case_embd": numb_param, } return input_dict @@ -117,7 +117,7 @@ def FittingParamDipole( exclude_types=[], precision="float64", embedding_width=None, - numb_param=0, # test numb_fparam, numb_aparam and numb_caseid together + numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): assert ( embedding_width is not None @@ -133,7 +133,7 @@ def FittingParamDipole( "precision": precision, "numb_fparam": numb_param, "numb_aparam": numb_param, - "numb_caseid": numb_param, + "dim_case_embd": numb_param, } return input_dict @@ -160,7 +160,7 @@ def FittingParamPolar( exclude_types=[], precision="float64", embedding_width=None, - numb_param=0, # test numb_fparam, numb_aparam and numb_caseid together + numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): assert embedding_width is not None, "embedding_width for polar fitting is required." input_dict = { @@ -174,7 +174,7 @@ def FittingParamPolar( "precision": precision, "numb_fparam": numb_param, "numb_aparam": numb_param, - "numb_caseid": numb_param, + "dim_case_embd": numb_param, } return input_dict @@ -201,7 +201,7 @@ def FittingParamProperty( exclude_types=[], precision="float64", embedding_width=None, - numb_param=0, # test numb_fparam, numb_aparam and numb_caseid together + numb_param=0, # test numb_fparam, numb_aparam and dim_case_embd together ): input_dict = { "ntypes": ntypes, @@ -213,7 +213,7 @@ def FittingParamProperty( "precision": precision, "numb_fparam": numb_param, "numb_aparam": numb_param, - "numb_caseid": numb_param, + "dim_case_embd": numb_param, } return input_dict