diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index e9c3033945..c22b9e9bfe 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -31,6 +31,7 @@ List, Optional, Tuple, + Union, ) from deepmd.dpmodel import ( @@ -53,188 +54,284 @@ ) -@BaseDescriptor.register("dpa2") -class DescrptDPA2(NativeOP, BaseDescriptor): +class RepinitArgs: def __init__( self, - # args for repinit - ntypes: int, - repinit_rcut: float, - repinit_rcut_smth: float, - repinit_nsel: int, - repformer_rcut: float, - repformer_rcut_smth: float, - repformer_nsel: int, - # kwargs for repinit - repinit_neuron: List[int] = [25, 50, 100], - repinit_axis_neuron: int = 16, - repinit_tebd_dim: int = 8, - repinit_tebd_input_mode: str = "concat", - repinit_set_davg_zero: bool = True, - repinit_activation_function="tanh", - repinit_resnet_dt: bool = False, - repinit_type_one_side: bool = False, - # kwargs for repformer - repformer_nlayers: int = 3, - repformer_g1_dim: int = 128, - repformer_g2_dim: int = 16, - repformer_axis_neuron: int = 4, - repformer_direct_dist: bool = False, - repformer_update_g1_has_conv: bool = True, - repformer_update_g1_has_drrd: bool = True, - repformer_update_g1_has_grrg: bool = True, - repformer_update_g1_has_attn: bool = True, - repformer_update_g2_has_g1g1: bool = True, - repformer_update_g2_has_attn: bool = True, - repformer_update_h2: bool = False, - repformer_attn1_hidden: int = 64, - repformer_attn1_nhead: int = 4, - repformer_attn2_hidden: int = 16, - repformer_attn2_nhead: int = 4, - repformer_attn2_has_gate: bool = False, - repformer_activation_function: str = "tanh", - repformer_update_style: str = "res_avg", - repformer_update_residual: float = 0.001, - repformer_update_residual_init: str = "norm", - repformer_set_davg_zero: bool = True, - repformer_trainable_ln: bool = True, - repformer_ln_eps: Optional[float] = 1e-5, - # kwargs for descriptor - concat_output_tebd: bool = True, - precision: str = "float64", - smooth: bool = True, - exclude_types: List[Tuple[int, int]] = [], - env_protection: float = 0.0, - trainable: bool = True, - seed: Optional[int] = None, - add_tebd_to_repinit_out: bool = False, + rcut: float, + rcut_smth: float, + nsel: int, + neuron: List[int] = [25, 50, 100], + axis_neuron: int = 16, + tebd_dim: int = 8, + tebd_input_mode: str = "concat", + set_davg_zero: bool = True, + activation_function="tanh", + resnet_dt: bool = False, + type_one_side: bool = False, ): - r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492. + r"""The constructor for the RepinitArgs class which defines the parameters of the repinit block in DPA2 descriptor. Parameters ---------- - repinit_rcut : float - (Used in the repinit block.) + rcut : float The cut-off radius. - repinit_rcut_smth : float - (Used in the repinit block.) + rcut_smth : float Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth. - repinit_nsel : int - (Used in the repinit block.) + nsel : int Maximally possible number of selected neighbors. - repinit_neuron : list, optional - (Used in the repinit block.) + neuron : list, optional Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built. - repinit_axis_neuron : int, optional - (Used in the repinit block.) + axis_neuron : int, optional Size of the submatrix of G (embedding matrix). - repinit_tebd_dim : int, optional - (Used in the repinit block.) + tebd_dim : int, optional The dimension of atom type embedding. - repinit_tebd_input_mode : str, optional - (Used in the repinit block.) + tebd_input_mode : str, optional The input mode of the type embedding. Supported modes are ['concat', 'strip']. - repinit_set_davg_zero : bool, optional - (Used in the repinit block.) + set_davg_zero : bool, optional Set the normalization average to zero. - repinit_activation_function : str, optional - (Used in the repinit block.) + activation_function : str, optional The activation function in the embedding net. - repinit_resnet_dt : bool, optional - (Used in the repinit block.) + resnet_dt : bool, optional Whether to use a "Timestep" in the skip connection. - repinit_type_one_side : bool, optional - (Used in the repinit block.) + type_one_side : bool, optional Whether to use one-side type embedding. - repformer_rcut : float - (Used in the repformer block.) + """ + self.rcut = rcut + self.rcut_smth = rcut_smth + self.nsel = nsel + self.neuron = neuron + self.axis_neuron = axis_neuron + self.tebd_dim = tebd_dim + self.tebd_input_mode = tebd_input_mode + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.resnet_dt = resnet_dt + self.type_one_side = type_one_side + + def __getitem__(self, key): + if hasattr(self, key): + return getattr(self, key) + else: + raise KeyError(key) + + def serialize(self) -> dict: + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.nsel, + "neuron": self.neuron, + "axis_neuron": self.axis_neuron, + "tebd_dim": self.tebd_dim, + "tebd_input_mode": self.tebd_input_mode, + "set_davg_zero": self.set_davg_zero, + "activation_function": self.activation_function, + "resnet_dt": self.resnet_dt, + "type_one_side": self.type_one_side, + } + + @classmethod + def deserialize(cls, data: dict) -> "RepinitArgs": + return cls(**data) + + +class RepformerArgs: + def __init__( + self, + rcut: float, + rcut_smth: float, + nsel: int, + nlayers: int = 3, + g1_dim: int = 128, + g2_dim: int = 16, + axis_neuron: int = 4, + direct_dist: bool = False, + update_g1_has_conv: bool = True, + update_g1_has_drrd: bool = True, + update_g1_has_grrg: bool = True, + update_g1_has_attn: bool = True, + update_g2_has_g1g1: bool = True, + update_g2_has_attn: bool = True, + update_h2: bool = False, + attn1_hidden: int = 64, + attn1_nhead: int = 4, + attn2_hidden: int = 16, + attn2_nhead: int = 4, + attn2_has_gate: bool = False, + activation_function: str = "tanh", + update_style: str = "res_avg", + update_residual: float = 0.001, + update_residual_init: str = "norm", + set_davg_zero: bool = True, + trainable_ln: bool = True, + ln_eps: Optional[float] = 1e-5, + ): + r"""The constructor for the RepformerArgs class which defines the parameters of the repformer block in DPA2 descriptor. + + Parameters + ---------- + rcut : float The cut-off radius. - repformer_rcut_smth : float - (Used in the repformer block.) + rcut_smth : float Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth. - repformer_nsel : int - (Used in the repformer block.) + nsel : int Maximally possible number of selected neighbors. - repformer_nlayers : int, optional - (Used in the repformer block.) + nlayers : int, optional Number of repformer layers. - repformer_g1_dim : int, optional - (Used in the repformer block.) + g1_dim : int, optional Dimension of the first graph convolution layer. - repformer_g2_dim : int, optional - (Used in the repformer block.) + g2_dim : int, optional Dimension of the second graph convolution layer. - repformer_axis_neuron : int, optional - (Used in the repformer block.) + axis_neuron : int, optional Size of the submatrix of G (embedding matrix). - repformer_direct_dist : bool, optional - (Used in the repformer block.) + direct_dist : bool, optional Whether to use direct distance information (1/r term) in the repformer block. - repformer_update_g1_has_conv : bool, optional - (Used in the repformer block.) + update_g1_has_conv : bool, optional Whether to update the g1 rep with convolution term. - repformer_update_g1_has_drrd : bool, optional - (Used in the repformer block.) + update_g1_has_drrd : bool, optional Whether to update the g1 rep with the drrd term. - repformer_update_g1_has_grrg : bool, optional - (Used in the repformer block.) + update_g1_has_grrg : bool, optional Whether to update the g1 rep with the grrg term. - repformer_update_g1_has_attn : bool, optional - (Used in the repformer block.) + update_g1_has_attn : bool, optional Whether to update the g1 rep with the localized self-attention. - repformer_update_g2_has_g1g1 : bool, optional - (Used in the repformer block.) + update_g2_has_g1g1 : bool, optional Whether to update the g2 rep with the g1xg1 term. - repformer_update_g2_has_attn : bool, optional - (Used in the repformer block.) + update_g2_has_attn : bool, optional Whether to update the g2 rep with the gated self-attention. - repformer_update_h2 : bool, optional - (Used in the repformer block.) + update_h2 : bool, optional Whether to update the h2 rep. - repformer_attn1_hidden : int, optional - (Used in the repformer block.) + attn1_hidden : int, optional The hidden dimension of localized self-attention to update the g1 rep. - repformer_attn1_nhead : int, optional - (Used in the repformer block.) + attn1_nhead : int, optional The number of heads in localized self-attention to update the g1 rep. - repformer_attn2_hidden : int, optional - (Used in the repformer block.) + attn2_hidden : int, optional The hidden dimension of gated self-attention to update the g2 rep. - repformer_attn2_nhead : int, optional - (Used in the repformer block.) + attn2_nhead : int, optional The number of heads in gated self-attention to update the g2 rep. - repformer_attn2_has_gate : bool, optional - (Used in the repformer block.) + attn2_has_gate : bool, optional Whether to use gate in the gated self-attention to update the g2 rep. - repformer_activation_function : str, optional - (Used in the repformer block.) + activation_function : str, optional The activation function in the embedding net. - repformer_update_style : str, optional - (Used in the repformer block.) + update_style : str, optional Style to update a representation. Supported options are: -'res_avg': Updates a rep `u` with: u = 1/\\sqrt{n+1} (u + u_1 + u_2 + ... + u_n) -'res_incr': Updates a rep `u` with: u = u + 1/\\sqrt{n} (u_1 + u_2 + ... + u_n) -'res_residual': Updates a rep `u` with: u = u + (r1*u_1 + r2*u_2 + ... + r3*u_n) - where `r1`, `r2` ... `r3` are residual weights defined by `repformer_update_residual` - and `repformer_update_residual_init`. - repformer_update_residual : float, optional - (Used in the repformer block.) + where `r1`, `r2` ... `r3` are residual weights defined by `update_residual` + and `update_residual_init`. + update_residual : float, optional When update using residual mode, the initial std of residual vector weights. - repformer_update_residual_init : str, optional - (Used in the repformer block.) + update_residual_init : str, optional When update using residual mode, the initialization mode of residual vector weights. - repformer_set_davg_zero : bool, optional - (Used in the repformer block.) + set_davg_zero : bool, optional Set the normalization average to zero. - repformer_trainable_ln : bool, optional - (Used in the repformer block.) + trainable_ln : bool, optional Whether to use trainable shift and scale weights in layer normalization. - repformer_ln_eps : float, optional - (Used in the repformer block.) + ln_eps : float, optional The epsilon value for layer normalization. + """ + self.rcut = rcut + self.rcut_smth = rcut_smth + self.nsel = nsel + self.nlayers = nlayers + self.g1_dim = g1_dim + self.g2_dim = g2_dim + self.axis_neuron = axis_neuron + self.direct_dist = direct_dist + self.update_g1_has_conv = update_g1_has_conv + self.update_g1_has_drrd = update_g1_has_drrd + self.update_g1_has_grrg = update_g1_has_grrg + self.update_g1_has_attn = update_g1_has_attn + self.update_g2_has_g1g1 = update_g2_has_g1g1 + self.update_g2_has_attn = update_g2_has_attn + self.update_h2 = update_h2 + self.attn1_hidden = attn1_hidden + self.attn1_nhead = attn1_nhead + self.attn2_hidden = attn2_hidden + self.attn2_nhead = attn2_nhead + self.attn2_has_gate = attn2_has_gate + self.activation_function = activation_function + self.update_style = update_style + self.update_residual = update_residual + self.update_residual_init = update_residual_init + self.set_davg_zero = set_davg_zero + self.trainable_ln = trainable_ln + # to keep consistent with default value in this backends + if ln_eps is None: + ln_eps = 1e-5 + self.ln_eps = ln_eps + + def __getitem__(self, key): + if hasattr(self, key): + return getattr(self, key) + else: + raise KeyError(key) + + def serialize(self) -> dict: + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "nsel": self.nsel, + "nlayers": self.nlayers, + "g1_dim": self.g1_dim, + "g2_dim": self.g2_dim, + "axis_neuron": self.axis_neuron, + "direct_dist": self.direct_dist, + "update_g1_has_conv": self.update_g1_has_conv, + "update_g1_has_drrd": self.update_g1_has_drrd, + "update_g1_has_grrg": self.update_g1_has_grrg, + "update_g1_has_attn": self.update_g1_has_attn, + "update_g2_has_g1g1": self.update_g2_has_g1g1, + "update_g2_has_attn": self.update_g2_has_attn, + "update_h2": self.update_h2, + "attn1_hidden": self.attn1_hidden, + "attn1_nhead": self.attn1_nhead, + "attn2_hidden": self.attn2_hidden, + "attn2_nhead": self.attn2_nhead, + "attn2_has_gate": self.attn2_has_gate, + "activation_function": self.activation_function, + "update_style": self.update_style, + "update_residual": self.update_residual, + "update_residual_init": self.update_residual_init, + "set_davg_zero": self.set_davg_zero, + "trainable_ln": self.trainable_ln, + "ln_eps": self.ln_eps, + } + + @classmethod + def deserialize(cls, data: dict) -> "RepformerArgs": + return cls(**data) + + +@BaseDescriptor.register("dpa2") +class DescrptDPA2(NativeOP, BaseDescriptor): + def __init__( + self, + ntypes: int, + # args for repinit + repinit: Union[RepinitArgs, dict], + # args for repformer + repformer: Union[RepformerArgs, dict], + # kwargs for descriptor + concat_output_tebd: bool = True, + precision: str = "float64", + smooth: bool = True, + exclude_types: List[Tuple[int, int]] = [], + env_protection: float = 0.0, + trainable: bool = True, + seed: Optional[int] = None, + add_tebd_to_repinit_out: bool = False, + ): + r"""The DPA-2 descriptor. see https://arxiv.org/abs/2312.15492. + + Parameters + ---------- + repinit : Union[RepinitArgs, dict] + The arguments used to initialize the repinit block, see docstr in `RepinitArgs` for details information. + repformer : Union[RepformerArgs, dict] + The arguments used to initialize the repformer block, see docstr in `RepformerArgs` for details information. concat_output_tebd : bool, optional Whether to concat type embedding at the output of the descriptor. precision : str, optional @@ -269,65 +366,76 @@ def __init__( The switch function for decaying inverse distance. """ - # to keep consistent with default value in this backends - if repformer_ln_eps is None: - repformer_ln_eps = 1e-5 + + def init_subclass_params(sub_data, sub_class): + if isinstance(sub_data, dict): + return sub_class(**sub_data) + elif isinstance(sub_data, sub_class): + return sub_data + else: + raise ValueError( + f"Input args must be a {sub_class.__name__} class or a dict!" + ) + + self.repinit_args = init_subclass_params(repinit, RepinitArgs) + self.repformer_args = init_subclass_params(repformer, RepformerArgs) + self.repinit = DescrptBlockSeAtten( - repinit_rcut, - repinit_rcut_smth, - repinit_nsel, + self.repinit_args.rcut, + self.repinit_args.rcut_smth, + self.repinit_args.nsel, ntypes, attn_layer=0, - neuron=repinit_neuron, - axis_neuron=repinit_axis_neuron, - tebd_dim=repinit_tebd_dim, - tebd_input_mode=repinit_tebd_input_mode, - set_davg_zero=repinit_set_davg_zero, + neuron=self.repinit_args.neuron, + axis_neuron=self.repinit_args.axis_neuron, + tebd_dim=self.repinit_args.tebd_dim, + tebd_input_mode=self.repinit_args.tebd_input_mode, + set_davg_zero=self.repinit_args.set_davg_zero, exclude_types=exclude_types, env_protection=env_protection, - activation_function=repinit_activation_function, + activation_function=self.repinit_args.activation_function, precision=precision, - resnet_dt=repinit_resnet_dt, + resnet_dt=self.repinit_args.resnet_dt, smooth=smooth, - type_one_side=repinit_type_one_side, + type_one_side=self.repinit_args.type_one_side, ) self.repformers = DescrptBlockRepformers( - repformer_rcut, - repformer_rcut_smth, - repformer_nsel, + self.repformer_args.rcut, + self.repformer_args.rcut_smth, + self.repformer_args.nsel, ntypes, - nlayers=repformer_nlayers, - g1_dim=repformer_g1_dim, - g2_dim=repformer_g2_dim, - axis_neuron=repformer_axis_neuron, - direct_dist=repformer_direct_dist, - update_g1_has_conv=repformer_update_g1_has_conv, - update_g1_has_drrd=repformer_update_g1_has_drrd, - update_g1_has_grrg=repformer_update_g1_has_grrg, - update_g1_has_attn=repformer_update_g1_has_attn, - update_g2_has_g1g1=repformer_update_g2_has_g1g1, - update_g2_has_attn=repformer_update_g2_has_attn, - update_h2=repformer_update_h2, - attn1_hidden=repformer_attn1_hidden, - attn1_nhead=repformer_attn1_nhead, - attn2_hidden=repformer_attn2_hidden, - attn2_nhead=repformer_attn2_nhead, - attn2_has_gate=repformer_attn2_has_gate, - activation_function=repformer_activation_function, - update_style=repformer_update_style, - update_residual=repformer_update_residual, - update_residual_init=repformer_update_residual_init, - set_davg_zero=repformer_set_davg_zero, + nlayers=self.repformer_args.nlayers, + g1_dim=self.repformer_args.g1_dim, + g2_dim=self.repformer_args.g2_dim, + axis_neuron=self.repformer_args.axis_neuron, + direct_dist=self.repformer_args.direct_dist, + update_g1_has_conv=self.repformer_args.update_g1_has_conv, + update_g1_has_drrd=self.repformer_args.update_g1_has_drrd, + update_g1_has_grrg=self.repformer_args.update_g1_has_grrg, + update_g1_has_attn=self.repformer_args.update_g1_has_attn, + update_g2_has_g1g1=self.repformer_args.update_g2_has_g1g1, + update_g2_has_attn=self.repformer_args.update_g2_has_attn, + update_h2=self.repformer_args.update_h2, + attn1_hidden=self.repformer_args.attn1_hidden, + attn1_nhead=self.repformer_args.attn1_nhead, + attn2_hidden=self.repformer_args.attn2_hidden, + attn2_nhead=self.repformer_args.attn2_nhead, + attn2_has_gate=self.repformer_args.attn2_has_gate, + activation_function=self.repformer_args.activation_function, + update_style=self.repformer_args.update_style, + update_residual=self.repformer_args.update_residual, + update_residual_init=self.repformer_args.update_residual_init, + set_davg_zero=self.repformer_args.set_davg_zero, smooth=smooth, exclude_types=exclude_types, env_protection=env_protection, precision=precision, - trainable_ln=repformer_trainable_ln, - ln_eps=repformer_ln_eps, + trainable_ln=self.repformer_args.trainable_ln, + ln_eps=self.repformer_args.ln_eps, ) self.type_embedding = TypeEmbedNet( ntypes=ntypes, - neuron=[repinit_tebd_dim], + neuron=[self.repinit_args.tebd_dim], padding=True, activation_function="Linear", precision=precision, @@ -352,7 +460,7 @@ def __init__( self.tebd_transform = None if self.add_tebd_to_repinit_out: self.tebd_transform = NativeLayer( - repinit_tebd_dim, + self.repinit_args.tebd_dim, self.repformers.dim_in, bias=False, precision=precision, @@ -360,7 +468,7 @@ def __init__( assert self.repinit.rcut > self.repformers.rcut assert self.repinit.sel[0] > self.repformers.sel[0] - self.tebd_dim = repinit_tebd_dim + self.tebd_dim = self.repinit_args.tebd_dim self.rcut = self.repinit.get_rcut() self.ntypes = ntypes self.sel = self.repinit.sel @@ -524,42 +632,8 @@ def serialize(self) -> dict: "type": "dpa2", "@version": 1, "ntypes": self.ntypes, - "repinit_rcut": repinit.rcut, - "repinit_rcut_smth": repinit.rcut_smth, - "repinit_nsel": repinit.sel, - "repformer_rcut": repformers.rcut, - "repformer_rcut_smth": repformers.rcut_smth, - "repformer_nsel": repformers.sel, - "repinit_neuron": repinit.neuron, - "repinit_axis_neuron": repinit.axis_neuron, - "repinit_tebd_dim": repinit.tebd_dim, - "repinit_tebd_input_mode": repinit.tebd_input_mode, - "repinit_set_davg_zero": repinit.set_davg_zero, - "repinit_activation_function": repinit.activation_function, - "repinit_resnet_dt": repinit.resnet_dt, - "repinit_type_one_side": repinit.type_one_side, - "repformer_nlayers": repformers.nlayers, - "repformer_g1_dim": repformers.g1_dim, - "repformer_g2_dim": repformers.g2_dim, - "repformer_axis_neuron": repformers.axis_neuron, - "repformer_direct_dist": repformers.direct_dist, - "repformer_update_g1_has_conv": repformers.update_g1_has_conv, - "repformer_update_g1_has_drrd": repformers.update_g1_has_drrd, - "repformer_update_g1_has_grrg": repformers.update_g1_has_grrg, - "repformer_update_g1_has_attn": repformers.update_g1_has_attn, - "repformer_update_g2_has_g1g1": repformers.update_g2_has_g1g1, - "repformer_update_g2_has_attn": repformers.update_g2_has_attn, - "repformer_update_h2": repformers.update_h2, - "repformer_attn1_hidden": repformers.attn1_hidden, - "repformer_attn1_nhead": repformers.attn1_nhead, - "repformer_attn2_hidden": repformers.attn2_hidden, - "repformer_attn2_nhead": repformers.attn2_nhead, - "repformer_attn2_has_gate": repformers.attn2_has_gate, - "repformer_activation_function": repformers.activation_function, - "repformer_update_style": repformers.update_style, - "repformer_set_davg_zero": repformers.set_davg_zero, - "repformer_trainable_ln": repformers.trainable_ln, - "repformer_ln_eps": repformers.ln_eps, + "repinit_args": self.repinit_args.serialize(), + "repformer_args": self.repformer_args.serialize(), "concat_output_tebd": self.concat_output_tebd, "precision": self.precision, "smooth": self.smooth, @@ -599,8 +673,8 @@ def serialize(self) -> dict: } data.update( { - "repinit": repinit_variable, - "repformers": repformers_variable, + "repinit_variable": repinit_variable, + "repformers_variable": repformers_variable, } ) return data @@ -611,12 +685,14 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": check_version_compatibility(data.pop("@version"), 1, 1) data.pop("@class") data.pop("type") - repinit_variable = data.pop("repinit").copy() - repformers_variable = data.pop("repformers").copy() + repinit_variable = data.pop("repinit_variable").copy() + repformers_variable = data.pop("repformers_variable").copy() type_embedding = data.pop("type_embedding") g1_shape_tranform = data.pop("g1_shape_tranform") tebd_transform = data.pop("tebd_transform", None) add_tebd_to_repinit_out = data["add_tebd_to_repinit_out"] + data["repinit"] = RepinitArgs(**data.pop("repinit_args")) + data["repformer"] = RepformerArgs(**data.pop("repformer_args")) obj = cls(**data) obj.type_embedding = TypeEmbedNet.deserialize(type_embedding) if add_tebd_to_repinit_out: @@ -628,7 +704,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": # deserialize repinit statistic_repinit = repinit_variable.pop("@variables") env_mat = repinit_variable.pop("env_mat") - tebd_input_mode = data["repinit_tebd_input_mode"] + tebd_input_mode = data["repinit"].tebd_input_mode obj.repinit.embeddings = NetworkCollection.deserialize( repinit_variable.pop("embeddings") ) @@ -666,18 +742,18 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy = update_sel.update_one_sel( + local_jdata_cpy["repinit"] = update_sel.update_one_sel( global_jdata, - local_jdata_cpy, + local_jdata_cpy["repinit"], True, - rcut_key="repinit_rcut", - sel_key="repinit_nsel", + rcut_key="rcut", + sel_key="nsel", ) - local_jdata_cpy = update_sel.update_one_sel( + local_jdata_cpy["repformer"] = update_sel.update_one_sel( global_jdata, - local_jdata_cpy, + local_jdata_cpy["repformer"], True, - rcut_key="repformer_rcut", - sel_key="repformer_nsel", + rcut_key="rcut", + sel_key="nsel", ) return local_jdata_cpy diff --git a/deepmd/pt/model/descriptor/dpa2.py b/deepmd/pt/model/descriptor/dpa2.py index baa1d10c41..c2910e9414 100644 --- a/deepmd/pt/model/descriptor/dpa2.py +++ b/deepmd/pt/model/descriptor/dpa2.py @@ -10,6 +10,10 @@ import torch +from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, +) from deepmd.dpmodel.utils import EnvMat as DPEnvMat from deepmd.pt.model.network.mlp import ( Identity, @@ -58,48 +62,11 @@ class DescrptDPA2(BaseDescriptor, torch.nn.Module): def __init__( self, - # args for repinit ntypes: int, - repinit_rcut: float, - repinit_rcut_smth: float, - repinit_nsel: int, - repformer_rcut: float, - repformer_rcut_smth: float, - repformer_nsel: int, - # kwargs for repinit - repinit_neuron: List[int] = [25, 50, 100], - repinit_axis_neuron: int = 16, - repinit_tebd_dim: int = 8, - repinit_tebd_input_mode: str = "concat", - repinit_set_davg_zero: bool = True, - repinit_activation_function="tanh", - repinit_resnet_dt: bool = False, - repinit_type_one_side: bool = False, - # kwargs for repformer - repformer_nlayers: int = 3, - repformer_g1_dim: int = 128, - repformer_g2_dim: int = 16, - repformer_axis_neuron: int = 4, - repformer_direct_dist: bool = False, - repformer_update_g1_has_conv: bool = True, - repformer_update_g1_has_drrd: bool = True, - repformer_update_g1_has_grrg: bool = True, - repformer_update_g1_has_attn: bool = True, - repformer_update_g2_has_g1g1: bool = True, - repformer_update_g2_has_attn: bool = True, - repformer_update_h2: bool = False, - repformer_attn1_hidden: int = 64, - repformer_attn1_nhead: int = 4, - repformer_attn2_hidden: int = 16, - repformer_attn2_nhead: int = 4, - repformer_attn2_has_gate: bool = False, - repformer_activation_function: str = "tanh", - repformer_update_style: str = "res_avg", - repformer_update_residual: float = 0.001, - repformer_update_residual_init: str = "norm", - repformer_set_davg_zero: bool = True, - repformer_trainable_ln: bool = True, - repformer_ln_eps: Optional[float] = 1e-5, + # args for repinit + repinit: Union[RepinitArgs, dict], + # args for repformer + repformer: Union[RepformerArgs, dict], # kwargs for descriptor concat_output_tebd: bool = True, precision: str = "float64", @@ -115,128 +82,10 @@ def __init__( Parameters ---------- - repinit_rcut : float - (Used in the repinit block.) - The cut-off radius. - repinit_rcut_smth : float - (Used in the repinit block.) - Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth. - repinit_nsel : int - (Used in the repinit block.) - Maximally possible number of selected neighbors. - repinit_neuron : list, optional - (Used in the repinit block.) - Number of neurons in each hidden layers of the embedding net. - When two layers are of the same size or one layer is twice as large as the previous layer, - a skip connection is built. - repinit_axis_neuron : int, optional - (Used in the repinit block.) - Size of the submatrix of G (embedding matrix). - repinit_tebd_dim : int, optional - (Used in the repinit block.) - The dimension of atom type embedding. - repinit_tebd_input_mode : str, optional - (Used in the repinit block.) - The input mode of the type embedding. Supported modes are ['concat', 'strip']. - repinit_set_davg_zero : bool, optional - (Used in the repinit block.) - Set the normalization average to zero. - repinit_activation_function : str, optional - (Used in the repinit block.) - The activation function in the embedding net. - repinit_resnet_dt : bool, optional - (Used in the repinit block.) - Whether to use a "Timestep" in the skip connection. - repinit_type_one_side : bool, optional - (Used in the repinit block.) - Whether to use one-side type embedding. - repformer_rcut : float - (Used in the repformer block.) - The cut-off radius. - repformer_rcut_smth : float - (Used in the repformer block.) - Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth. - repformer_nsel : int - (Used in the repformer block.) - Maximally possible number of selected neighbors. - repformer_nlayers : int, optional - (Used in the repformer block.) - Number of repformer layers. - repformer_g1_dim : int, optional - (Used in the repformer block.) - Dimension of the first graph convolution layer. - repformer_g2_dim : int, optional - (Used in the repformer block.) - Dimension of the second graph convolution layer. - repformer_axis_neuron : int, optional - (Used in the repformer block.) - Size of the submatrix of G (embedding matrix). - repformer_direct_dist : bool, optional - (Used in the repformer block.) - Whether to use direct distance information (1/r term) in the repformer block. - repformer_update_g1_has_conv : bool, optional - (Used in the repformer block.) - Whether to update the g1 rep with convolution term. - repformer_update_g1_has_drrd : bool, optional - (Used in the repformer block.) - Whether to update the g1 rep with the drrd term. - repformer_update_g1_has_grrg : bool, optional - (Used in the repformer block.) - Whether to update the g1 rep with the grrg term. - repformer_update_g1_has_attn : bool, optional - (Used in the repformer block.) - Whether to update the g1 rep with the localized self-attention. - repformer_update_g2_has_g1g1 : bool, optional - (Used in the repformer block.) - Whether to update the g2 rep with the g1xg1 term. - repformer_update_g2_has_attn : bool, optional - (Used in the repformer block.) - Whether to update the g2 rep with the gated self-attention. - repformer_update_h2 : bool, optional - (Used in the repformer block.) - Whether to update the h2 rep. - repformer_attn1_hidden : int, optional - (Used in the repformer block.) - The hidden dimension of localized self-attention to update the g1 rep. - repformer_attn1_nhead : int, optional - (Used in the repformer block.) - The number of heads in localized self-attention to update the g1 rep. - repformer_attn2_hidden : int, optional - (Used in the repformer block.) - The hidden dimension of gated self-attention to update the g2 rep. - repformer_attn2_nhead : int, optional - (Used in the repformer block.) - The number of heads in gated self-attention to update the g2 rep. - repformer_attn2_has_gate : bool, optional - (Used in the repformer block.) - Whether to use gate in the gated self-attention to update the g2 rep. - repformer_activation_function : str, optional - (Used in the repformer block.) - The activation function in the embedding net. - repformer_update_style : str, optional - (Used in the repformer block.) - Style to update a representation. - Supported options are: - -'res_avg': Updates a rep `u` with: u = 1/\\sqrt{n+1} (u + u_1 + u_2 + ... + u_n) - -'res_incr': Updates a rep `u` with: u = u + 1/\\sqrt{n} (u_1 + u_2 + ... + u_n) - -'res_residual': Updates a rep `u` with: u = u + (r1*u_1 + r2*u_2 + ... + r3*u_n) - where `r1`, `r2` ... `r3` are residual weights defined by `repformer_update_residual` - and `repformer_update_residual_init`. - repformer_update_residual : float, optional - (Used in the repformer block.) - When update using residual mode, the initial std of residual vector weights. - repformer_update_residual_init : str, optional - (Used in the repformer block.) - When update using residual mode, the initialization mode of residual vector weights. - repformer_set_davg_zero : bool, optional - (Used in the repformer block.) - Set the normalization average to zero. - repformer_trainable_ln : bool, optional - (Used in the repformer block.) - Whether to use trainable shift and scale weights in layer normalization. - repformer_ln_eps : float, optional - (Used in the repformer block.) - The epsilon value for layer normalization. + repinit : Union[RepinitArgs, dict] + The arguments used to initialize the repinit block, see docstr in `RepinitArgs` for details information. + repformer : Union[RepformerArgs, dict] + The arguments used to initialize the repformer block, see docstr in `RepformerArgs` for details information. concat_output_tebd : bool, optional Whether to concat type embedding at the output of the descriptor. precision : str, optional @@ -272,65 +121,76 @@ def __init__( """ super().__init__() - # to keep consistent with default value in this backends - if repformer_ln_eps is None: - repformer_ln_eps = 1e-5 + + def init_subclass_params(sub_data, sub_class): + if isinstance(sub_data, dict): + return sub_class(**sub_data) + elif isinstance(sub_data, sub_class): + return sub_data + else: + raise ValueError( + f"Input args must be a {sub_class.__name__} class or a dict!" + ) + + self.repinit_args = init_subclass_params(repinit, RepinitArgs) + self.repformer_args = init_subclass_params(repformer, RepformerArgs) + self.repinit = DescrptBlockSeAtten( - repinit_rcut, - repinit_rcut_smth, - repinit_nsel, + self.repinit_args.rcut, + self.repinit_args.rcut_smth, + self.repinit_args.nsel, ntypes, attn_layer=0, - neuron=repinit_neuron, - axis_neuron=repinit_axis_neuron, - tebd_dim=repinit_tebd_dim, - tebd_input_mode=repinit_tebd_input_mode, - set_davg_zero=repinit_set_davg_zero, + neuron=self.repinit_args.neuron, + axis_neuron=self.repinit_args.axis_neuron, + tebd_dim=self.repinit_args.tebd_dim, + tebd_input_mode=self.repinit_args.tebd_input_mode, + set_davg_zero=self.repinit_args.set_davg_zero, exclude_types=exclude_types, env_protection=env_protection, - activation_function=repinit_activation_function, + activation_function=self.repinit_args.activation_function, precision=precision, - resnet_dt=repinit_resnet_dt, + resnet_dt=self.repinit_args.resnet_dt, smooth=smooth, - type_one_side=repinit_type_one_side, + type_one_side=self.repinit_args.type_one_side, ) self.repformers = DescrptBlockRepformers( - repformer_rcut, - repformer_rcut_smth, - repformer_nsel, + self.repformer_args.rcut, + self.repformer_args.rcut_smth, + self.repformer_args.nsel, ntypes, - nlayers=repformer_nlayers, - g1_dim=repformer_g1_dim, - g2_dim=repformer_g2_dim, - axis_neuron=repformer_axis_neuron, - direct_dist=repformer_direct_dist, - update_g1_has_conv=repformer_update_g1_has_conv, - update_g1_has_drrd=repformer_update_g1_has_drrd, - update_g1_has_grrg=repformer_update_g1_has_grrg, - update_g1_has_attn=repformer_update_g1_has_attn, - update_g2_has_g1g1=repformer_update_g2_has_g1g1, - update_g2_has_attn=repformer_update_g2_has_attn, - update_h2=repformer_update_h2, - attn1_hidden=repformer_attn1_hidden, - attn1_nhead=repformer_attn1_nhead, - attn2_hidden=repformer_attn2_hidden, - attn2_nhead=repformer_attn2_nhead, - attn2_has_gate=repformer_attn2_has_gate, - activation_function=repformer_activation_function, - update_style=repformer_update_style, - update_residual=repformer_update_residual, - update_residual_init=repformer_update_residual_init, - set_davg_zero=repformer_set_davg_zero, + nlayers=self.repformer_args.nlayers, + g1_dim=self.repformer_args.g1_dim, + g2_dim=self.repformer_args.g2_dim, + axis_neuron=self.repformer_args.axis_neuron, + direct_dist=self.repformer_args.direct_dist, + update_g1_has_conv=self.repformer_args.update_g1_has_conv, + update_g1_has_drrd=self.repformer_args.update_g1_has_drrd, + update_g1_has_grrg=self.repformer_args.update_g1_has_grrg, + update_g1_has_attn=self.repformer_args.update_g1_has_attn, + update_g2_has_g1g1=self.repformer_args.update_g2_has_g1g1, + update_g2_has_attn=self.repformer_args.update_g2_has_attn, + update_h2=self.repformer_args.update_h2, + attn1_hidden=self.repformer_args.attn1_hidden, + attn1_nhead=self.repformer_args.attn1_nhead, + attn2_hidden=self.repformer_args.attn2_hidden, + attn2_nhead=self.repformer_args.attn2_nhead, + attn2_has_gate=self.repformer_args.attn2_has_gate, + activation_function=self.repformer_args.activation_function, + update_style=self.repformer_args.update_style, + update_residual=self.repformer_args.update_residual, + update_residual_init=self.repformer_args.update_residual_init, + set_davg_zero=self.repformer_args.set_davg_zero, smooth=smooth, exclude_types=exclude_types, env_protection=env_protection, precision=precision, - trainable_ln=repformer_trainable_ln, - ln_eps=repformer_ln_eps, + trainable_ln=self.repformer_args.trainable_ln, + ln_eps=self.repformer_args.ln_eps, old_impl=old_impl, ) self.type_embedding = TypeEmbedNet( - ntypes, repinit_tebd_dim, precision=precision + ntypes, self.repinit_args.tebd_dim, precision=precision ) self.concat_output_tebd = concat_output_tebd self.precision = precision @@ -353,7 +213,7 @@ def __init__( self.tebd_transform = None if self.add_tebd_to_repinit_out: self.tebd_transform = MLPLayer( - repinit_tebd_dim, + self.repinit_args.tebd_dim, self.repformers.dim_in, bias=False, precision=precision, @@ -361,7 +221,7 @@ def __init__( assert self.repinit.rcut > self.repformers.rcut assert self.repinit.sel[0] > self.repformers.sel[0] - self.tebd_dim = repinit_tebd_dim + self.tebd_dim = self.repinit_args.tebd_dim self.rcut = self.repinit.get_rcut() self.rcut_smth = self.repinit.get_rcut_smth() self.ntypes = ntypes @@ -499,42 +359,8 @@ def serialize(self) -> dict: "type": "dpa2", "@version": 1, "ntypes": self.ntypes, - "repinit_rcut": repinit.rcut, - "repinit_rcut_smth": repinit.rcut_smth, - "repinit_nsel": repinit.sel, - "repformer_rcut": repformers.rcut, - "repformer_rcut_smth": repformers.rcut_smth, - "repformer_nsel": repformers.sel, - "repinit_neuron": repinit.neuron, - "repinit_axis_neuron": repinit.axis_neuron, - "repinit_tebd_dim": repinit.tebd_dim, - "repinit_tebd_input_mode": repinit.tebd_input_mode, - "repinit_set_davg_zero": repinit.set_davg_zero, - "repinit_activation_function": repinit.activation_function, - "repinit_resnet_dt": repinit.resnet_dt, - "repinit_type_one_side": repinit.type_one_side, - "repformer_nlayers": repformers.nlayers, - "repformer_g1_dim": repformers.g1_dim, - "repformer_g2_dim": repformers.g2_dim, - "repformer_axis_neuron": repformers.axis_neuron, - "repformer_direct_dist": repformers.direct_dist, - "repformer_update_g1_has_conv": repformers.update_g1_has_conv, - "repformer_update_g1_has_drrd": repformers.update_g1_has_drrd, - "repformer_update_g1_has_grrg": repformers.update_g1_has_grrg, - "repformer_update_g1_has_attn": repformers.update_g1_has_attn, - "repformer_update_g2_has_g1g1": repformers.update_g2_has_g1g1, - "repformer_update_g2_has_attn": repformers.update_g2_has_attn, - "repformer_update_h2": repformers.update_h2, - "repformer_attn1_hidden": repformers.attn1_hidden, - "repformer_attn1_nhead": repformers.attn1_nhead, - "repformer_attn2_hidden": repformers.attn2_hidden, - "repformer_attn2_nhead": repformers.attn2_nhead, - "repformer_attn2_has_gate": repformers.attn2_has_gate, - "repformer_activation_function": repformers.activation_function, - "repformer_update_style": repformers.update_style, - "repformer_set_davg_zero": repformers.set_davg_zero, - "repformer_trainable_ln": repformers.trainable_ln, - "repformer_ln_eps": repformers.ln_eps, + "repinit_args": self.repinit_args.serialize(), + "repformer_args": self.repformer_args.serialize(), "concat_output_tebd": self.concat_output_tebd, "precision": self.precision, "smooth": self.smooth, @@ -574,8 +400,8 @@ def serialize(self) -> dict: } data.update( { - "repinit": repinit_variable, - "repformers": repformers_variable, + "repinit_variable": repinit_variable, + "repformers_variable": repformers_variable, } ) return data @@ -586,12 +412,14 @@ def deserialize(cls, data: dict) -> "DescrptDPA2": check_version_compatibility(data.pop("@version"), 1, 1) data.pop("@class") data.pop("type") - repinit_variable = data.pop("repinit").copy() - repformers_variable = data.pop("repformers").copy() + repinit_variable = data.pop("repinit_variable").copy() + repformers_variable = data.pop("repformers_variable").copy() type_embedding = data.pop("type_embedding") g1_shape_tranform = data.pop("g1_shape_tranform") tebd_transform = data.pop("tebd_transform", None) add_tebd_to_repinit_out = data["add_tebd_to_repinit_out"] + data["repinit"] = RepinitArgs(**data.pop("repinit_args")) + data["repformer"] = RepformerArgs(**data.pop("repformer_args")) obj = cls(**data) obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize( type_embedding @@ -608,7 +436,7 @@ def t_cvt(xx): # deserialize repinit statistic_repinit = repinit_variable.pop("@variables") env_mat = repinit_variable.pop("env_mat") - tebd_input_mode = data["repinit_tebd_input_mode"] + tebd_input_mode = data["repinit"].tebd_input_mode obj.repinit.filter_layers = NetworkCollection.deserialize( repinit_variable.pop("embeddings") ) @@ -737,18 +565,18 @@ def update_sel(cls, global_jdata: dict, local_jdata: dict): """ local_jdata_cpy = local_jdata.copy() update_sel = UpdateSel() - local_jdata_cpy = update_sel.update_one_sel( + local_jdata_cpy["repinit"] = update_sel.update_one_sel( global_jdata, - local_jdata_cpy, + local_jdata_cpy["repinit"], True, - rcut_key="repinit_rcut", - sel_key="repinit_nsel", + rcut_key="rcut", + sel_key="nsel", ) - local_jdata_cpy = update_sel.update_one_sel( + local_jdata_cpy["repformer"] = update_sel.update_one_sel( global_jdata, - local_jdata_cpy, + local_jdata_cpy["repformer"], True, - rcut_key="repformer_rcut", - sel_key="repformer_nsel", + rcut_key="rcut", + sel_key="nsel", ) return local_jdata_cpy diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index b1f5cda0b2..6929712eec 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -618,109 +618,9 @@ def descrpt_se_atten_v2_args(): @descrpt_args_plugin.register("dpa2", doc=doc_only_pt_supported) def descrpt_dpa2_args(): # repinit args - doc_repinit = "(Used in the repinit block.) " - doc_repinit_rcut = f"{doc_repinit}The cut-off radius." - doc_repinit_rcut_smth = f"{doc_repinit}Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." - doc_repinit_nsel = f"{doc_repinit}Maximally possible number of selected neighbors." - doc_repinit_neuron = ( - f"{doc_repinit}Number of neurons in each hidden layers of the embedding net." - f"When two layers are of the same size or one layer is twice as large as the previous layer, " - f"a skip connection is built." - ) - doc_repinit_axis_neuron = ( - f"{doc_repinit}Size of the submatrix of G (embedding matrix)." - ) - doc_repinit_tebd_dim = f"{doc_repinit}The dimension of atom type embedding." - doc_repinit_tebd_input_mode = ( - f"{doc_repinit}The input mode of the type embedding. Supported modes are ['concat', 'strip']." - "- 'concat': Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. " - "When `type_one_side` is False, the input is `input_ij = concat([r_ij, tebd_j, tebd_i])`. When `type_one_side` is True, the input is `input_ij = concat([r_ij, tebd_j])`. " - "The output is `out_ij = embeding(input_ij)` for the pair-wise representation of atom i with neighbor j." - "- 'strip': Use a separated embedding network for the type embedding and combine the output with the radial embedding network output. " - f"When `type_one_side` is False, the input is `input_t = concat([tebd_j, tebd_i])`. {doc_only_pt_supported} When `type_one_side` is True, the input is `input_t = tebd_j`. " - "The output is `out_ij = embeding_t(input_t) * embeding_s(r_ij) + embeding_s(r_ij)` for the pair-wise representation of atom i with neighbor j." - ) - doc_repinit_set_davg_zero = ( - f"{doc_repinit}Set the normalization average to zero. " - f"This option should be set when `atom_ener` in the energy fitting is used." - ) - doc_repinit_activation_function = f"{doc_repinit}The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}." - doc_repinit_type_one_side = ( - f"{doc_repinit}" - + r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." - ) - doc_repinit_resnet_dt = ( - f'{doc_repinit}Whether to use a "Timestep" in the skip connection.' - ) - + doc_repinit = "The arguments used to initialize the repinit block." # repformer args - doc_repformer = "(Used in the repformer block.) " - doc_repformer_rcut = f"{doc_repformer}The cut-off radius." - doc_repformer_rcut_smth = f"{doc_repformer}Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." - doc_repformer_nsel = ( - f"{doc_repformer}Maximally possible number of selected neighbors." - ) - doc_repformer_nlayers = f"{doc_repformer}The number of repformer layers." - doc_repformer_g1_dim = ( - f"{doc_repformer}The dimension of invariant single-atom representation." - ) - doc_repformer_g2_dim = ( - f"{doc_repformer}The dimension of invariant pair-atom representation." - ) - doc_repformer_axis_neuron = f"{doc_repformer}The number of dimension of submatrix in the symmetrization ops." - doc_repformer_direct_dist = f"{doc_repformer}Whether or not use direct distance as input for the embedding net to get g2 instead of smoothed 1/r." - doc_repformer_update_g1_has_conv = ( - f"{doc_repformer}Update the g1 rep with convolution term." - ) - doc_repformer_update_g1_has_drrd = ( - f"{doc_repformer}Update the g1 rep with the drrd term." - ) - doc_repformer_update_g1_has_grrg = ( - f"{doc_repformer}Update the g1 rep with the grrg term." - ) - doc_repformer_update_g1_has_attn = ( - f"{doc_repformer}Update the g1 rep with the localized self-attention." - ) - doc_repformer_update_g2_has_g1g1 = ( - f"{doc_repformer}Update the g2 rep with the g1xg1 term." - ) - doc_repformer_update_g2_has_attn = ( - f"{doc_repformer}Update the g2 rep with the gated self-attention." - ) - doc_repformer_update_h2 = f"{doc_repformer}Update the h2 rep." - doc_repformer_attn1_hidden = f"{doc_repformer}The hidden dimension of localized self-attention to update the g1 rep." - doc_repformer_attn1_nhead = f"{doc_repformer}The number of heads in localized self-attention to update the g1 rep." - doc_repformer_attn2_hidden = f"{doc_repformer}The hidden dimension of gated self-attention to update the g2 rep." - doc_repformer_attn2_nhead = f"{doc_repformer}The number of heads in gated self-attention to update the g2 rep." - doc_repformer_attn2_has_gate = f"{doc_repformer}Whether to use gate in the gated self-attention to update the g2 rep." - doc_repformer_activation_function = f"{doc_repformer}The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}." - doc_repformer_update_style = ( - f"{doc_repformer}Style to update a representation. " - f"Supported options are: " - "-'res_avg': Updates a rep `u` with: u = 1/\\sqrt{n+1} (u + u_1 + u_2 + ... + u_n) " - "-'res_incr': Updates a rep `u` with: u = u + 1/\\sqrt{n} (u_1 + u_2 + ... + u_n)" - "-'res_residual': Updates a rep `u` with: u = u + (r1*u_1 + r2*u_2 + ... + r3*u_n) " - "where `r1`, `r2` ... `r3` are residual weights defined by `repformer_update_residual` " - "and `repformer_update_residual_init`." - ) - doc_repformer_update_residual = ( - f"{doc_repformer}When update using residual mode, " - "the initial std of residual vector weights." - ) - doc_repformer_update_residual_init = ( - f"{doc_repformer}When update using residual mode, " - "the initialization mode of residual vector weights." - "Supported modes are: ['norm', 'const']." - ) - doc_repformer_set_davg_zero = ( - f"{doc_repformer}Set the normalization average to zero. " - f"This option should be set when `atom_ener` in the energy fitting is used." - ) - doc_repformer_trainable_ln = ( - "Whether to use trainable shift and scale weights in layer normalization." - ) - doc_repformer_ln_eps = "The epsilon value for layer normalization. The default value for TensorFlow is set to 1e-3 to keep consistent with keras while set to 1e-5 in PyTorch and DP implementation." - + doc_repformer = "The arguments used to initialize the repformer block." # descriptor args doc_concat_output_tebd = ( "Whether to concat type embedding at the output of the descriptor." @@ -736,270 +636,367 @@ def descrpt_dpa2_args(): doc_add_tebd_to_repinit_out = "Add type embedding to the output representation from repinit before inputting it into repformer." return [ # repinit args - Argument("repinit_rcut", float, doc=doc_repinit_rcut), - Argument("repinit_rcut_smth", float, doc=doc_repinit_rcut_smth), - Argument("repinit_nsel", int, doc=doc_repinit_nsel), + Argument("repinit", dict, dpa2_repinit_args(), doc=doc_repinit), + # repformer args + Argument("repformer", dict, dpa2_repformer_args(), doc=doc_repformer), + # descriptor args + Argument( + "concat_output_tebd", + bool, + optional=True, + default=True, + doc=doc_concat_output_tebd, + ), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument("smooth", bool, optional=True, default=True, doc=doc_smooth), + Argument( + "exclude_types", + List[List[int]], + optional=True, + default=[], + doc=doc_exclude_types, + ), + Argument( + "env_protection", + float, + optional=True, + default=0.0, + doc=doc_only_pt_supported + doc_env_protection, + ), + Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), + Argument("seed", [int, None], optional=True, doc=doc_seed), + Argument( + "add_tebd_to_repinit_out", + bool, + optional=True, + default=False, + alias=["repformer_add_type_ebd_to_seq"], + doc=doc_add_tebd_to_repinit_out, + ), + ] + + +# repinit for dpa2 +def dpa2_repinit_args(): + # repinit args + doc_rcut = "The cut-off radius." + doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." + doc_nsel = "Maximally possible number of selected neighbors." + doc_neuron = ( + "Number of neurons in each hidden layers of the embedding net." + "When two layers are of the same size or one layer is twice as large as the previous layer, " + "a skip connection is built." + ) + doc_axis_neuron = "Size of the submatrix of G (embedding matrix)." + doc_tebd_dim = "The dimension of atom type embedding." + doc_tebd_input_mode = ( + "The input mode of the type embedding. Supported modes are ['concat', 'strip']." + "- 'concat': Concatenate the type embedding with the smoothed radial information as the union input for the embedding network. " + "When `type_one_side` is False, the input is `input_ij = concat([r_ij, tebd_j, tebd_i])`. When `type_one_side` is True, the input is `input_ij = concat([r_ij, tebd_j])`. " + "The output is `out_ij = embeding(input_ij)` for the pair-wise representation of atom i with neighbor j." + "- 'strip': Use a separated embedding network for the type embedding and combine the output with the radial embedding network output. " + f"When `type_one_side` is False, the input is `input_t = concat([tebd_j, tebd_i])`. {doc_only_pt_supported} When `type_one_side` is True, the input is `input_t = tebd_j`. " + "The output is `out_ij = embeding_t(input_t) * embeding_s(r_ij) + embeding_s(r_ij)` for the pair-wise representation of atom i with neighbor j." + ) + doc_set_davg_zero = ( + "Set the normalization average to zero. " + "This option should be set when `atom_ener` in the energy fitting is used." + ) + doc_activation_function = f"The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}." + doc_type_one_side = r"If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters." + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection.' + + return [ + # repinit args + Argument("rcut", float, doc=doc_rcut), + Argument("rcut_smth", float, doc=doc_rcut_smth), + Argument("nsel", int, doc=doc_nsel), Argument( - "repinit_neuron", + "neuron", list, optional=True, default=[25, 50, 100], - doc=doc_repinit_neuron, + doc=doc_neuron, ), Argument( - "repinit_axis_neuron", + "axis_neuron", int, optional=True, default=16, - doc=doc_repinit_axis_neuron, + doc=doc_axis_neuron, ), Argument( - "repinit_tebd_dim", + "tebd_dim", int, optional=True, default=8, - alias=["tebd_dim"], - doc=doc_repinit_tebd_dim, + doc=doc_tebd_dim, ), Argument( - "repinit_tebd_input_mode", + "tebd_input_mode", str, optional=True, default="concat", - doc=doc_repinit_tebd_input_mode, + doc=doc_tebd_input_mode, ), Argument( - "repinit_set_davg_zero", + "set_davg_zero", bool, optional=True, default=True, - doc=doc_repinit_set_davg_zero, + doc=doc_set_davg_zero, ), Argument( - "repinit_activation_function", + "activation_function", str, optional=True, default="tanh", - alias=["repinit_activation"], - doc=doc_repinit_activation_function, + doc=doc_activation_function, ), Argument( - "repinit_type_one_side", + "type_one_side", bool, optional=True, default=False, - doc=doc_repinit_type_one_side, + doc=doc_type_one_side, ), Argument( - "repinit_resnet_dt", + "resnet_dt", bool, optional=True, default=False, - doc=doc_repinit_resnet_dt, + doc=doc_resnet_dt, ), + ] + + +# repformer for dpa2 +def dpa2_repformer_args(): + # repformer args + doc_rcut = "The cut-off radius." + doc_rcut_smth = "Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`." + doc_nsel = "Maximally possible number of selected neighbors." + doc_nlayers = "The number of repformer layers." + doc_g1_dim = "The dimension of invariant single-atom representation." + doc_g2_dim = "The dimension of invariant pair-atom representation." + doc_axis_neuron = "The number of dimension of submatrix in the symmetrization ops." + doc_direct_dist = "Whether or not use direct distance as input for the embedding net to get g2 instead of smoothed 1/r." + doc_update_g1_has_conv = "Update the g1 rep with convolution term." + doc_update_g1_has_drrd = "Update the g1 rep with the drrd term." + doc_update_g1_has_grrg = "Update the g1 rep with the grrg term." + doc_update_g1_has_attn = "Update the g1 rep with the localized self-attention." + doc_update_g2_has_g1g1 = "Update the g2 rep with the g1xg1 term." + doc_update_g2_has_attn = "Update the g2 rep with the gated self-attention." + doc_update_h2 = "Update the h2 rep." + doc_attn1_hidden = ( + "The hidden dimension of localized self-attention to update the g1 rep." + ) + doc_attn1_nhead = ( + "The number of heads in localized self-attention to update the g1 rep." + ) + doc_attn2_hidden = ( + "The hidden dimension of gated self-attention to update the g2 rep." + ) + doc_attn2_nhead = ( + "The number of heads in gated self-attention to update the g2 rep." + ) + doc_attn2_has_gate = ( + "Whether to use gate in the gated self-attention to update the g2 rep." + ) + doc_activation_function = f"The activation function in the embedding net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())}." + doc_update_style = ( + "Style to update a representation. " + "Supported options are: " + "-'res_avg': Updates a rep `u` with: u = 1/\\sqrt{n+1} (u + u_1 + u_2 + ... + u_n) " + "-'res_incr': Updates a rep `u` with: u = u + 1/\\sqrt{n} (u_1 + u_2 + ... + u_n)" + "-'res_residual': Updates a rep `u` with: u = u + (r1*u_1 + r2*u_2 + ... + r3*u_n) " + "where `r1`, `r2` ... `r3` are residual weights defined by `update_residual` " + "and `update_residual_init`." + ) + doc_update_residual = ( + "When update using residual mode, " + "the initial std of residual vector weights." + ) + doc_update_residual_init = ( + "When update using residual mode, " + "the initialization mode of residual vector weights." + "Supported modes are: ['norm', 'const']." + ) + doc_set_davg_zero = ( + "Set the normalization average to zero. " + "This option should be set when `atom_ener` in the energy fitting is used." + ) + doc_trainable_ln = ( + "Whether to use trainable shift and scale weights in layer normalization." + ) + doc_ln_eps = "The epsilon value for layer normalization. The default value for TensorFlow is set to 1e-3 to keep consistent with keras while set to 1e-5 in PyTorch and DP implementation." + + return [ # repformer args - Argument("repformer_rcut", float, doc=doc_repformer_rcut), - Argument("repformer_rcut_smth", float, doc=doc_repformer_rcut_smth), - Argument("repformer_nsel", int, doc=doc_repformer_nsel), + Argument("rcut", float, doc=doc_rcut), + Argument("rcut_smth", float, doc=doc_rcut_smth), + Argument("nsel", int, doc=doc_nsel), Argument( - "repformer_nlayers", + "nlayers", int, optional=True, default=3, - doc=doc_repformer_nlayers, + doc=doc_nlayers, ), Argument( - "repformer_g1_dim", + "g1_dim", int, optional=True, default=128, - doc=doc_repformer_g1_dim, - ), - Argument( - "repformer_g2_dim", int, optional=True, default=16, doc=doc_repformer_g2_dim + doc=doc_g1_dim, ), + Argument("g2_dim", int, optional=True, default=16, doc=doc_g2_dim), Argument( - "repformer_axis_neuron", + "axis_neuron", int, optional=True, default=4, - alias=["repformer_axis_dim"], - doc=doc_repformer_axis_neuron, + doc=doc_axis_neuron, ), Argument( - "repformer_direct_dist", + "direct_dist", bool, optional=True, default=False, - doc=doc_repformer_direct_dist, + doc=doc_direct_dist, ), Argument( - "repformer_update_g1_has_conv", + "update_g1_has_conv", bool, optional=True, default=True, - doc=doc_repformer_update_g1_has_conv, + doc=doc_update_g1_has_conv, ), Argument( - "repformer_update_g1_has_drrd", + "update_g1_has_drrd", bool, optional=True, default=True, - doc=doc_repformer_update_g1_has_drrd, + doc=doc_update_g1_has_drrd, ), Argument( - "repformer_update_g1_has_grrg", + "update_g1_has_grrg", bool, optional=True, default=True, - doc=doc_repformer_update_g1_has_grrg, + doc=doc_update_g1_has_grrg, ), Argument( - "repformer_update_g1_has_attn", + "update_g1_has_attn", bool, optional=True, default=True, - doc=doc_repformer_update_g1_has_attn, + doc=doc_update_g1_has_attn, ), Argument( - "repformer_update_g2_has_g1g1", + "update_g2_has_g1g1", bool, optional=True, default=True, - doc=doc_repformer_update_g2_has_g1g1, + doc=doc_update_g2_has_g1g1, ), Argument( - "repformer_update_g2_has_attn", + "update_g2_has_attn", bool, optional=True, default=True, - doc=doc_repformer_update_g2_has_attn, + doc=doc_update_g2_has_attn, ), Argument( - "repformer_update_h2", + "update_h2", bool, optional=True, default=False, - doc=doc_repformer_update_h2, + doc=doc_update_h2, ), Argument( - "repformer_attn1_hidden", + "attn1_hidden", int, optional=True, default=64, - doc=doc_repformer_attn1_hidden, + doc=doc_attn1_hidden, ), Argument( - "repformer_attn1_nhead", + "attn1_nhead", int, optional=True, default=4, - doc=doc_repformer_attn1_nhead, + doc=doc_attn1_nhead, ), Argument( - "repformer_attn2_hidden", + "attn2_hidden", int, optional=True, default=16, - doc=doc_repformer_attn2_hidden, + doc=doc_attn2_hidden, ), Argument( - "repformer_attn2_nhead", + "attn2_nhead", int, optional=True, default=4, - doc=doc_repformer_attn2_nhead, + doc=doc_attn2_nhead, ), Argument( - "repformer_attn2_has_gate", + "attn2_has_gate", bool, optional=True, default=False, - doc=doc_repformer_attn2_has_gate, + doc=doc_attn2_has_gate, ), Argument( - "repformer_activation_function", + "activation_function", str, optional=True, default="tanh", - alias=["repformer_activation"], - doc=doc_repformer_activation_function, + doc=doc_activation_function, ), Argument( - "repformer_update_style", + "update_style", str, optional=True, default="res_avg", - doc=doc_repformer_update_style, + doc=doc_update_style, ), Argument( - "repformer_update_residual", + "update_residual", float, optional=True, default=0.001, - doc=doc_repformer_update_residual, + doc=doc_update_residual, ), Argument( - "repformer_update_residual_init", + "update_residual_init", str, optional=True, default="norm", - doc=doc_repformer_update_residual_init, + doc=doc_update_residual_init, ), Argument( - "repformer_set_davg_zero", + "set_davg_zero", bool, optional=True, default=True, - doc=doc_repformer_set_davg_zero, + doc=doc_set_davg_zero, ), Argument( - "repformer_trainable_ln", + "trainable_ln", bool, optional=True, default=True, - doc=doc_repformer_trainable_ln, + doc=doc_trainable_ln, ), Argument( - "repformer_ln_eps", + "ln_eps", float, optional=True, default=None, - doc=doc_repformer_ln_eps, - ), - # descriptor args - Argument( - "concat_output_tebd", - bool, - optional=True, - default=True, - doc=doc_concat_output_tebd, - ), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument("smooth", bool, optional=True, default=False, doc=doc_smooth), - Argument( - "exclude_types", - List[List[int]], - optional=True, - default=[], - doc=doc_exclude_types, - ), - Argument( - "env_protection", - float, - optional=True, - default=0.0, - doc=doc_only_pt_supported + doc_env_protection, - ), - Argument("trainable", bool, optional=True, default=True, doc=doc_trainable), - Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument( - "add_tebd_to_repinit_out", - bool, - optional=True, - default=False, - alias=["repformer_add_type_ebd_to_seq"], - doc=doc_add_tebd_to_repinit_out, + doc=doc_ln_eps, ), ] diff --git a/examples/water/dpa2/input_torch.json b/examples/water/dpa2/input_torch.json index 7096e1d40b..ba8f2e5967 100644 --- a/examples/water/dpa2/input_torch.json +++ b/examples/water/dpa2/input_torch.json @@ -7,36 +7,40 @@ ], "descriptor": { "type": "dpa2", - "repinit_tebd_dim": 8, - "repinit_rcut": 9.0, - "repinit_rcut_smth": 8.0, - "repinit_nsel": 120, - "repformer_rcut": 4.0, - "repformer_rcut_smth": 3.5, - "repformer_nsel": 40, - "repinit_neuron": [ - 25, - 50, - 100 - ], - "repinit_axis_neuron": 12, - "repinit_activation_function": "tanh", - "repformer_nlayers": 12, - "repformer_g1_dim": 128, - "repformer_g2_dim": 32, - "repformer_attn2_hidden": 32, - "repformer_attn2_nhead": 4, - "repformer_attn1_hidden": 128, - "repformer_attn1_nhead": 4, - "repformer_axis_neuron": 4, - "repformer_update_h2": false, - "repformer_update_g1_has_conv": true, - "repformer_update_g1_has_grrg": true, - "repformer_update_g1_has_drrd": true, - "repformer_update_g1_has_attn": true, - "repformer_update_g2_has_g1g1": true, - "repformer_update_g2_has_attn": true, - "repformer_attn2_has_gate": true, + "repinit": { + "tebd_dim": 8, + "rcut": 9.0, + "rcut_smth": 8.0, + "nsel": 120, + "neuron": [ + 25, + 50, + 100 + ], + "axis_neuron": 12, + "activation_function": "tanh" + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 3.5, + "nsel": 40, + "nlayers": 12, + "g1_dim": 128, + "g2_dim": 32, + "attn2_hidden": 32, + "attn2_nhead": 4, + "attn1_hidden": 128, + "attn1_nhead": 4, + "axis_neuron": 4, + "update_h2": false, + "update_g1_has_conv": true, + "update_g1_has_grrg": true, + "update_g1_has_drrd": true, + "update_g1_has_attn": true, + "update_g2_has_g1g1": true, + "update_g2_has_attn": true, + "attn2_has_gate": true + }, "add_tebd_to_repinit_out": false }, "fitting_net": { diff --git a/source/tests/common/dpmodel/test_descriptor_dpa2.py b/source/tests/common/dpmodel/test_descriptor_dpa2.py index 3ae0689dad..cf6789faf6 100644 --- a/source/tests/common/dpmodel/test_descriptor_dpa2.py +++ b/source/tests/common/dpmodel/test_descriptor_dpa2.py @@ -6,6 +6,10 @@ from deepmd.dpmodel.descriptor import ( DescrptDPA2, ) +from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, +) from .case_single_frame_with_nlist import ( TestCaseSingleFrameWithNlist, @@ -28,14 +32,21 @@ def test_self_consistency( dstd = 0.1 + np.abs(dstd) dstd_2 = 0.1 + np.abs(dstd_2) + repinit = RepinitArgs( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + nsel=self.sel_mix, + ) + repformer = RepformerArgs( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth, + nsel=nnei // 2, + ) + em0 = DescrptDPA2( ntypes=self.nt, - repinit_rcut=self.rcut, - repinit_rcut_smth=self.rcut_smth, - repinit_nsel=self.sel_mix, - repformer_rcut=self.rcut / 2, - repformer_rcut_smth=self.rcut_smth, - repformer_nsel=nnei // 2, + repinit=repinit, + repformer=repformer, ) em0.repinit.mean = davg diff --git a/source/tests/consistent/descriptor/test_dpa2.py b/source/tests/consistent/descriptor/test_dpa2.py index dd868dbd59..a3cf5303bd 100644 --- a/source/tests/consistent/descriptor/test_dpa2.py +++ b/source/tests/consistent/descriptor/test_dpa2.py @@ -32,6 +32,10 @@ # not implemented DescrptDPA2TF = None +from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, +) from deepmd.utils.argcheck import ( descrpt_se_atten_args, ) @@ -88,45 +92,53 @@ def data(self) -> dict: ) = self.param return { "ntypes": self.ntypes, - "repinit_rcut": 6.00, - "repinit_rcut_smth": 5.80, - "repinit_nsel": 10, - "repformer_rcut": 4.00, - "repformer_rcut_smth": 3.50, - "repformer_nsel": 8, # kwargs for repinit - "repinit_neuron": [6, 12, 24], - "repinit_axis_neuron": 3, - "repinit_tebd_dim": 4, - "repinit_tebd_input_mode": repinit_tebd_input_mode, - "repinit_set_davg_zero": repinit_set_davg_zero, - "repinit_activation_function": "tanh", - "repinit_type_one_side": repinit_type_one_side, + "repinit": RepinitArgs( + **{ + "rcut": 6.00, + "rcut_smth": 5.80, + "nsel": 10, + "neuron": [6, 12, 24], + "axis_neuron": 3, + "tebd_dim": 4, + "tebd_input_mode": repinit_tebd_input_mode, + "set_davg_zero": repinit_set_davg_zero, + "activation_function": "tanh", + "type_one_side": repinit_type_one_side, + } + ), # kwargs for repformer - "repformer_nlayers": 3, - "repformer_g1_dim": 20, - "repformer_g2_dim": 10, - "repformer_axis_neuron": 3, - "repformer_direct_dist": repformer_direct_dist, - "repformer_update_g1_has_conv": repformer_update_g1_has_conv, - "repformer_update_g1_has_drrd": repformer_update_g1_has_drrd, - "repformer_update_g1_has_grrg": repformer_update_g1_has_grrg, - "repformer_update_g1_has_attn": repformer_update_g1_has_attn, - "repformer_update_g2_has_g1g1": repformer_update_g2_has_g1g1, - "repformer_update_g2_has_attn": repformer_update_g2_has_attn, - "repformer_update_h2": repformer_update_h2, - "repformer_attn1_hidden": 12, - "repformer_attn1_nhead": 2, - "repformer_attn2_hidden": 10, - "repformer_attn2_nhead": 2, - "repformer_attn2_has_gate": repformer_attn2_has_gate, - "repformer_activation_function": "tanh", - "repformer_update_style": repformer_update_style, - "repformer_update_residual": 0.001, - "repformer_update_residual_init": repformer_update_residual_init, - "repformer_set_davg_zero": True, - "repformer_trainable_ln": repformer_trainable_ln, - "repformer_ln_eps": repformer_ln_eps, + "repformer": RepformerArgs( + **{ + "rcut": 4.00, + "rcut_smth": 3.50, + "nsel": 8, + "nlayers": 3, + "g1_dim": 20, + "g2_dim": 10, + "axis_neuron": 3, + "direct_dist": repformer_direct_dist, + "update_g1_has_conv": repformer_update_g1_has_conv, + "update_g1_has_drrd": repformer_update_g1_has_drrd, + "update_g1_has_grrg": repformer_update_g1_has_grrg, + "update_g1_has_attn": repformer_update_g1_has_attn, + "update_g2_has_g1g1": repformer_update_g2_has_g1g1, + "update_g2_has_attn": repformer_update_g2_has_attn, + "update_h2": repformer_update_h2, + "attn1_hidden": 12, + "attn1_nhead": 2, + "attn2_hidden": 10, + "attn2_nhead": 2, + "attn2_has_gate": repformer_attn2_has_gate, + "activation_function": "tanh", + "update_style": repformer_update_style, + "update_residual": 0.001, + "update_residual_init": repformer_update_residual_init, + "set_davg_zero": True, + "trainable_ln": repformer_trainable_ln, + "ln_eps": repformer_ln_eps, + } + ), # kwargs for descriptor "concat_output_tebd": True, "precision": precision, diff --git a/source/tests/pt/model/models/dpa2.json b/source/tests/pt/model/models/dpa2.json index 966e298dfa..ca1948492a 100644 --- a/source/tests/pt/model/models/dpa2.json +++ b/source/tests/pt/model/models/dpa2.json @@ -5,35 +5,40 @@ ], "descriptor": { "type": "dpa2", - "repinit_rcut": 6.0, - "repinit_rcut_smth": 2.0, - "repinit_nsel": 30, - "repformer_rcut": 4.0, - "repformer_rcut_smth": 0.5, - "repformer_nsel": 10, - "repinit_neuron": [ - 2, - 4, - 8 - ], - "repinit_axis_neuron": 4, - "repinit_activation_function": "tanh", - "repformer_nlayers": 12, - "repformer_g1_dim": 8, - "repformer_g2_dim": 5, - "repformer_attn2_hidden": 3, - "repformer_attn2_nhead": 1, - "repformer_attn1_hidden": 5, - "repformer_attn1_nhead": 1, - "repformer_axis_neuron": 4, - "repformer_update_h2": false, - "repformer_update_g1_has_conv": true, - "repformer_update_g1_has_grrg": true, - "repformer_update_g1_has_drrd": true, - "repformer_update_g1_has_attn": true, - "repformer_update_g2_has_g1g1": true, - "repformer_update_g2_has_attn": true, - "repformer_attn2_has_gate": true, + "repinit": { + "rcut": 6.0, + "rcut_smth": 2.0, + "nsel": 30, + "neuron": [ + 2, + 4, + 8 + ], + "axis_neuron": 4, + "activation_function": "tanh" + + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "nlayers": 12, + "g1_dim": 8, + "g2_dim": 5, + "attn2_hidden": 3, + "attn2_nhead": 1, + "attn1_hidden": 5, + "attn1_nhead": 1, + "axis_neuron": 4, + "update_h2": false, + "update_g1_has_conv": true, + "update_g1_has_grrg": true, + "update_g1_has_drrd": true, + "update_g1_has_attn": true, + "update_g2_has_g1g1": true, + "update_g2_has_attn": true, + "attn2_has_gate": true + }, "add_tebd_to_repinit_out": false }, "fitting_net": { diff --git a/source/tests/pt/model/test_autodiff.py b/source/tests/pt/model/test_autodiff.py index 91fc3cabf6..268591948a 100644 --- a/source/tests/pt/model/test_autodiff.py +++ b/source/tests/pt/model/test_autodiff.py @@ -197,13 +197,6 @@ def setUp(self): class TestEnergyModelDPA2Force(unittest.TestCase, ForceTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -211,13 +204,6 @@ def setUp(self): class TestEnergyModelDPAUniVirial(unittest.TestCase, VirialTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_dpa2.py b/source/tests/pt/model/test_dpa2.py index 8e957202da..0bff470f76 100644 --- a/source/tests/pt/model/test_dpa2.py +++ b/source/tests/pt/model/test_dpa2.py @@ -6,6 +6,10 @@ import torch from deepmd.dpmodel.descriptor.dpa2 import DescrptDPA2 as DPDescrptDPA2 +from deepmd.dpmodel.descriptor.dpa2 import ( + RepformerArgs, + RepinitArgs, +) from deepmd.pt.model.descriptor.dpa2 import ( DescrptDPA2, ) @@ -84,37 +88,42 @@ def test_consistency( if prec == "float64": atol = 1e-11 # marginal GPU test cases... + repinit = RepinitArgs( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + nsel=self.sel_mix, + tebd_input_mode=riti, + set_davg_zero=riz, + ) + repformer = RepformerArgs( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth, + nsel=nnei // 2, + nlayers=3, + g1_dim=20, + g2_dim=10, + axis_neuron=4, + update_g1_has_conv=rp1c, + update_g1_has_drrd=rp1d, + update_g1_has_grrg=rp1g, + update_g1_has_attn=rp1a, + update_g2_has_g1g1=rp2g, + update_g2_has_attn=rp2a, + update_h2=rph, + attn1_hidden=20, + attn1_nhead=2, + attn2_hidden=10, + attn2_nhead=2, + attn2_has_gate=rp2gate, + update_style=rus, + set_davg_zero=rpz, + ) + # dpa2 new impl dd0 = DescrptDPA2( self.nt, - repinit_rcut=self.rcut, - repinit_rcut_smth=self.rcut_smth, - repinit_nsel=self.sel_mix, - repformer_rcut=self.rcut / 2, - repformer_rcut_smth=self.rcut_smth, - repformer_nsel=nnei // 2, - # kwargs for repinit - repinit_tebd_input_mode=riti, - repinit_set_davg_zero=riz, - # kwargs for repformer - repformer_nlayers=3, - repformer_g1_dim=20, - repformer_g2_dim=10, - repformer_axis_neuron=4, - repformer_update_g1_has_conv=rp1c, - repformer_update_g1_has_drrd=rp1d, - repformer_update_g1_has_grrg=rp1g, - repformer_update_g1_has_attn=rp1a, - repformer_update_g2_has_g1g1=rp2g, - repformer_update_g2_has_attn=rp2a, - repformer_update_h2=rph, - repformer_attn1_hidden=20, - repformer_attn1_nhead=2, - repformer_attn2_hidden=10, - repformer_attn2_nhead=2, - repformer_attn2_has_gate=rp2gate, - repformer_update_style=rus, - repformer_set_davg_zero=rpz, + repinit=repinit, + repformer=repformer, # kwargs for descriptor smooth=sm, exclude_types=[], @@ -162,34 +171,8 @@ def test_consistency( if prec == "float64" and rus == "res_avg": dd3 = DescrptDPA2( self.nt, - repinit_rcut=self.rcut, - repinit_rcut_smth=self.rcut_smth, - repinit_nsel=self.sel_mix, - repformer_rcut=self.rcut / 2, - repformer_rcut_smth=self.rcut_smth, - repformer_nsel=nnei // 2, - # kwargs for repinit - repinit_tebd_input_mode=riti, - repinit_set_davg_zero=riz, - # kwargs for repformer - repformer_nlayers=3, - repformer_g1_dim=20, - repformer_g2_dim=10, - repformer_axis_neuron=4, - repformer_update_g1_has_conv=rp1c, - repformer_update_g1_has_drrd=rp1d, - repformer_update_g1_has_grrg=rp1g, - repformer_update_g1_has_attn=rp1a, - repformer_update_g2_has_g1g1=rp2g, - repformer_update_g2_has_attn=rp2a, - repformer_update_h2=rph, - repformer_attn1_hidden=20, - repformer_attn1_nhead=2, - repformer_attn2_hidden=10, - repformer_attn2_nhead=2, - repformer_attn2_has_gate=rp2gate, - repformer_update_style="res_avg", - repformer_set_davg_zero=rpz, + repinit=repinit, + repformer=repformer, # kwargs for descriptor smooth=sm, exclude_types=[], @@ -290,37 +273,42 @@ def test_jit( dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) + repinit = RepinitArgs( + rcut=self.rcut, + rcut_smth=self.rcut_smth, + nsel=self.sel_mix, + tebd_input_mode=riti, + set_davg_zero=riz, + ) + repformer = RepformerArgs( + rcut=self.rcut / 2, + rcut_smth=self.rcut_smth, + nsel=nnei // 2, + nlayers=3, + g1_dim=20, + g2_dim=10, + axis_neuron=4, + update_g1_has_conv=rp1c, + update_g1_has_drrd=rp1d, + update_g1_has_grrg=rp1g, + update_g1_has_attn=rp1a, + update_g2_has_g1g1=rp2g, + update_g2_has_attn=rp2a, + update_h2=rph, + attn1_hidden=20, + attn1_nhead=2, + attn2_hidden=10, + attn2_nhead=2, + attn2_has_gate=rp2gate, + update_style=rus, + set_davg_zero=rpz, + ) + # dpa2 new impl dd0 = DescrptDPA2( self.nt, - repinit_rcut=self.rcut, - repinit_rcut_smth=self.rcut_smth, - repinit_nsel=self.sel_mix, - repformer_rcut=self.rcut / 2, - repformer_rcut_smth=self.rcut_smth, - repformer_nsel=nnei // 2, - # kwargs for repinit - repinit_tebd_input_mode=riti, - repinit_set_davg_zero=riz, - # kwargs for repformer - repformer_nlayers=3, - repformer_g1_dim=20, - repformer_g2_dim=10, - repformer_axis_neuron=4, - repformer_update_g1_has_conv=rp1c, - repformer_update_g1_has_drrd=rp1d, - repformer_update_g1_has_grrg=rp1g, - repformer_update_g1_has_attn=rp1a, - repformer_update_g2_has_g1g1=rp2g, - repformer_update_g2_has_attn=rp2a, - repformer_update_h2=rph, - repformer_attn1_hidden=20, - repformer_attn1_nhead=2, - repformer_attn2_hidden=10, - repformer_attn2_nhead=2, - repformer_attn2_has_gate=rp2gate, - repformer_update_style=rus, - repformer_set_davg_zero=rpz, + repinit=repinit, + repformer=repformer, # kwargs for descriptor smooth=sm, exclude_types=[], diff --git a/source/tests/pt/model/test_ener_spin_model.py b/source/tests/pt/model/test_ener_spin_model.py index 2bd5c22aaf..8b435717db 100644 --- a/source/tests/pt/model/test_ener_spin_model.py +++ b/source/tests/pt/model/test_ener_spin_model.py @@ -408,8 +408,8 @@ def setUp(self): SpinTest.setUp(self) model_params = copy.deepcopy(model_spin) model_params["descriptor"] = copy.deepcopy(model_dpa2["descriptor"]) - self.rcut = model_params["descriptor"]["repinit_rcut"] - self.nsel = model_params["descriptor"]["repinit_nsel"] + self.rcut = model_params["descriptor"]["repinit"]["rcut"] + self.nsel = model_params["descriptor"]["repinit"]["nsel"] self.type_map = model_params["type_map"] # not implement serialize and deserialize self.serial_test = False diff --git a/source/tests/pt/model/test_forward_lower.py b/source/tests/pt/model/test_forward_lower.py index c43c4baea2..5d2bfe599e 100644 --- a/source/tests/pt/model/test_forward_lower.py +++ b/source/tests/pt/model/test_forward_lower.py @@ -143,13 +143,6 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, ForwardLowerTest): def setUp(self): self.prec = 1e-10 - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_jit.py b/source/tests/pt/model/test_jit.py index 41a5902a5a..248ccf9173 100644 --- a/source/tests/pt/model/test_jit.py +++ b/source/tests/pt/model/test_jit.py @@ -104,15 +104,6 @@ def setUp(self): self.config["training"]["training_data"]["systems"] = data_file self.config["training"]["validation_data"]["systems"] = data_file self.config["model"] = deepcopy(model_dpa2) - # self.config["model"]["descriptor"]["rcut"] = self.config["model"]["descriptor"][ - # "repinit_rcut" - # ] - # self.config["model"]["descriptor"]["rcut_smth"] = self.config["model"][ - # "descriptor" - # ]["repinit_rcut_smth"] - # self.config["model"]["descriptor"]["sel"] = self.config["model"]["descriptor"][ - # "repinit_nsel" - # ] self.config["training"]["numb_steps"] = 10 self.config["training"]["save_freq"] = 10 diff --git a/source/tests/pt/model/test_permutation.py b/source/tests/pt/model/test_permutation.py index b2bfcb8a4d..f8d75fc99f 100644 --- a/source/tests/pt/model/test_permutation.py +++ b/source/tests/pt/model/test_permutation.py @@ -122,31 +122,35 @@ "type_map": ["O", "H", "B"], "descriptor": { "type": "dpa2", - "repinit_rcut": 6.0, - "repinit_rcut_smth": 2.0, - "repinit_nsel": 100, - "repformer_rcut": 4.0, - "repformer_rcut_smth": 0.5, - "repformer_nsel": 40, - "repinit_neuron": [2, 4, 8], - "repinit_axis_neuron": 4, - "repinit_activation_function": "tanh", - "repformer_nlayers": 12, - "repformer_g1_dim": 8, - "repformer_g2_dim": 5, - "repformer_attn2_hidden": 3, - "repformer_attn2_nhead": 1, - "repformer_attn1_hidden": 5, - "repformer_attn1_nhead": 1, - "repformer_axis_neuron": 4, - "repformer_update_h2": False, - "repformer_update_g1_has_conv": True, - "repformer_update_g1_has_grrg": True, - "repformer_update_g1_has_drrd": True, - "repformer_update_g1_has_attn": True, - "repformer_update_g2_has_g1g1": True, - "repformer_update_g2_has_attn": True, - "repformer_attn2_has_gate": True, + "repinit": { + "rcut": 6.0, + "rcut_smth": 2.0, + "nsel": 100, + "neuron": [2, 4, 8], + "axis_neuron": 4, + "activation_function": "tanh", + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 40, + "nlayers": 12, + "g1_dim": 8, + "g2_dim": 5, + "attn2_hidden": 3, + "attn2_nhead": 1, + "attn1_hidden": 5, + "attn1_nhead": 1, + "axis_neuron": 4, + "update_h2": False, + "update_g1_has_conv": True, + "update_g1_has_grrg": True, + "update_g1_has_drrd": True, + "update_g1_has_attn": True, + "update_g2_has_g1g1": True, + "update_g2_has_attn": True, + "attn2_has_gate": True, + }, "add_tebd_to_repinit_out": False, }, "fitting_net": { @@ -207,31 +211,35 @@ }, { "type": "dpa2", - "repinit_rcut": 6.0, - "repinit_rcut_smth": 2.0, - "repinit_nsel": 30, - "repformer_rcut": 4.0, - "repformer_rcut_smth": 0.5, - "repformer_nsel": 10, - "repinit_neuron": [2, 4, 8], - "repinit_axis_neuron": 4, - "repinit_activation_function": "tanh", - "repformer_nlayers": 12, - "repformer_g1_dim": 8, - "repformer_g2_dim": 5, - "repformer_attn2_hidden": 3, - "repformer_attn2_nhead": 1, - "repformer_attn1_hidden": 5, - "repformer_attn1_nhead": 1, - "repformer_axis_neuron": 4, - "repformer_update_h2": False, - "repformer_update_g1_has_conv": True, - "repformer_update_g1_has_grrg": True, - "repformer_update_g1_has_drrd": True, - "repformer_update_g1_has_attn": True, - "repformer_update_g2_has_g1g1": True, - "repformer_update_g2_has_attn": True, - "repformer_attn2_has_gate": True, + "repinit": { + "rcut": 6.0, + "rcut_smth": 2.0, + "nsel": 30, + "neuron": [2, 4, 8], + "axis_neuron": 4, + "activation_function": "tanh", + }, + "repformer": { + "rcut": 4.0, + "rcut_smth": 0.5, + "nsel": 10, + "nlayers": 12, + "g1_dim": 8, + "g2_dim": 5, + "attn2_hidden": 3, + "attn2_nhead": 1, + "attn1_hidden": 5, + "attn1_nhead": 1, + "axis_neuron": 4, + "update_h2": False, + "update_g1_has_conv": True, + "update_g1_has_grrg": True, + "update_g1_has_drrd": True, + "update_g1_has_attn": True, + "update_g2_has_g1g1": True, + "update_g2_has_attn": True, + "attn2_has_gate": True, + }, "add_tebd_to_repinit_out": False, }, ], @@ -319,13 +327,6 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, PermutationTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -333,13 +334,6 @@ def setUp(self): class TestForceModelDPA2(unittest.TestCase, PermutationTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True diff --git a/source/tests/pt/model/test_permutation_denoise.py b/source/tests/pt/model/test_permutation_denoise.py index 3b6be0c495..fad57a807b 100644 --- a/source/tests/pt/model/test_permutation_denoise.py +++ b/source/tests/pt/model/test_permutation_denoise.py @@ -76,13 +76,6 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, PermutationDenoiseTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model( diff --git a/source/tests/pt/model/test_rot.py b/source/tests/pt/model/test_rot.py index cbf09ecf40..ec8bdebe74 100644 --- a/source/tests/pt/model/test_rot.py +++ b/source/tests/pt/model/test_rot.py @@ -156,13 +156,6 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, RotTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -170,13 +163,6 @@ def setUp(self): class TestForceModelDPA2(unittest.TestCase, RotTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True diff --git a/source/tests/pt/model/test_rot_denoise.py b/source/tests/pt/model/test_rot_denoise.py index e4ae02f630..6b2e4a8176 100644 --- a/source/tests/pt/model/test_rot_denoise.py +++ b/source/tests/pt/model/test_rot_denoise.py @@ -107,13 +107,6 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, RotDenoiseTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_smooth.py b/source/tests/pt/model/test_smooth.py index 1a75caebdc..975c8d5387 100644 --- a/source/tests/pt/model/test_smooth.py +++ b/source/tests/pt/model/test_smooth.py @@ -188,19 +188,8 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, SmoothTest): def setUp(self): model_params = copy.deepcopy(model_dpa2) - model_params["descriptor"]["repinit_rcut"] = 8 - model_params["descriptor"]["repinit_rcut_smth"] = 3.5 - model_params_sample = copy.deepcopy(model_params) - ####################################################### - # dirty hack here! the interface of dataload should be - # redesigned to support specifying rcut and sel - ####################################################### - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] + model_params["descriptor"]["repinit"]["rcut"] = 8 + model_params["descriptor"]["repinit"]["rcut_smth"] = 3.5 self.type_split = True self.model = get_model(model_params).to(env.DEVICE) self.epsilon, self.aprec = 1e-5, 1e-4 @@ -210,13 +199,6 @@ class TestEnergyModelDPA2_1(unittest.TestCase, SmoothTest): def setUp(self): model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "ener" - model_params_sample = copy.deepcopy(model_params) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] self.type_split = True self.test_virial = False self.model = get_model(model_params).to(env.DEVICE) @@ -227,13 +209,6 @@ class TestEnergyModelDPA2_2(unittest.TestCase, SmoothTest): def setUp(self): model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "ener" - model_params_sample = copy.deepcopy(model_params) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] self.type_split = True self.test_virial = False self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_smooth_denoise.py b/source/tests/pt/model/test_smooth_denoise.py index 777d288f3c..200069daa1 100644 --- a/source/tests/pt/model/test_smooth_denoise.py +++ b/source/tests/pt/model/test_smooth_denoise.py @@ -98,13 +98,6 @@ def compare(ret0, ret1): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, SmoothDenoiseTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) model_params["descriptor"]["sel"] = 8 model_params["descriptor"]["rcut_smth"] = 3.5 @@ -118,13 +111,6 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2_1(unittest.TestCase, SmoothDenoiseTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) # model_params["descriptor"]["combine_grrg"] = True self.type_split = True diff --git a/source/tests/pt/model/test_trans.py b/source/tests/pt/model/test_trans.py index a0aeefd6b3..d2d790d9e3 100644 --- a/source/tests/pt/model/test_trans.py +++ b/source/tests/pt/model/test_trans.py @@ -100,13 +100,6 @@ def setUp(self): class TestEnergyModelDPA2(unittest.TestCase, TransTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) @@ -114,13 +107,6 @@ def setUp(self): class TestForceModelDPA2(unittest.TestCase, TransTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) model_params["fitting_net"]["type"] = "direct_force_ener" self.type_split = True diff --git a/source/tests/pt/model/test_trans_denoise.py b/source/tests/pt/model/test_trans_denoise.py index 9ba93a244a..3b9a950f4b 100644 --- a/source/tests/pt/model/test_trans_denoise.py +++ b/source/tests/pt/model/test_trans_denoise.py @@ -66,13 +66,6 @@ def setUp(self): @unittest.skip("support of the denoise is temporally disabled") class TestDenoiseModelDPA2(unittest.TestCase, TransDenoiseTest): def setUp(self): - model_params_sample = copy.deepcopy(model_dpa2) - model_params_sample["descriptor"]["rcut"] = model_params_sample["descriptor"][ - "repinit_rcut" - ] - model_params_sample["descriptor"]["sel"] = model_params_sample["descriptor"][ - "repinit_nsel" - ] model_params = copy.deepcopy(model_dpa2) self.type_split = True self.model = get_model(model_params).to(env.DEVICE) diff --git a/source/tests/pt/model/test_unused_params.py b/source/tests/pt/model/test_unused_params.py index 730eb00a0b..b4c69ab69f 100644 --- a/source/tests/pt/model/test_unused_params.py +++ b/source/tests/pt/model/test_unused_params.py @@ -41,15 +41,15 @@ def test_unused(self): # skip the case g2 is not envolved continue model = copy.deepcopy(model_dpa2) - model["descriptor"]["repformer_nlayers"] = 2 + model["descriptor"]["repformer"]["nlayers"] = 2 # model["descriptor"]["combine_grrg"] = cmbg2 - model["descriptor"]["repformer_update_g1_has_conv"] = conv - model["descriptor"]["repformer_update_g1_has_drrd"] = drrd - model["descriptor"]["repformer_update_g1_has_grrg"] = grrg - model["descriptor"]["repformer_update_g1_has_attn"] = attn1 - model["descriptor"]["repformer_update_g2_has_g1g1"] = g1g1 - model["descriptor"]["repformer_update_g2_has_attn"] = attn2 - model["descriptor"]["repformer_update_h2"] = h2 + model["descriptor"]["repformer"]["update_g1_has_conv"] = conv + model["descriptor"]["repformer"]["update_g1_has_drrd"] = drrd + model["descriptor"]["repformer"]["update_g1_has_grrg"] = grrg + model["descriptor"]["repformer"]["update_g1_has_attn"] = attn1 + model["descriptor"]["repformer"]["update_g2_has_g1g1"] = g1g1 + model["descriptor"]["repformer"]["update_g2_has_attn"] = attn2 + model["descriptor"]["repformer"]["update_h2"] = h2 model["fitting_net"]["neuron"] = [12, 12, 12] self._test_unused(model)