diff --git a/deepmd/pd/loss/__init__.py b/deepmd/pd/loss/__init__.py index 0e978b95c2..5e411b0f7f 100644 --- a/deepmd/pd/loss/__init__.py +++ b/deepmd/pd/loss/__init__.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from .ener import ( + EnergyHessianStdLoss, EnergyStdLoss, ) from .loss import ( @@ -7,6 +8,7 @@ ) __all__ = [ + "EnergyHessianStdLoss", "EnergyStdLoss", "TaskLoss", ] diff --git a/deepmd/pd/loss/ener.py b/deepmd/pd/loss/ener.py index 1560e4d59d..09ec5ff49e 100644 --- a/deepmd/pd/loss/ener.py +++ b/deepmd/pd/loss/ener.py @@ -56,7 +56,7 @@ def __init__( use_huber=False, huber_delta=0.01, **kwargs, - ): + ) -> None: r"""Construct a layer to compute loss on energy, force and virial. Parameters @@ -287,9 +287,9 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): rmse_f.detach(), find_force ) else: - l1_force_loss = F.l1_loss(force_label, force_pred, reduction="none") + l1_force_loss = F.l1_loss(force_label, force_pred, reduction="mean") more_loss["mae_f"] = self.display_if_exist( - l1_force_loss.mean().detach(), find_force + l1_force_loss.detach(), find_force ) l1_force_loss = l1_force_loss.sum(-1).mean(-1).sum() loss += (pref_f * l1_force_loss).to(GLOBAL_PD_FLOAT_PRECISION) @@ -324,20 +324,19 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): drdq_reshape = drdq.reshape( [-1, natoms * 3, self.numb_generalized_coord] ) + gen_force_label = paddle.einsum( + "bij,bi->bj", drdq_reshape, force_label_reshape_nframes + ) + # gen_force_label = ( + # drdq_reshape * force_label_reshape_nframes.unsqueeze(-1) + # ).sum([-2]) - # gen_force_label = paddle.einsum( - # "bij,bi->bj", drdq_reshape, force_label_reshape_nframes - # ) - gen_force_label = ( - drdq_reshape * force_label_reshape_nframes.unsqueeze(-1) - ).sum([-2]) - - # gen_force = paddle.einsum( - # "bij,bi->bj", drdq_reshape, force_reshape_nframes - # ) - gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum( - [-2] + gen_force = paddle.einsum( + "bij,bi->bj", drdq_reshape, force_reshape_nframes ) + # gen_force = (drdq_reshape * force_reshape_nframes.unsqueeze(-1)).sum( + # [-2] + # ) diff_gen_force = gen_force_label - gen_force l2_gen_force_loss = paddle.square(diff_gen_force).mean() @@ -534,3 +533,75 @@ def deserialize(cls, data: dict) -> "TaskLoss": check_version_compatibility(data.pop("@version"), 2, 1) data.pop("@class") return cls(**data) + + +class EnergyHessianStdLoss(EnergyStdLoss): + def __init__( + self, + start_pref_h=0.0, + limit_pref_h=0.0, + **kwargs, + ): + r"""Enable the layer to compute loss on hessian. + + Parameters + ---------- + start_pref_h : float + The prefactor of hessian loss at the start of the training. + limit_pref_h : float + The prefactor of hessian loss at the end of the training. + **kwargs + Other keyword arguments. + """ + super().__init__(**kwargs) + self.has_h = (start_pref_h != 0.0 and limit_pref_h != 0.0) or self.inference + + self.start_pref_h = start_pref_h + self.limit_pref_h = limit_pref_h + + def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + model_pred, loss, more_loss = super().forward( + input_dict, model, label, natoms, learning_rate, mae=mae + ) + coef = learning_rate / self.starter_learning_rate + pref_h = self.limit_pref_h + (self.start_pref_h - self.limit_pref_h) * coef + + if self.has_h and "hessian" in model_pred and "hessian" in label: + find_hessian = label.get("find_hessian", 0.0) + pref_h = pref_h * find_hessian + diff_h = label["hessian"].reshape( + [-1], + ) - model_pred["hessian"].reshape( + [-1], + ) + l2_hessian_loss = paddle.mean(paddle.square(diff_h)) + if not self.inference: + more_loss["l2_hessian_loss"] = self.display_if_exist( + l2_hessian_loss.detach(), find_hessian + ) + loss += pref_h * l2_hessian_loss + rmse_h = l2_hessian_loss.sqrt() + more_loss["rmse_h"] = self.display_if_exist(rmse_h.detach(), find_hessian) + if mae: + mae_h = paddle.mean(paddle.abs(diff_h)) + more_loss["mae_h"] = self.display_if_exist(mae_h.detach(), find_hessian) + + if not self.inference: + more_loss["rmse"] = paddle.sqrt(loss.detach()) + return model_pred, loss, more_loss + + @property + def label_requirement(self) -> list[DataRequirementItem]: + """Add hessian label requirement needed for this loss calculation.""" + label_requirement = super().label_requirement + if self.has_h: + label_requirement.append( + DataRequirementItem( + "hessian", + ndof=1, # 9=3*3 --> 3N*3N=ndof*natoms*natoms + atomic=True, + must=False, + high_prec=False, + ) + ) + return label_requirement diff --git a/deepmd/pd/model/descriptor/__init__.py b/deepmd/pd/model/descriptor/__init__.py index caa1437fc7..28a99a8230 100644 --- a/deepmd/pd/model/descriptor/__init__.py +++ b/deepmd/pd/model/descriptor/__init__.py @@ -12,6 +12,9 @@ from .dpa2 import ( DescrptDPA2, ) +from .dpa3 import ( + DescrptDPA3, +) from .env_mat import ( prod_env_mat, ) @@ -39,6 +42,7 @@ "DescrptBlockSeTTebd", "DescrptDPA1", "DescrptDPA2", + "DescrptDPA3", "DescrptSeA", "DescrptSeAttenV2", "DescrptSeTTebd", diff --git a/deepmd/pd/model/descriptor/descriptor.py b/deepmd/pd/model/descriptor/descriptor.py index c6d2d1efe6..3050b7dca3 100644 --- a/deepmd/pd/model/descriptor/descriptor.py +++ b/deepmd/pd/model/descriptor/descriptor.py @@ -6,12 +6,16 @@ ) from typing import ( Callable, + NoReturn, Optional, Union, ) import paddle +from deepmd.pd.model.network.network import ( + TypeEmbedNet, +) from deepmd.pd.utils import ( env, ) @@ -99,7 +103,7 @@ def compute_input_stats( self, merged: Union[Callable[[], list[dict]], list[dict]], path: Optional[DPPath] = None, - ): + ) -> NoReturn: """ Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. @@ -122,7 +126,7 @@ def get_stats(self) -> dict[str, StatItem]: """Get the statistics of the descriptor.""" raise NotImplementedError - def share_params(self, base_class, shared_level, resume=False): + def share_params(self, base_class, shared_level, resume=False) -> None: """ Share the parameters of self to the base_class with shared_level during multitask training. If not start from checkpoint (resume is False), @@ -134,7 +138,10 @@ def share_params(self, base_class, shared_level, resume=False): if shared_level == 0: # link buffers if hasattr(self, "mean"): - if not resume: + if not resume and ( + not getattr(self, "set_stddev_constant", False) + or not getattr(self, "set_davg_zero", False) + ): # in case of change params during resume base_env = EnvMatStatSe(base_class) base_env.stats = base_class.stats @@ -172,6 +179,7 @@ def forward( extended_atype: paddle.Tensor, extended_atype_embd: Optional[paddle.Tensor] = None, mapping: Optional[paddle.Tensor] = None, + type_embedding: Optional[paddle.Tensor] = None, ): """Calculate DescriptorBlock.""" pass @@ -185,7 +193,15 @@ def need_sorted_nlist_for_lower(self) -> bool: """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" -def extend_descrpt_stat(des, type_map, des_with_stat=None): +def make_default_type_embedding( + ntypes, +): + aux = {} + aux["tebd_dim"] = 8 + return TypeEmbedNet(ntypes, aux["tebd_dim"]), aux + + +def extend_descrpt_stat(des, type_map, des_with_stat=None) -> None: r""" Extend the statistics of a descriptor block with types from newly provided `type_map`. diff --git a/deepmd/pd/model/descriptor/dpa1.py b/deepmd/pd/model/descriptor/dpa1.py index 42cba44eb4..6af494d5b4 100644 --- a/deepmd/pd/model/descriptor/dpa1.py +++ b/deepmd/pd/model/descriptor/dpa1.py @@ -584,7 +584,7 @@ def enable_compression( The overflow check frequency """ # do some checks before the mocel compression process - raise NotImplementedError("Model compression is not supported in paddle yet.") + raise ValueError("Compression is already enabled.") def forward( self, diff --git a/deepmd/pd/model/descriptor/dpa2.py b/deepmd/pd/model/descriptor/dpa2.py index 9cd3d77f89..0e3b24397f 100644 --- a/deepmd/pd/model/descriptor/dpa2.py +++ b/deepmd/pd/model/descriptor/dpa2.py @@ -408,7 +408,9 @@ def share_params(self, base_class, shared_level, resume=False) -> None: # shared_level: 1 # share all parameters in type_embedding elif shared_level == 1: - self._modules["type_embedding"] = base_class._modules["type_embedding"] + self._sub_layers["type_embedding"] = base_class._sub_layers[ + "type_embedding" + ] # Other shared levels else: raise NotImplementedError @@ -899,4 +901,4 @@ def enable_compression( The overflow check frequency """ # do some checks before the mocel compression process - raise NotImplementedError("enable_compression is not implemented yet") + raise ValueError("Compression is already enabled.") diff --git a/deepmd/pd/model/descriptor/dpa3.py b/deepmd/pd/model/descriptor/dpa3.py new file mode 100644 index 0000000000..8538d7068c --- /dev/null +++ b/deepmd/pd/model/descriptor/dpa3.py @@ -0,0 +1,571 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Callable, + Optional, + Union, +) + +import paddle + +from deepmd.dpmodel.descriptor.dpa3 import ( + RepFlowArgs, +) +from deepmd.dpmodel.utils import EnvMat as DPEnvMat +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.pd.model.network.mlp import ( + MLPLayer, +) +from deepmd.pd.model.network.network import ( + TypeEmbedNet, + TypeEmbedNetConsistent, +) +from deepmd.pd.utils import ( + env, +) +from deepmd.pd.utils.env import ( + PRECISION_DICT, +) +from deepmd.pd.utils.update_sel import ( + UpdateSel, +) +from deepmd.pd.utils.utils import ( + to_numpy_array, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) +from deepmd.utils.finetune import ( + get_index_between_two_maps, + map_pair_exclude_types, +) +from deepmd.utils.path import ( + DPPath, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + +from .base_descriptor import ( + BaseDescriptor, +) +from .descriptor import ( + extend_descrpt_stat, +) +from .repflow_layer import ( + RepFlowLayer, +) +from .repflows import ( + DescrptBlockRepflows, +) + + +@BaseDescriptor.register("dpa3") +class DescrptDPA3(BaseDescriptor, paddle.nn.Layer): + r"""The DPA-3 descriptor. + + Parameters + ---------- + repflow : Union[RepFlowArgs, dict] + The arguments used to initialize the repflow block, see docstr in `RepFlowArgs` for details information. + concat_output_tebd : bool, optional + Whether to concat type embedding at the output of the descriptor. + activation_function : str, optional + The activation function in the embedding net. + precision : str, optional + The precision of the embedding net parameters. + exclude_types : list[list[int]], optional + The excluded pairs of types which have no interaction with each other. + For example, `[[0, 1]]` means no interaction between type 0 and type 1. + env_protection : float, optional + Protection parameter to prevent division by zero errors during environment matrix calculations. + For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. + trainable : bool, optional + If the parameters are trainable. + seed : int, optional + Random seed for parameter initialization. + use_econf_tebd : bool, Optional + Whether to use electronic configuration type embedding. + use_tebd_bias : bool, Optional + Whether to use bias in the type embedding layer. + type_map : list[str], Optional + A list of strings. Give the name to each type of atoms. + """ + + def __init__( + self, + ntypes: int, + # args for repflow + repflow: Union[RepFlowArgs, dict], + # kwargs for descriptor + concat_output_tebd: bool = False, + activation_function: str = "silu", + precision: str = "float64", + exclude_types: list[tuple[int, int]] = [], + env_protection: float = 0.0, + trainable: bool = True, + seed: Optional[Union[int, list[int]]] = None, + use_econf_tebd: bool = False, + use_tebd_bias: bool = False, + type_map: Optional[list[str]] = None, + ) -> None: + super().__init__() + + 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.repflow_args = init_subclass_params(repflow, RepFlowArgs) + self.activation_function = activation_function + + self.repflows = DescrptBlockRepflows( + self.repflow_args.e_rcut, + self.repflow_args.e_rcut_smth, + self.repflow_args.e_sel, + self.repflow_args.a_rcut, + self.repflow_args.a_rcut_smth, + self.repflow_args.a_sel, + ntypes, + nlayers=self.repflow_args.nlayers, + n_dim=self.repflow_args.n_dim, + e_dim=self.repflow_args.e_dim, + a_dim=self.repflow_args.a_dim, + a_compress_rate=self.repflow_args.a_compress_rate, + a_compress_e_rate=self.repflow_args.a_compress_e_rate, + a_compress_use_split=self.repflow_args.a_compress_use_split, + n_multi_edge_message=self.repflow_args.n_multi_edge_message, + axis_neuron=self.repflow_args.axis_neuron, + update_angle=self.repflow_args.update_angle, + activation_function=self.activation_function, + update_style=self.repflow_args.update_style, + update_residual=self.repflow_args.update_residual, + update_residual_init=self.repflow_args.update_residual_init, + fix_stat_std=self.repflow_args.fix_stat_std, + optim_update=self.repflow_args.optim_update, + smooth_edge_update=self.repflow_args.smooth_edge_update, + exclude_types=exclude_types, + env_protection=env_protection, + precision=precision, + seed=child_seed(seed, 1), + ) + + self.use_econf_tebd = use_econf_tebd + self.use_tebd_bias = use_tebd_bias + self.type_map = type_map + self.tebd_dim = self.repflow_args.n_dim + self.type_embedding = TypeEmbedNet( + ntypes, + self.tebd_dim, + precision=precision, + seed=child_seed(seed, 2), + use_econf_tebd=self.use_econf_tebd, + use_tebd_bias=use_tebd_bias, + type_map=type_map, + ) + self.concat_output_tebd = concat_output_tebd + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.exclude_types = exclude_types + self.env_protection = env_protection + self.trainable = trainable + + assert self.repflows.e_rcut >= self.repflows.a_rcut, ( + f"Edge radial cutoff (e_rcut: {self.repflows.e_rcut}) " + f"must be greater than or equal to angular cutoff (a_rcut: {self.repflows.a_rcut})!" + ) + assert self.repflows.e_sel >= self.repflows.a_sel, ( + f"Edge sel number (e_sel: {self.repflows.e_sel}) " + f"must be greater than or equal to angular sel (a_sel: {self.repflows.a_sel})!" + ) + + self.rcut = self.repflows.get_rcut() + self.rcut_smth = self.repflows.get_rcut_smth() + self.sel = self.repflows.get_sel() + self.ntypes = ntypes + + # set trainable + for param in self.parameters(): + param.requires_grad = trainable + self.compress = False + + def get_rcut(self) -> float: + """Returns the cut-off radius.""" + return self.rcut + + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.rcut_smth + + def get_nsel(self) -> int: + """Returns the number of selected atoms in the cut-off radius.""" + return sum(self.sel) + + def get_sel(self) -> list[int]: + """Returns the number of selected atoms for each type.""" + return self.sel + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.ntypes + + def get_type_map(self) -> list[str]: + """Get the name to each type of atoms.""" + return self.type_map + + def get_dim_out(self) -> int: + """Returns the output dimension of this descriptor.""" + ret = self.repflows.dim_out + if self.concat_output_tebd: + ret += self.tebd_dim + return ret + + def get_dim_emb(self) -> int: + """Returns the embedding dimension of this descriptor.""" + return self.repflows.dim_emb + + def mixed_types(self) -> bool: + """If true, the descriptor + 1. assumes total number of atoms aligned across frames; + 2. requires a neighbor list that does not distinguish different atomic types. + + If false, the descriptor + 1. assumes total number of atoms of each atom type aligned across frames; + 2. requires a neighbor list that distinguishes different atomic types. + + """ + return True + + def has_message_passing(self) -> bool: + """Returns whether the descriptor has message passing.""" + return self.repflows.has_message_passing() + + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor needs sorted nlist when using `forward_lower`.""" + return True + + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.repflows.get_env_protection() + + def share_params(self, base_class, shared_level, resume=False) -> None: + """ + Share the parameters of self to the base_class with shared_level during multitask training. + If not start from checkpoint (resume is False), + some separated parameters (e.g. mean and stddev) will be re-calculated across different classes. + """ + assert self.__class__ == base_class.__class__, ( + "Only descriptors of the same type can share params!" + ) + # For DPA3 descriptors, the user-defined share-level + # shared_level: 0 + # share all parameters in type_embedding, repflow + if shared_level == 0: + self._sub_layers["type_embedding"] = base_class._sub_layers[ + "type_embedding" + ] + self.repflows.share_params(base_class.repflows, 0, resume=resume) + # shared_level: 1 + # share all parameters in type_embedding + elif shared_level == 1: + self._sub_layers["type_embedding"] = base_class._sub_layers[ + "type_embedding" + ] + # Other shared levels + else: + raise NotImplementedError + + def change_type_map( + self, type_map: list[str], model_with_new_type_stat=None + ) -> None: + """Change the type related params to new ones, according to `type_map` and the original one in the model. + If there are new types in `type_map`, statistics will be updated accordingly to `model_with_new_type_stat` for these new types. + """ + assert self.type_map is not None, ( + "'type_map' must be defined when performing type changing!" + ) + remap_index, has_new_type = get_index_between_two_maps(self.type_map, type_map) + self.type_map = type_map + self.type_embedding.change_type_map(type_map=type_map) + self.exclude_types = map_pair_exclude_types(self.exclude_types, remap_index) + self.ntypes = len(type_map) + repflow = self.repflows + if has_new_type: + # the avg and std of new types need to be updated + extend_descrpt_stat( + repflow, + type_map, + des_with_stat=model_with_new_type_stat.repflows + if model_with_new_type_stat is not None + else None, + ) + repflow.ntypes = self.ntypes + repflow.reinit_exclude(self.exclude_types) + repflow["davg"] = repflow["davg"][remap_index] + repflow["dstd"] = repflow["dstd"][remap_index] + + @property + def dim_out(self): + return self.get_dim_out() + + @property + def dim_emb(self): + """Returns the embedding dimension g2.""" + return self.get_dim_emb() + + def compute_input_stats( + self, + merged: Union[Callable[[], list[dict]], list[dict]], + path: Optional[DPPath] = None, + ) -> None: + """ + Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. + + Parameters + ---------- + merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + path : Optional[DPPath] + The path to the stat file. + + """ + descrpt_list = [self.repflows] + for ii, descrpt in enumerate(descrpt_list): + descrpt.compute_input_stats(merged, path) + + def set_stat_mean_and_stddev( + self, + mean: list[paddle.Tensor], + stddev: list[paddle.Tensor], + ) -> None: + """Update mean and stddev for descriptor.""" + descrpt_list = [self.repflows] + for ii, descrpt in enumerate(descrpt_list): + descrpt.mean = mean[ii] + descrpt.stddev = stddev[ii] + + def get_stat_mean_and_stddev( + self, + ) -> tuple[list[paddle.Tensor], list[paddle.Tensor]]: + """Get mean and stddev for descriptor.""" + mean_list = [self.repflows.mean] + stddev_list = [self.repflows.stddev] + return mean_list, stddev_list + + def serialize(self) -> dict: + repflows = self.repflows + data = { + "@class": "Descriptor", + "type": "dpa3", + "@version": 1, + "ntypes": self.ntypes, + "repflow_args": self.repflow_args.serialize(), + "concat_output_tebd": self.concat_output_tebd, + "activation_function": self.activation_function, + "precision": self.precision, + "exclude_types": self.exclude_types, + "env_protection": self.env_protection, + "trainable": self.trainable, + "use_econf_tebd": self.use_econf_tebd, + "use_tebd_bias": self.use_tebd_bias, + "type_map": self.type_map, + "type_embedding": self.type_embedding.embedding.serialize(), + } + repflow_variable = { + "edge_embd": repflows.edge_embd.serialize(), + "angle_embd": repflows.angle_embd.serialize(), + "repflow_layers": [layer.serialize() for layer in repflows.layers], + "env_mat": DPEnvMat(repflows.rcut, repflows.rcut_smth).serialize(), + "@variables": { + "davg": to_numpy_array(repflows["davg"]), + "dstd": to_numpy_array(repflows["dstd"]), + }, + } + data.update( + { + "repflow_variable": repflow_variable, + } + ) + return data + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA3": + data = data.copy() + version = data.pop("@version") + check_version_compatibility(version, 1, 1) + data.pop("@class") + data.pop("type") + repflow_variable = data.pop("repflow_variable").copy() + type_embedding = data.pop("type_embedding") + data["repflow"] = RepFlowArgs(**data.pop("repflow_args")) + obj = cls(**data) + obj.type_embedding.embedding = TypeEmbedNetConsistent.deserialize( + type_embedding + ) + + def t_cvt(xx): + return paddle.to_tensor(xx, dtype=obj.repflows.prec, place=env.DEVICE) + + # deserialize repflow + statistic_repflows = repflow_variable.pop("@variables") + env_mat = repflow_variable.pop("env_mat") + repflow_layers = repflow_variable.pop("repflow_layers") + obj.repflows.edge_embd = MLPLayer.deserialize(repflow_variable.pop("edge_embd")) + obj.repflows.angle_embd = MLPLayer.deserialize( + repflow_variable.pop("angle_embd") + ) + obj.repflows["davg"] = t_cvt(statistic_repflows["davg"]) + obj.repflows["dstd"] = t_cvt(statistic_repflows["dstd"]) + obj.repflows.layers = paddle.nn.LayerList( + [RepFlowLayer.deserialize(layer) for layer in repflow_layers] + ) + return obj + + def forward( + self, + extended_coord: paddle.Tensor, + extended_atype: paddle.Tensor, + nlist: paddle.Tensor, + mapping: Optional[paddle.Tensor] = None, + comm_dict: Optional[dict[str, paddle.Tensor]] = None, + ): + """Compute the descriptor. + + Parameters + ---------- + extended_coord + The extended coordinates of atoms. shape: nf x (nallx3) + extended_atype + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + mapping + The index mapping, mapps extended region index to local region. + comm_dict + The data needed for communication for parallel inference. + + Returns + ------- + node_ebd + The output descriptor. shape: nf x nloc x n_dim (or n_dim + tebd_dim) + rot_mat + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x e_dim x 3 + edge_ebd + The edge embedding. + shape: nf x nloc x nnei x e_dim + h2 + The rotationally equivariant pair-partical representation. + shape: nf x nloc x nnei x 3 + sw + The smooth switch function. shape: nf x nloc x nnei + + """ + # cast the input to internal precsion + extended_coord = extended_coord.to(dtype=self.prec) + nframes, nloc, nnei = nlist.shape + nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 + + node_ebd_ext = self.type_embedding(extended_atype) + node_ebd_inp = node_ebd_ext[:, :nloc, :] + # repflows + node_ebd, edge_ebd, h2, rot_mat, sw = self.repflows( + nlist, + extended_coord, + extended_atype, + node_ebd_ext, + mapping, + comm_dict=comm_dict, + ) + if self.concat_output_tebd: + node_ebd = paddle.concat([node_ebd, node_ebd_inp], axis=-1) + return ( + node_ebd.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + rot_mat.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + edge_ebd.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + h2.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + sw.to(dtype=env.GLOBAL_PD_FLOAT_PRECISION), + ) + + @classmethod + def update_sel( + cls, + train_data: DeepmdDataSystem, + type_map: Optional[list[str]], + local_jdata: dict, + ) -> tuple[dict, Optional[float]]: + """Update the selection and perform neighbor statistics. + + Parameters + ---------- + train_data : DeepmdDataSystem + data used to do neighbor statistics + type_map : list[str], optional + The name of each type of atoms + local_jdata : dict + The local data refer to the current class + + Returns + ------- + dict + The updated local data + float + The minimum distance between two atoms + """ + local_jdata_cpy = local_jdata.copy() + update_sel = UpdateSel() + min_nbor_dist, repflow_e_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repflow"]["e_rcut"], + local_jdata_cpy["repflow"]["e_sel"], + True, + ) + local_jdata_cpy["repflow"]["e_sel"] = repflow_e_sel[0] + + min_nbor_dist, repflow_a_sel = update_sel.update_one_sel( + train_data, + type_map, + local_jdata_cpy["repflow"]["a_rcut"], + local_jdata_cpy["repflow"]["a_sel"], + True, + ) + local_jdata_cpy["repflow"]["a_sel"] = repflow_a_sel[0] + + return local_jdata_cpy, min_nbor_dist + + def enable_compression( + self, + min_nbor_dist: float, + table_extrapolate: float = 5, + table_stride_1: float = 0.01, + table_stride_2: float = 0.1, + check_frequency: int = -1, + ) -> None: + """Receive the statistics (distance, max_nbor_size and env_mat_range) of the training data. + + Parameters + ---------- + min_nbor_dist + The nearest distance between atoms + table_extrapolate + The scale of model extrapolation + table_stride_1 + The uniform stride of the first table + table_stride_2 + The uniform stride of the second table + check_frequency + The overflow check frequency + """ + raise NotImplementedError("Compression is unsupported for DPA3.") diff --git a/deepmd/pd/model/descriptor/repflow_layer.py b/deepmd/pd/model/descriptor/repflow_layer.py new file mode 100644 index 0000000000..d8ff384107 --- /dev/null +++ b/deepmd/pd/model/descriptor/repflow_layer.py @@ -0,0 +1,932 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Optional, + Union, +) + +import paddle +import paddle.nn as nn + +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.pd.model.descriptor.repformer_layer import ( + _apply_nlist_mask, + _apply_switch, + _make_nei_g1, + get_residual, +) +from deepmd.pd.model.network.mlp import ( + MLPLayer, +) +from deepmd.pd.utils.env import ( + PRECISION_DICT, +) +from deepmd.pd.utils.utils import ( + ActivationFn, + to_numpy_array, + to_paddle_tensor, +) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +class RepFlowLayer(paddle.nn.Layer): + def __init__( + self, + e_rcut: float, + e_rcut_smth: float, + e_sel: int, + a_rcut: float, + a_rcut_smth: float, + a_sel: int, + ntypes: int, + n_dim: int = 128, + e_dim: int = 16, + a_dim: int = 64, + a_compress_rate: int = 0, + a_compress_use_split: bool = False, + a_compress_e_rate: int = 1, + n_multi_edge_message: int = 1, + axis_neuron: int = 4, + update_angle: bool = True, + optim_update: bool = True, + smooth_edge_update: bool = False, + activation_function: str = "silu", + update_style: str = "res_residual", + update_residual: float = 0.1, + update_residual_init: str = "const", + precision: str = "float64", + seed: Optional[Union[int, list[int]]] = None, + ) -> None: + super().__init__() + self.epsilon = 1e-4 # protection of 1./nnei + self.e_rcut = float(e_rcut) + self.e_rcut_smth = float(e_rcut_smth) + self.ntypes = ntypes + e_sel = [e_sel] if isinstance(e_sel, int) else e_sel + self.nnei = sum(e_sel) + assert len(e_sel) == 1 + self.e_sel = e_sel + self.sec = self.e_sel + self.a_rcut = a_rcut + self.a_rcut_smth = a_rcut_smth + self.a_sel = a_sel + self.n_dim = n_dim + self.e_dim = e_dim + self.a_dim = a_dim + self.a_compress_rate = a_compress_rate + if a_compress_rate != 0: + assert (a_dim * a_compress_e_rate) % (2 * a_compress_rate) == 0, ( + f"For a_compress_rate of {a_compress_rate}, a_dim*a_compress_e_rate must be divisible by {2 * a_compress_rate}. " + f"Currently, a_dim={a_dim} and a_compress_e_rate={a_compress_e_rate} is not valid." + ) + self.n_multi_edge_message = n_multi_edge_message + assert self.n_multi_edge_message >= 1, "n_multi_edge_message must >= 1!" + self.axis_neuron = axis_neuron + self.update_angle = update_angle + self.activation_function = activation_function + self.act = ActivationFn(activation_function) + self.update_style = update_style + self.update_residual = update_residual + self.update_residual_init = update_residual_init + self.a_compress_e_rate = a_compress_e_rate + self.a_compress_use_split = a_compress_use_split + self.precision = precision + self.seed = seed + self.prec = PRECISION_DICT[precision] + self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update + + assert update_residual_init in [ + "norm", + "const", + ], "'update_residual_init' only support 'norm' or 'const'!" + + self.update_residual = update_residual + self.update_residual_init = update_residual_init + self.n_residual = [] + self.e_residual = [] + self.a_residual = [] + self.edge_info_dim = self.n_dim * 2 + self.e_dim + + # node self mlp + self.node_self_mlp = MLPLayer( + n_dim, + n_dim, + precision=precision, + seed=child_seed(seed, 0), + ) + if self.update_style == "res_residual": + self.n_residual.append( + get_residual( + n_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(seed, 1), + ) + ) + + # node sym (grrg + drrd) + self.n_sym_dim = n_dim * self.axis_neuron + e_dim * self.axis_neuron + self.node_sym_linear = MLPLayer( + self.n_sym_dim, + n_dim, + precision=precision, + seed=child_seed(seed, 2), + ) + if self.update_style == "res_residual": + self.n_residual.append( + get_residual( + n_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(seed, 3), + ) + ) + + # node edge message + self.node_edge_linear = MLPLayer( + self.edge_info_dim, + self.n_multi_edge_message * n_dim, + precision=precision, + seed=child_seed(seed, 4), + ) + if self.update_style == "res_residual": + for head_index in range(self.n_multi_edge_message): + self.n_residual.append( + get_residual( + n_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(child_seed(seed, 5), head_index), + ) + ) + + # edge self message + self.edge_self_linear = MLPLayer( + self.edge_info_dim, + e_dim, + precision=precision, + seed=child_seed(seed, 6), + ) + if self.update_style == "res_residual": + self.e_residual.append( + get_residual( + e_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(seed, 7), + ) + ) + + if self.update_angle: + self.angle_dim = self.a_dim + if self.a_compress_rate == 0: + # angle + node + edge * 2 + self.angle_dim += self.n_dim + 2 * self.e_dim + self.a_compress_n_linear = None + self.a_compress_e_linear = None + self.e_a_compress_dim = e_dim + self.n_a_compress_dim = n_dim + else: + # angle + a_dim/c + a_dim/2c * 2 * e_rate + self.angle_dim += (1 + self.a_compress_e_rate) * ( + self.a_dim // self.a_compress_rate + ) + self.e_a_compress_dim = ( + self.a_dim // (2 * self.a_compress_rate) * self.a_compress_e_rate + ) + self.n_a_compress_dim = self.a_dim // self.a_compress_rate + if not self.a_compress_use_split: + self.a_compress_n_linear = MLPLayer( + self.n_dim, + self.n_a_compress_dim, + precision=precision, + bias=False, + seed=child_seed(seed, 8), + ) + self.a_compress_e_linear = MLPLayer( + self.e_dim, + self.e_a_compress_dim, + precision=precision, + bias=False, + seed=child_seed(seed, 9), + ) + else: + self.a_compress_n_linear = None + self.a_compress_e_linear = None + + # edge angle message + self.edge_angle_linear1 = MLPLayer( + self.angle_dim, + self.e_dim, + precision=precision, + seed=child_seed(seed, 10), + ) + self.edge_angle_linear2 = MLPLayer( + self.e_dim, + self.e_dim, + precision=precision, + seed=child_seed(seed, 11), + ) + if self.update_style == "res_residual": + self.e_residual.append( + get_residual( + self.e_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(seed, 12), + ) + ) + + # angle self message + self.angle_self_linear = MLPLayer( + self.angle_dim, + self.a_dim, + precision=precision, + seed=child_seed(seed, 13), + ) + if self.update_style == "res_residual": + self.a_residual.append( + get_residual( + self.a_dim, + self.update_residual, + self.update_residual_init, + precision=precision, + seed=child_seed(seed, 14), + ) + ) + else: + self.angle_self_linear = None + self.edge_angle_linear1 = None + self.edge_angle_linear2 = None + self.a_compress_n_linear = None + self.a_compress_e_linear = None + self.angle_dim = 0 + + self.n_residual = nn.ParameterList(self.n_residual) + self.e_residual = nn.ParameterList(self.e_residual) + self.a_residual = nn.ParameterList(self.a_residual) + + @staticmethod + def _cal_hg( + edge_ebd: paddle.Tensor, + h2: paddle.Tensor, + nlist_mask: paddle.Tensor, + sw: paddle.Tensor, + ) -> paddle.Tensor: + """ + Calculate the transposed rotation matrix. + + Parameters + ---------- + edge_ebd + Neighbor-wise/Pair-wise edge embeddings, with shape nb x nloc x nnei x e_dim. + h2 + Neighbor-wise/Pair-wise equivariant rep tensors, with shape nb x nloc x nnei x 3. + nlist_mask + Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei. + sw + The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, + and remains 0 beyond rcut, with shape nb x nloc x nnei. + + Returns + ------- + hg + The transposed rotation matrix, with shape nb x nloc x 3 x e_dim. + """ + # edge_ebd: nb x nloc x nnei x e_dim + # h2: nb x nloc x nnei x 3 + # msk: nb x nloc x nnei + nb, nloc, nnei, _ = edge_ebd.shape + e_dim = edge_ebd.shape[-1] + # nb x nloc x nnei x e_dim + edge_ebd = _apply_nlist_mask(edge_ebd, nlist_mask) + edge_ebd = _apply_switch(edge_ebd, sw) + invnnei = paddle.rsqrt( + float(nnei) + * paddle.ones((nb, nloc, 1, 1), dtype=edge_ebd.dtype).to( + device=edge_ebd.place + ) + ) + # nb x nloc x 3 x e_dim + h2g2 = paddle.matmul(paddle.matrix_transpose(h2), edge_ebd) * invnnei + return h2g2 + + @staticmethod + def _cal_grrg(h2g2: paddle.Tensor, axis_neuron: int) -> paddle.Tensor: + """ + Calculate the atomic invariant rep. + + Parameters + ---------- + h2g2 + The transposed rotation matrix, with shape nb x nloc x 3 x e_dim. + axis_neuron + Size of the submatrix. + + Returns + ------- + grrg + Atomic invariant rep, with shape nb x nloc x (axis_neuron x e_dim) + """ + # nb x nloc x 3 x e_dim + nb, nloc, _, e_dim = h2g2.shape + # nb x nloc x 3 x axis + h2g2m = h2g2[..., :axis_neuron] + # nb x nloc x axis x e_dim + g1_13 = paddle.matmul(paddle.matrix_transpose(h2g2m), h2g2) / (3.0**1) + # nb x nloc x (axisxng2) + g1_13 = g1_13.reshape([nb, nloc, axis_neuron * e_dim]) + return g1_13 + + def symmetrization_op( + self, + edge_ebd: paddle.Tensor, + h2: paddle.Tensor, + nlist_mask: paddle.Tensor, + sw: paddle.Tensor, + axis_neuron: int, + ) -> paddle.Tensor: + """ + Symmetrization operator to obtain atomic invariant rep. + + Parameters + ---------- + edge_ebd + Neighbor-wise/Pair-wise invariant rep tensors, with shape nb x nloc x nnei x e_dim. + h2 + Neighbor-wise/Pair-wise equivariant rep tensors, with shape nb x nloc x nnei x 3. + nlist_mask + Neighbor list mask, where zero means no neighbor, with shape nb x nloc x nnei. + sw + The switch function, which equals 1 within the rcut_smth range, smoothly decays from 1 to 0 between rcut_smth and rcut, + and remains 0 beyond rcut, with shape nb x nloc x nnei. + axis_neuron + Size of the submatrix. + + Returns + ------- + grrg + Atomic invariant rep, with shape nb x nloc x (axis_neuron x e_dim) + """ + # edge_ebd: nb x nloc x nnei x e_dim + # h2: nb x nloc x nnei x 3 + # msk: nb x nloc x nnei + nb, nloc, nnei, _ = edge_ebd.shape + # nb x nloc x 3 x e_dim + h2g2 = self._cal_hg( + edge_ebd, + h2, + nlist_mask, + sw, + ) + # nb x nloc x (axisxng2) + g1_13 = self._cal_grrg(h2g2, axis_neuron) + return g1_13 + + def optim_angle_update( + self, + angle_ebd: paddle.Tensor, + node_ebd: paddle.Tensor, + edge_ebd: paddle.Tensor, + feat: str = "edge", + ) -> paddle.Tensor: + if feat == "edge": + assert self.edge_angle_linear1 is not None + matrix, bias = self.edge_angle_linear1.matrix, self.edge_angle_linear1.bias + elif feat == "angle": + assert self.angle_self_linear is not None + matrix, bias = self.angle_self_linear.matrix, self.angle_self_linear.bias + else: + raise NotImplementedError + assert bias is not None + + angle_dim = angle_ebd.shape[-1] + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + # angle_dim, node_dim, edge_dim, edge_dim + sub_angle, sub_node, sub_edge_ij, sub_edge_ik = paddle.split( + matrix, [angle_dim, node_dim, edge_dim, edge_dim] + ) + + # nf * nloc * a_sel * a_sel * angle_dim + sub_angle_update = paddle.matmul(angle_ebd, sub_angle) + # nf * nloc * angle_dim + sub_node_update = paddle.matmul(node_ebd, sub_node) + # nf * nloc * a_nnei * angle_dim + sub_edge_update_ij = paddle.matmul(edge_ebd, sub_edge_ij) + sub_edge_update_ik = paddle.matmul(edge_ebd, sub_edge_ik) + + result_update = ( + bias + + sub_node_update.unsqueeze(2).unsqueeze(3) + + sub_edge_update_ij.unsqueeze(2) + + sub_edge_update_ik.unsqueeze(3) + + sub_angle_update + ) + return result_update + + def optim_edge_update( + self, + node_ebd: paddle.Tensor, + node_ebd_ext: paddle.Tensor, + edge_ebd: paddle.Tensor, + nlist: paddle.Tensor, + feat: str = "node", + ) -> paddle.Tensor: + if feat == "node": + matrix, bias = self.node_edge_linear.matrix, self.node_edge_linear.bias + elif feat == "edge": + matrix, bias = self.edge_self_linear.matrix, self.edge_self_linear.bias + else: + raise NotImplementedError + assert bias is not None + + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + # node_dim, node_dim, edge_dim + node, node_ext, edge = paddle.split(matrix, [node_dim, node_dim, edge_dim]) + + # nf * nloc * node/edge_dim + sub_node_update = paddle.matmul(node_ebd, node) + # nf * nall * node/edge_dim + sub_node_ext_update = paddle.matmul(node_ebd_ext, node_ext) + # nf * nloc * nnei * node/edge_dim + sub_node_ext_update = _make_nei_g1(sub_node_ext_update, nlist) + # nf * nloc * nnei * node/edge_dim + sub_edge_update = paddle.matmul(edge_ebd, edge) + + result_update = ( + bias + sub_node_update.unsqueeze(2) + sub_edge_update + sub_node_ext_update + ) + return result_update + + def forward( + self, + node_ebd_ext: paddle.Tensor, # nf x nall x n_dim + edge_ebd: paddle.Tensor, # nf x nloc x nnei x e_dim + h2: paddle.Tensor, # nf x nloc x nnei x 3 + angle_ebd: paddle.Tensor, # nf x nloc x a_nnei x a_nnei x a_dim + nlist: paddle.Tensor, # nf x nloc x nnei + nlist_mask: paddle.Tensor, # nf x nloc x nnei + sw: paddle.Tensor, # switch func, nf x nloc x nnei + a_nlist: paddle.Tensor, # nf x nloc x a_nnei + a_nlist_mask: paddle.Tensor, # nf x nloc x a_nnei + a_sw: paddle.Tensor, # switch func, nf x nloc x a_nnei + ): + """ + Parameters + ---------- + node_ebd_ext : nf x nall x n_dim + Extended node embedding. + edge_ebd : nf x nloc x nnei x e_dim + Edge embedding. + h2 : nf x nloc x nnei x 3 + Pair-atom channel, equivariant. + angle_ebd : nf x nloc x a_nnei x a_nnei x a_dim + Angle embedding. + nlist : nf x nloc x nnei + Neighbor list. (padded neis are set to 0) + nlist_mask : nf x nloc x nnei + Masks of the neighbor list. real nei 1 otherwise 0 + sw : nf x nloc x nnei + Switch function. + a_nlist : nf x nloc x a_nnei + Neighbor list for angle. (padded neis are set to 0) + a_nlist_mask : nf x nloc x a_nnei + Masks of the neighbor list for angle. real nei 1 otherwise 0 + a_sw : nf x nloc x a_nnei + Switch function for angle. + + Returns + ------- + n_updated: nf x nloc x n_dim + Updated node embedding. + e_updated: nf x nloc x nnei x e_dim + Updated edge embedding. + a_updated : nf x nloc x a_nnei x a_nnei x a_dim + Updated angle embedding. + """ + nb, nloc, nnei, _ = edge_ebd.shape + nall = node_ebd_ext.shape[1] + node_ebd = node_ebd_ext[:, :nloc, :] + if paddle.in_dynamic_mode(): + assert [nb, nloc] == node_ebd.shape[:2] + if paddle.in_dynamic_mode(): + assert [nb, nloc, nnei] == h2.shape[:3] + del a_nlist # may be used in the future + + n_update_list: list[paddle.Tensor] = [node_ebd] + e_update_list: list[paddle.Tensor] = [edge_ebd] + a_update_list: list[paddle.Tensor] = [angle_ebd] + + # node self mlp + node_self_mlp = self.act(self.node_self_mlp(node_ebd)) + n_update_list.append(node_self_mlp) + + nei_node_ebd = _make_nei_g1(node_ebd_ext, nlist) + + # node sym (grrg + drrd) + node_sym_list: list[paddle.Tensor] = [] + node_sym_list.append( + self.symmetrization_op( + edge_ebd, + h2, + nlist_mask, + sw, + self.axis_neuron, + ) + ) + node_sym_list.append( + self.symmetrization_op( + nei_node_ebd, + h2, + nlist_mask, + sw, + self.axis_neuron, + ) + ) + node_sym = self.act(self.node_sym_linear(paddle.concat(node_sym_list, axis=-1))) + n_update_list.append(node_sym) + + if not self.optim_update: + # nb x nloc x nnei x (n_dim * 2 + e_dim) + edge_info = paddle.concat( + [ + paddle.tile(node_ebd.unsqueeze(-2), [1, 1, self.nnei, 1]), + nei_node_ebd, + edge_ebd, + ], + axis=-1, + ) + else: + edge_info = None + + # node edge message + # nb x nloc x nnei x (h * n_dim) + if not self.optim_update: + assert edge_info is not None + node_edge_update = self.act( + self.node_edge_linear(edge_info) + ) * sw.unsqueeze(-1) + else: + node_edge_update = self.act( + self.optim_edge_update( + node_ebd, + node_ebd_ext, + edge_ebd, + nlist, + "node", + ) + ) * sw.unsqueeze(-1) + + node_edge_update = paddle.sum(node_edge_update, axis=-2) / self.nnei + if self.n_multi_edge_message > 1: + # nb x nloc x nnei x h x n_dim + node_edge_update_mul_head = node_edge_update.reshape( + [nb, nloc, self.n_multi_edge_message, self.n_dim] + ) + for head_index in range(self.n_multi_edge_message): + n_update_list.append(node_edge_update_mul_head[:, :, head_index, :]) + else: + n_update_list.append(node_edge_update) + # update node_ebd + n_updated = self.list_update(n_update_list, "node") + + # edge self message + if not self.optim_update: + assert edge_info is not None + edge_self_update = self.act(self.edge_self_linear(edge_info)) + else: + edge_self_update = self.act( + self.optim_edge_update( + node_ebd, + node_ebd_ext, + edge_ebd, + nlist, + "edge", + ) + ) + e_update_list.append(edge_self_update) + + if self.update_angle: + if paddle.in_dynamic_mode(): + assert self.angle_self_linear is not None + if paddle.in_dynamic_mode(): + assert self.edge_angle_linear1 is not None + if paddle.in_dynamic_mode(): + assert self.edge_angle_linear2 is not None + # get angle info + if self.a_compress_rate != 0: + if not self.a_compress_use_split: + if paddle.in_dynamic_mode(): + assert self.a_compress_n_linear is not None + if paddle.in_dynamic_mode(): + assert self.a_compress_e_linear is not None + node_ebd_for_angle = self.a_compress_n_linear(node_ebd) + edge_ebd_for_angle = self.a_compress_e_linear(edge_ebd) + else: + # use the first a_compress_dim dim for node and edge + node_ebd_for_angle = node_ebd[:, :, : self.n_a_compress_dim] + edge_ebd_for_angle = edge_ebd[:, :, :, : self.e_a_compress_dim] + else: + node_ebd_for_angle = node_ebd + edge_ebd_for_angle = edge_ebd + + # nb x nloc x a_nnei x e_dim + edge_for_angle = edge_ebd_for_angle[:, :, : self.a_sel, :] + # nb x nloc x a_nnei x e_dim + edge_for_angle = paddle.where( + a_nlist_mask.unsqueeze(-1), + edge_for_angle, + paddle.zeros_like(edge_for_angle), + ).astype(edge_for_angle.dtype) + if not self.optim_update: + # nb x nloc x a_nnei x a_nnei x n_dim + node_for_angle_info = paddle.tile( + node_ebd_for_angle.unsqueeze(2).unsqueeze(2), + [1, 1, self.a_sel, self.a_sel, 1], + ) + # nb x nloc x (a_nnei) x a_nnei x edge_ebd + edge_for_angle_i = paddle.tile( + edge_for_angle.unsqueeze(2), (1, 1, self.a_sel, 1, 1) + ) + # nb x nloc x a_nnei x (a_nnei) x e_dim + edge_for_angle_j = paddle.tile( + edge_for_angle.unsqueeze(3), (1, 1, 1, self.a_sel, 1) + ) + # nb x nloc x a_nnei x a_nnei x (e_dim + e_dim) + edge_for_angle_info = paddle.concat( + [edge_for_angle_i, edge_for_angle_j], axis=-1 + ) + angle_info_list = [angle_ebd] + angle_info_list.append(node_for_angle_info) + angle_info_list.append(edge_for_angle_info) + # nb x nloc x a_nnei x a_nnei x (a + n_dim + e_dim*2) or (a + a/c + a/c) + angle_info = paddle.concat(angle_info_list, axis=-1) + else: + angle_info = None + + # edge angle message + # nb x nloc x a_nnei x a_nnei x e_dim + if not self.optim_update: + assert angle_info is not None + edge_angle_update = self.act(self.edge_angle_linear1(angle_info)) + else: + edge_angle_update = self.act( + self.optim_angle_update( + angle_ebd, + node_ebd_for_angle, + edge_for_angle, + "edge", + ) + ) + + # nb x nloc x a_nnei x a_nnei x e_dim + weighted_edge_angle_update = ( + a_sw[..., None, None] * a_sw[..., None, :, None] * edge_angle_update + ) + # nb x nloc x a_nnei x e_dim + reduced_edge_angle_update = paddle.sum( + weighted_edge_angle_update, axis=-2 + ) / (self.a_sel**0.5) + # nb x nloc x nnei x e_dim + padding_edge_angle_update = paddle.concat( + [ + reduced_edge_angle_update, + paddle.zeros( + [nb, nloc, self.nnei - self.a_sel, self.e_dim], + dtype=edge_ebd.dtype, + ).to(device=edge_ebd.place), + ], + axis=2, + ) + if not self.smooth_edge_update: + # will be deprecated in the future + full_mask = paddle.concat( + [ + a_nlist_mask, + paddle.zeros( + [nb, nloc, self.nnei - self.a_sel], + dtype=a_nlist_mask.dtype, + ).to(a_nlist_mask.place), + ], + axis=-1, + ) + padding_edge_angle_update = paddle.where( + full_mask.unsqueeze(-1), padding_edge_angle_update, edge_ebd + ) + e_update_list.append( + self.act(self.edge_angle_linear2(padding_edge_angle_update)) + ) + # update edge_ebd + e_updated = self.list_update(e_update_list, "edge") + + # angle self message + # nb x nloc x a_nnei x a_nnei x dim_a + if not self.optim_update: + assert angle_info is not None + angle_self_update = self.act(self.angle_self_linear(angle_info)) + else: + angle_self_update = self.act( + self.optim_angle_update( + angle_ebd, + node_ebd_for_angle, + edge_for_angle, + "angle", + ) + ) + a_update_list.append(angle_self_update) + else: + # update edge_ebd + e_updated = self.list_update(e_update_list, "edge") + + # update angle_ebd + a_updated = self.list_update(a_update_list, "angle") + return n_updated, e_updated, a_updated + + def list_update_res_avg( + self, + update_list: list[paddle.Tensor], + ) -> paddle.Tensor: + nitem = len(update_list) + uu = update_list[0] + for ii in range(1, nitem): + uu = uu + update_list[ii] + return uu / (float(nitem) ** 0.5) + + def list_update_res_incr(self, update_list: list[paddle.Tensor]) -> paddle.Tensor: + nitem = len(update_list) + uu = update_list[0] + scale = 1.0 / (float(nitem - 1) ** 0.5) if nitem > 1 else 0.0 + for ii in range(1, nitem): + uu = uu + scale * update_list[ii] + return uu + + def list_update_res_residual( + self, update_list: list[paddle.Tensor], update_name: str = "node" + ) -> paddle.Tensor: + nitem = len(update_list) + uu = update_list[0] + # make jit happy + if update_name == "node": + for ii, vv in enumerate(self.n_residual): + uu = uu + vv * update_list[ii + 1] + elif update_name == "edge": + for ii, vv in enumerate(self.e_residual): + uu = uu + vv * update_list[ii + 1] + elif update_name == "angle": + for ii, vv in enumerate(self.a_residual): + uu = uu + vv * update_list[ii + 1] + else: + raise NotImplementedError + return uu + + def list_update( + self, update_list: list[paddle.Tensor], update_name: str = "node" + ) -> paddle.Tensor: + if self.update_style == "res_avg": + return self.list_update_res_avg(update_list) + elif self.update_style == "res_incr": + return self.list_update_res_incr(update_list) + elif self.update_style == "res_residual": + return self.list_update_res_residual(update_list, update_name=update_name) + else: + raise RuntimeError(f"unknown update style {self.update_style}") + + def serialize(self) -> dict: + """Serialize the networks to a dict. + + Returns + ------- + dict + The serialized networks. + """ + data = { + "@class": "RepFlowLayer", + "@version": 1, + "e_rcut": self.e_rcut, + "e_rcut_smth": self.e_rcut_smth, + "e_sel": self.e_sel, + "a_rcut": self.a_rcut, + "a_rcut_smth": self.a_rcut_smth, + "a_sel": self.a_sel, + "ntypes": self.ntypes, + "n_dim": self.n_dim, + "e_dim": self.e_dim, + "a_dim": self.a_dim, + "a_compress_rate": self.a_compress_rate, + "a_compress_e_rate": self.a_compress_e_rate, + "a_compress_use_split": self.a_compress_use_split, + "n_multi_edge_message": self.n_multi_edge_message, + "axis_neuron": self.axis_neuron, + "activation_function": self.activation_function, + "update_angle": self.update_angle, + "update_style": self.update_style, + "update_residual": self.update_residual, + "update_residual_init": self.update_residual_init, + "precision": self.precision, + "optim_update": self.optim_update, + "smooth_edge_update": self.smooth_edge_update, + "node_self_mlp": self.node_self_mlp.serialize(), + "node_sym_linear": self.node_sym_linear.serialize(), + "node_edge_linear": self.node_edge_linear.serialize(), + "edge_self_linear": self.edge_self_linear.serialize(), + } + if self.update_angle: + data.update( + { + "edge_angle_linear1": self.edge_angle_linear1.serialize(), + "edge_angle_linear2": self.edge_angle_linear2.serialize(), + "angle_self_linear": self.angle_self_linear.serialize(), + } + ) + if self.a_compress_rate != 0 and not self.a_compress_use_split: + data.update( + { + "a_compress_n_linear": self.a_compress_n_linear.serialize(), + "a_compress_e_linear": self.a_compress_e_linear.serialize(), + } + ) + if self.update_style == "res_residual": + data.update( + { + "@variables": { + "n_residual": [to_numpy_array(t) for t in self.n_residual], + "e_residual": [to_numpy_array(t) for t in self.e_residual], + "a_residual": [to_numpy_array(t) for t in self.a_residual], + } + } + ) + return data + + @classmethod + def deserialize(cls, data: dict) -> "RepFlowLayer": + """Deserialize the networks from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = data.copy() + check_version_compatibility(data.pop("@version"), 1, 1) + data.pop("@class") + update_angle = data["update_angle"] + a_compress_rate = data["a_compress_rate"] + a_compress_use_split = data["a_compress_use_split"] + node_self_mlp = data.pop("node_self_mlp") + node_sym_linear = data.pop("node_sym_linear") + node_edge_linear = data.pop("node_edge_linear") + edge_self_linear = data.pop("edge_self_linear") + edge_angle_linear1 = data.pop("edge_angle_linear1", None) + edge_angle_linear2 = data.pop("edge_angle_linear2", None) + angle_self_linear = data.pop("angle_self_linear", None) + a_compress_n_linear = data.pop("a_compress_n_linear", None) + a_compress_e_linear = data.pop("a_compress_e_linear", None) + update_style = data["update_style"] + variables = data.pop("@variables", {}) + n_residual = variables.get("n_residual", data.pop("n_residual", [])) + e_residual = variables.get("e_residual", data.pop("e_residual", [])) + a_residual = variables.get("a_residual", data.pop("a_residual", [])) + + obj = cls(**data) + obj.node_self_mlp = MLPLayer.deserialize(node_self_mlp) + obj.node_sym_linear = MLPLayer.deserialize(node_sym_linear) + obj.node_edge_linear = MLPLayer.deserialize(node_edge_linear) + obj.edge_self_linear = MLPLayer.deserialize(edge_self_linear) + + if update_angle: + if paddle.in_dynamic_mode(): + assert isinstance(edge_angle_linear1, dict) + if paddle.in_dynamic_mode(): + assert isinstance(edge_angle_linear2, dict) + if paddle.in_dynamic_mode(): + assert isinstance(angle_self_linear, dict) + obj.edge_angle_linear1 = MLPLayer.deserialize(edge_angle_linear1) + obj.edge_angle_linear2 = MLPLayer.deserialize(edge_angle_linear2) + obj.angle_self_linear = MLPLayer.deserialize(angle_self_linear) + if a_compress_rate != 0 and not a_compress_use_split: + if paddle.in_dynamic_mode(): + assert isinstance(a_compress_n_linear, dict) + if paddle.in_dynamic_mode(): + assert isinstance(a_compress_e_linear, dict) + obj.a_compress_n_linear = MLPLayer.deserialize(a_compress_n_linear) + obj.a_compress_e_linear = MLPLayer.deserialize(a_compress_e_linear) + + if update_style == "res_residual": + for ii, t in enumerate(obj.n_residual): + t.data = to_paddle_tensor(n_residual[ii]) + for ii, t in enumerate(obj.e_residual): + t.data = to_paddle_tensor(e_residual[ii]) + for ii, t in enumerate(obj.a_residual): + t.data = to_paddle_tensor(a_residual[ii]) + return obj diff --git a/deepmd/pd/model/descriptor/repflows.py b/deepmd/pd/model/descriptor/repflows.py new file mode 100644 index 0000000000..f655c4b47f --- /dev/null +++ b/deepmd/pd/model/descriptor/repflows.py @@ -0,0 +1,546 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Callable, + Optional, + Union, +) + +import paddle + +from deepmd.dpmodel.utils.seed import ( + child_seed, +) +from deepmd.pd.model.descriptor.descriptor import ( + DescriptorBlock, +) +from deepmd.pd.model.descriptor.env_mat import ( + prod_env_mat, +) +from deepmd.pd.model.network.mlp import ( + MLPLayer, +) +from deepmd.pd.utils import ( + env, +) +from deepmd.pd.utils.env import ( + PRECISION_DICT, +) +from deepmd.pd.utils.env_mat_stat import ( + EnvMatStatSe, +) +from deepmd.pd.utils.exclude_mask import ( + PairExcludeMask, +) +from deepmd.pd.utils.utils import ( + ActivationFn, +) +from deepmd.utils.env_mat_stat import ( + StatItem, +) +from deepmd.utils.path import ( + DPPath, +) + +from .repflow_layer import ( + RepFlowLayer, +) + + +@DescriptorBlock.register("se_repflow") +class DescrptBlockRepflows(DescriptorBlock): + r""" + The repflow descriptor block. + + Parameters + ---------- + n_dim : int, optional + The dimension of node representation. + e_dim : int, optional + The dimension of edge representation. + a_dim : int, optional + The dimension of angle representation. + nlayers : int, optional + Number of repflow layers. + e_rcut : float, optional + The edge cut-off radius. + e_rcut_smth : float, optional + Where to start smoothing for edge. For example the 1/r term is smoothed from rcut to rcut_smth. + e_sel : int, optional + Maximally possible number of selected edge neighbors. + a_rcut : float, optional + The angle cut-off radius. + a_rcut_smth : float, optional + Where to start smoothing for angle. For example the 1/r term is smoothed from rcut to rcut_smth. + a_sel : int, optional + Maximally possible number of selected angle neighbors. + a_compress_rate : int, optional + The compression rate for angular messages. The default value is 0, indicating no compression. + If a non-zero integer c is provided, the node and edge dimensions will be compressed + to a_dim/c and a_dim/2c, respectively, within the angular message. + a_compress_e_rate : int, optional + The extra compression rate for edge in angular message compression. The default value is 1. + When using angular message compression with a_compress_rate c and a_compress_e_rate c_e, + the edge dimension will be compressed to (c_e * a_dim / 2c) within the angular message. + a_compress_use_split : bool, optional + Whether to split first sub-vectors instead of linear mapping during angular message compression. + The default value is False. + n_multi_edge_message : int, optional + The head number of multiple edge messages to update node feature. + Default is 1, indicating one head edge message. + axis_neuron : int, optional + The number of dimension of submatrix in the symmetrization ops. + update_angle : bool, optional + Where to update the angle rep. If not, only node and edge rep will be used. + 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 `update_residual` + and `update_residual_init`. + update_residual : float, optional + When update using residual mode, the initial std of residual vector weights. + update_residual_init : str, optional + When update using residual mode, the initialization mode of residual vector weights. + fix_stat_std : float, optional + If non-zero (default is 0.3), use this constant as the normalization standard deviation + instead of computing it from data statistics. + smooth_edge_update : bool, optional + Whether to make edge update smooth. + If True, the edge update from angle message will not use self as padding. + optim_update : bool, optional + Whether to enable the optimized update method. + Uses a more efficient process when enabled. Defaults to True + ntypes : int + Number of element types + activation_function : str, optional + The activation function in the embedding net. + set_davg_zero : bool, optional + Set the normalization average to zero. + precision : str, optional + The precision of the embedding net parameters. + exclude_types : list[list[int]], optional + The excluded pairs of types which have no interaction with each other. + For example, `[[0, 1]]` means no interaction between type 0 and type 1. + env_protection : float, optional + Protection parameter to prevent division by zero errors during environment matrix calculations. + For example, when using paddings, there may be zero distances of neighbors, which may make division by zero error during environment matrix calculations without protection. + seed : int, optional + Random seed for parameter initialization. + """ + + def __init__( + self, + e_rcut, + e_rcut_smth, + e_sel: int, + a_rcut, + a_rcut_smth, + a_sel: int, + ntypes: int, + nlayers: int = 6, + n_dim: int = 128, + e_dim: int = 64, + a_dim: int = 64, + a_compress_rate: int = 0, + a_compress_e_rate: int = 1, + a_compress_use_split: bool = False, + n_multi_edge_message: int = 1, + axis_neuron: int = 4, + update_angle: bool = True, + activation_function: str = "silu", + update_style: str = "res_residual", + update_residual: float = 0.1, + update_residual_init: str = "const", + set_davg_zero: bool = True, + exclude_types: list[tuple[int, int]] = [], + env_protection: float = 0.0, + precision: str = "float64", + fix_stat_std: float = 0.3, + smooth_edge_update: bool = False, + optim_update: bool = True, + seed: Optional[Union[int, list[int]]] = None, + ) -> None: + super().__init__() + self.e_rcut = float(e_rcut) + self.e_rcut_smth = float(e_rcut_smth) + self.e_sel = e_sel + self.a_rcut = float(a_rcut) + self.a_rcut_smth = float(a_rcut_smth) + self.a_sel = a_sel + self.ntypes = ntypes + self.nlayers = nlayers + # for other common desciptor method + sel = [e_sel] if isinstance(e_sel, int) else e_sel + self.nnei = sum(sel) + self.ndescrpt = self.nnei * 4 # use full descriptor. + assert len(sel) == 1 + self.sel = sel + self.rcut = e_rcut + self.rcut_smth = e_rcut_smth + self.sec = self.sel + self.split_sel = self.sel + self.a_compress_rate = a_compress_rate + self.a_compress_e_rate = a_compress_e_rate + self.n_multi_edge_message = n_multi_edge_message + self.axis_neuron = axis_neuron + self.set_davg_zero = set_davg_zero + self.fix_stat_std = fix_stat_std + self.set_stddev_constant = fix_stat_std != 0.0 + self.a_compress_use_split = a_compress_use_split + self.optim_update = optim_update + self.smooth_edge_update = smooth_edge_update + + self.n_dim = n_dim + self.e_dim = e_dim + self.a_dim = a_dim + self.update_angle = update_angle + + self.activation_function = activation_function + self.update_style = update_style + self.update_residual = update_residual + self.update_residual_init = update_residual_init + self.act = ActivationFn(activation_function) + self.prec = PRECISION_DICT[precision] + + # order matters, placed after the assignment of self.ntypes + self.reinit_exclude(exclude_types) + self.env_protection = env_protection + self.precision = precision + self.epsilon = 1e-4 + self.seed = seed + + self.edge_embd = MLPLayer( + 1, self.e_dim, precision=precision, seed=child_seed(seed, 0) + ) + self.angle_embd = MLPLayer( + 1, self.a_dim, precision=precision, bias=False, seed=child_seed(seed, 1) + ) + layers = [] + for ii in range(nlayers): + layers.append( + RepFlowLayer( + e_rcut=self.e_rcut, + e_rcut_smth=self.e_rcut_smth, + e_sel=self.sel, + a_rcut=self.a_rcut, + a_rcut_smth=self.a_rcut_smth, + a_sel=self.a_sel, + ntypes=self.ntypes, + n_dim=self.n_dim, + e_dim=self.e_dim, + a_dim=self.a_dim, + a_compress_rate=self.a_compress_rate, + a_compress_use_split=self.a_compress_use_split, + a_compress_e_rate=self.a_compress_e_rate, + n_multi_edge_message=self.n_multi_edge_message, + axis_neuron=self.axis_neuron, + update_angle=self.update_angle, + activation_function=self.activation_function, + update_style=self.update_style, + update_residual=self.update_residual, + update_residual_init=self.update_residual_init, + precision=precision, + optim_update=self.optim_update, + smooth_edge_update=self.smooth_edge_update, + seed=child_seed(child_seed(seed, 1), ii), + ) + ) + self.layers = paddle.nn.LayerList(layers) + + wanted_shape = (self.ntypes, self.nnei, 4) + mean = paddle.zeros(wanted_shape, dtype=self.prec).to(device=env.DEVICE) + stddev = paddle.ones(wanted_shape, dtype=self.prec).to(device=env.DEVICE) + if self.set_stddev_constant: + stddev = stddev * self.fix_stat_std + self.register_buffer("mean", mean) + self.register_buffer("stddev", stddev) + self.stats = None + + def get_rcut(self) -> float: + """Returns the cut-off radius.""" + return self.e_rcut + + def get_rcut_smth(self) -> float: + """Returns the radius where the neighbor information starts to smoothly decay to 0.""" + return self.e_rcut_smth + + def get_nsel(self) -> int: + """Returns the number of selected atoms in the cut-off radius.""" + return sum(self.sel) + + def get_sel(self) -> list[int]: + """Returns the number of selected atoms for each type.""" + return self.sel + + def get_ntypes(self) -> int: + """Returns the number of element types.""" + return self.ntypes + + def get_dim_out(self) -> int: + """Returns the output dimension.""" + return self.dim_out + + def get_dim_in(self) -> int: + """Returns the input dimension.""" + return self.dim_in + + def get_dim_emb(self) -> int: + """Returns the embedding dimension e_dim.""" + return self.e_dim + + def __setitem__(self, key, value) -> None: + if key in ("avg", "data_avg", "davg"): + self.mean = value + elif key in ("std", "data_std", "dstd"): + self.stddev = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.mean + elif key in ("std", "data_std", "dstd"): + return self.stddev + else: + raise KeyError(key) + + def mixed_types(self) -> bool: + """If true, the descriptor + 1. assumes total number of atoms aligned across frames; + 2. requires a neighbor list that does not distinguish different atomic types. + + If false, the descriptor + 1. assumes total number of atoms of each atom type aligned across frames; + 2. requires a neighbor list that distinguishes different atomic types. + + """ + return True + + def get_env_protection(self) -> float: + """Returns the protection of building environment matrix.""" + return self.env_protection + + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.n_dim + + @property + def dim_in(self): + """Returns the atomic input dimension of this descriptor.""" + return self.n_dim + + @property + def dim_emb(self): + """Returns the embedding dimension e_dim.""" + return self.get_dim_emb() + + def reinit_exclude( + self, + exclude_types: list[tuple[int, int]] = [], + ) -> None: + self.exclude_types = exclude_types + self.emask = PairExcludeMask(self.ntypes, exclude_types=exclude_types) + + def forward( + self, + nlist: paddle.Tensor, + extended_coord: paddle.Tensor, + extended_atype: paddle.Tensor, + extended_atype_embd: Optional[paddle.Tensor] = None, + mapping: Optional[paddle.Tensor] = None, + comm_dict: Optional[dict[str, paddle.Tensor]] = None, + ): + if comm_dict is None: + assert mapping is not None + assert extended_atype_embd is not None + nframes, nloc, nnei = nlist.shape + nall = extended_coord.reshape([nframes, -1]).shape[1] // 3 + atype = extended_atype[:, :nloc] + # nb x nloc x nnei + exclude_mask = self.emask(nlist, extended_atype) + nlist = paddle.where(exclude_mask != 0, nlist, -1) + # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1 + dmatrix, diff, sw = prod_env_mat( + extended_coord, + nlist, + atype, + self.mean, + self.stddev, + self.e_rcut, + self.e_rcut_smth, + protection=self.env_protection, + ) + nlist_mask = nlist != -1 + sw = paddle.squeeze(sw, -1) + # beyond the cutoff sw should be 0.0 + sw = sw.masked_fill(~nlist_mask, 0.0) + + # [nframes, nloc, tebd_dim] + if comm_dict is None: + if paddle.in_dynamic_mode(): + assert isinstance(extended_atype_embd, paddle.Tensor) + atype_embd = extended_atype_embd[:, :nloc, :] + if paddle.in_dynamic_mode(): + assert atype_embd.shape == [nframes, nloc, self.n_dim] + else: + atype_embd = extended_atype_embd + if paddle.in_dynamic_mode(): + assert isinstance(atype_embd, paddle.Tensor) + node_ebd = self.act(atype_embd) + n_dim = node_ebd.shape[-1] + # nb x nloc x nnei x 1, nb x nloc x nnei x 3 + edge_input, h2 = paddle.split(dmatrix, [1, 3], axis=-1) + # nb x nloc x nnei x e_dim + edge_ebd = self.act(self.edge_embd(edge_input)) + + # get angle nlist (maybe smaller) + a_dist_mask = (paddle.linalg.norm(diff, axis=-1) < self.a_rcut)[ + :, :, : self.a_sel + ] + a_nlist = nlist[:, :, : self.a_sel] + a_nlist = paddle.where(a_dist_mask, a_nlist, -1) + _, a_diff, a_sw = prod_env_mat( + extended_coord, + a_nlist, + atype, + self.mean[:, : self.a_sel], + self.stddev[:, : self.a_sel], + self.a_rcut, + self.a_rcut_smth, + protection=self.env_protection, + ) + a_nlist_mask = a_nlist != -1 + a_sw = paddle.squeeze(a_sw, -1) + # beyond the cutoff sw should be 0.0 + a_sw = a_sw.masked_fill(~a_nlist_mask, 0.0) + a_nlist[a_nlist == -1] = 0 + + # nf x nloc x a_nnei x 3 + normalized_diff_i = a_diff / ( + paddle.linalg.norm(a_diff, axis=-1, keepdim=True) + 1e-6 + ) + # nf x nloc x 3 x a_nnei + normalized_diff_j = paddle.transpose(normalized_diff_i, [0, 1, 3, 2]) + # nf x nloc x a_nnei x a_nnei + # 1 - 1e-6 for paddle.acos stability + cosine_ij = paddle.matmul(normalized_diff_i, normalized_diff_j) * (1 - 1e-6) + # nf x nloc x a_nnei x a_nnei x 1 + cosine_ij = cosine_ij.unsqueeze(-1) / (paddle.pi**0.5) + # nf x nloc x a_nnei x a_nnei x a_dim + angle_ebd = self.angle_embd(cosine_ij).reshape( + [nframes, nloc, self.a_sel, self.a_sel, self.a_dim] + ) + + # set all padding positions to index of 0 + # if the a neighbor is real or not is indicated by nlist_mask + nlist[nlist == -1] = 0 + # nb x nall x n_dim + if comm_dict is None: + assert mapping is not None + mapping = ( + mapping.reshape([nframes, nall]) + .unsqueeze(-1) + .expand([-1, -1, self.n_dim]) + ) + for idx, ll in enumerate(self.layers): + # node_ebd: nb x nloc x n_dim + # node_ebd_ext: nb x nall x n_dim + if comm_dict is None: + assert mapping is not None + node_ebd_ext = paddle.take_along_axis( + node_ebd, mapping, 1, broadcast=False + ) + else: + raise NotImplementedError("Not implemented") + node_ebd, edge_ebd, angle_ebd = ll.forward( + node_ebd_ext, + edge_ebd, + h2, + angle_ebd, + nlist, + nlist_mask, + sw, + a_nlist, + a_nlist_mask, + a_sw, + ) + + # nb x nloc x 3 x e_dim + h2g2 = RepFlowLayer._cal_hg(edge_ebd, h2, nlist_mask, sw) + # (nb x nloc) x e_dim x 3 + rot_mat = paddle.transpose(h2g2, (0, 1, 3, 2)) + + return ( + node_ebd, + edge_ebd, + h2, + rot_mat.reshape([nframes, nloc, self.dim_emb, 3]), + sw, + ) + + def compute_input_stats( + self, + merged: Union[Callable[[], list[dict]], list[dict]], + path: Optional[DPPath] = None, + ) -> None: + """ + Compute the input statistics (e.g. mean and stddev) for the descriptors from packed data. + + Parameters + ---------- + merged : Union[Callable[[], list[dict]], list[dict]] + - list[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `paddle.Tensor` + originating from the `i`-th data system. + - Callable[[], list[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + path : Optional[DPPath] + The path to the stat file. + + """ + if self.set_stddev_constant and self.set_davg_zero: + return + env_mat_stat = EnvMatStatSe(self) + if path is not None: + path = path / env_mat_stat.get_hash() + if path is None or not path.is_dir(): + if callable(merged): + # only get data for once + sampled = merged() + else: + sampled = merged + else: + sampled = [] + env_mat_stat.load_or_compute_stats(sampled, path) + self.stats = env_mat_stat.stats + mean, stddev = env_mat_stat() + if not self.set_davg_zero: + paddle.assign( + paddle.to_tensor(mean, dtype=self.mean.dtype, place=env.DEVICE), + self.mean, + ) + if not self.set_stddev_constant: + paddle.assign( + paddle.to_tensor(stddev, dtype=self.stddev.dtype, place=env.DEVICE), + self.stddev, + ) + + def get_stats(self) -> dict[str, StatItem]: + """Get the statistics of the descriptor.""" + if self.stats is None: + raise RuntimeError( + "The statistics of the descriptor has not been computed." + ) + return self.stats + + def has_message_passing(self) -> bool: + """Returns whether the descriptor block has message passing.""" + return True + + def need_sorted_nlist_for_lower(self) -> bool: + """Returns whether the descriptor block needs sorted nlist when using `forward_lower`.""" + return True diff --git a/deepmd/pd/model/descriptor/repformer_layer.py b/deepmd/pd/model/descriptor/repformer_layer.py index a09c5cbe17..b4d93d8301 100644 --- a/deepmd/pd/model/descriptor/repformer_layer.py +++ b/deepmd/pd/model/descriptor/repformer_layer.py @@ -110,7 +110,7 @@ def _make_nei_g1( # index: nb x (nloc x nnei) x ng1 index = nlist.reshape([nb, nloc * nnei]).unsqueeze(-1).expand([-1, -1, ng1]) # gg1 : nb x (nloc x nnei) x ng1 - gg1 = paddle.take_along_axis(g1_ext, axis=1, indices=index) + gg1 = paddle.take_along_axis(g1_ext, indices=index, axis=1, broadcast=False) # gg1 : nb x nloc x nnei x ng1 gg1 = gg1.reshape([nb, nloc, nnei, ng1]) return gg1 @@ -163,7 +163,7 @@ def __init__( attnw_shift: float = 20.0, precision: str = "float64", seed: Optional[Union[int, list[int]]] = None, - ): + ) -> None: """Return neighbor-wise multi-head self-attention maps, with gate mechanism.""" super().__init__() self.input_dim = input_dim @@ -990,7 +990,7 @@ def _cal_hg( else: invnnei = paddle.rsqrt( float(nnei) - * paddle.ones((nb, nloc, 1, 1), dtype=g2.dtype).to(device=g2.place) + * paddle.ones([nb, nloc, 1, 1], dtype=g2.dtype).to(device=g2.place) ) # nb x nloc x 3 x ng2 h2g2 = paddle.matmul(paddle.transpose(h2, [0, 1, 3, 2]), g2) * invnnei @@ -1016,7 +1016,6 @@ def _cal_grrg(h2g2: paddle.Tensor, axis_neuron: int) -> paddle.Tensor: # nb x nloc x 3 x ng2 nb, nloc, _, ng2 = h2g2.shape # nb x nloc x 3 x axis - # h2g2m = paddle.split(h2g2, decomp.sec(h2g2.shape[-1], axis_neuron), axis=-1)[0] h2g2m = h2g2[..., :axis_neuron] # use slice instead of split # nb x nloc x axis x ng2 g1_13 = paddle.matmul(paddle.transpose(h2g2m, [0, 1, 3, 2]), h2g2) / (3.0**1) diff --git a/deepmd/pd/model/descriptor/repformers.py b/deepmd/pd/model/descriptor/repformers.py index aa90bb32c1..32f88dd1d3 100644 --- a/deepmd/pd/model/descriptor/repformers.py +++ b/deepmd/pd/model/descriptor/repformers.py @@ -414,7 +414,6 @@ def forward( if not self.direct_dist: g2, h2 = paddle.split(dmatrix, [1, 3], axis=-1) else: - # g2, h2 = paddle.linalg.norm(diff, axis=-1, keepdim=True), diff g2, h2 = paddle.linalg.norm(diff, axis=-1, keepdim=True), diff g2 = g2 / self.rcut h2 = h2 / self.rcut @@ -437,65 +436,12 @@ def forward( # g1_ext: nb x nall x ng1 if comm_dict is None: assert mapping is not None - g1_ext = paddle.take_along_axis(g1, axis=1, indices=mapping) + g1_ext = paddle.take_along_axis( + g1, axis=1, indices=mapping, broadcast=False + ) else: raise NotImplementedError("Not implemented yet") - # has_spin = "has_spin" in comm_dict - # if not has_spin: - # n_padding = nall - nloc - # g1 = paddle.nn.functional.pad( - # g1.squeeze(0), (0, 0, 0, n_padding), value=0.0 - # ) - # real_nloc = nloc - # real_nall = nall - # else: - # # for spin - # real_nloc = nloc // 2 - # real_nall = nall // 2 - # real_n_padding = real_nall - real_nloc - # g1_real, g1_virtual = paddle.split( - # g1, [real_nloc, real_nloc], axis=1 - # ) - # # mix_g1: nb x real_nloc x (ng1 * 2) - # mix_g1 = paddle.concat([g1_real, g1_virtual], axis=2) - # # nb x real_nall x (ng1 * 2) - # g1 = paddle.nn.functional.pad( - # mix_g1.squeeze(0), (0, 0, 0, real_n_padding), value=0.0 - # ) - - # assert "send_list" in comm_dict - # assert "send_proc" in comm_dict - # assert "recv_proc" in comm_dict - # assert "send_num" in comm_dict - # assert "recv_num" in comm_dict - # assert "communicator" in comm_dict - # ret = paddle.ops.deepmd.border_op( - # comm_dict["send_list"], - # comm_dict["send_proc"], - # comm_dict["recv_proc"], - # comm_dict["send_num"], - # comm_dict["recv_num"], - # g1, - # comm_dict["communicator"], - # paddle.to_tensor( - # real_nloc, - # dtype=paddle.int32, - # place=env.DEVICE, - # ), # should be int of c++ - # paddle.to_tensor( - # real_nall - real_nloc, - # dtype=paddle.int32, - # place=env.DEVICE, - # ), # should be int of c++ - # ) - # g1_ext = ret[0].unsqueeze(0) - # if has_spin: - # g1_real_ext, g1_virtual_ext = paddle.split( - # g1_ext, [ng1, ng1], axis=2 - # ) - # g1_ext = concat_switch_virtual( - # g1_real_ext, g1_virtual_ext, real_nloc - # ) + g1, g2, h2 = ll.forward( g1_ext, g2, diff --git a/deepmd/pd/model/descriptor/se_a.py b/deepmd/pd/model/descriptor/se_a.py index b08be0ced6..8d84620caa 100644 --- a/deepmd/pd/model/descriptor/se_a.py +++ b/deepmd/pd/model/descriptor/se_a.py @@ -672,8 +672,8 @@ def enable_compression( place="cpu", ) tensor_data_ii = table_data[net].to(device=env.DEVICE, dtype=self.prec) - self.compress_data[embedding_idx] = tensor_data_ii - self.compress_info[embedding_idx] = info_ii + paddle.assign(tensor_data_ii, self.compress_data[embedding_idx]) + paddle.assign(info_ii, self.compress_info[embedding_idx]) self.compress = True def forward( diff --git a/deepmd/pd/model/descriptor/se_atten.py b/deepmd/pd/model/descriptor/se_atten.py index aff71d23ec..6bec47b12e 100644 --- a/deepmd/pd/model/descriptor/se_atten.py +++ b/deepmd/pd/model/descriptor/se_atten.py @@ -428,7 +428,10 @@ def enable_compression( dtype=self.prec, place="cpu", ) - self.compress_data[0] = table_data[net].to(device=env.DEVICE, dtype=self.prec) + paddle.assign( + table_data[net].to(device=env.DEVICE, dtype=self.prec), + self.compress_data[0], + ) self.compress = True def forward( diff --git a/deepmd/pd/model/network/layernorm.py b/deepmd/pd/model/network/layernorm.py index 9b174a49f3..7cfffe1b93 100644 --- a/deepmd/pd/model/network/layernorm.py +++ b/deepmd/pd/model/network/layernorm.py @@ -57,7 +57,7 @@ def __init__( shape=[num_in], dtype=self.prec, default_initializer=nn.initializer.Assign( - empty_t((num_in,), self.prec), + empty_t([num_in], self.prec), ), ) self.bias = self.create_parameter( diff --git a/deepmd/pd/model/network/mlp.py b/deepmd/pd/model/network/mlp.py index 370b0fa8fa..41286fbbae 100644 --- a/deepmd/pd/model/network/mlp.py +++ b/deepmd/pd/model/network/mlp.py @@ -101,7 +101,7 @@ def __init__( (num_in, num_out), dtype=self.prec, default_initializer=nn.initializer.Assign( - empty_t((num_in, num_out), self.prec) + empty_t([num_in, num_out], self.prec) ), ) random_generator = get_generator(seed) diff --git a/deepmd/pd/train/training.py b/deepmd/pd/train/training.py index 9389b8f4b5..2d36c95e5b 100644 --- a/deepmd/pd/train/training.py +++ b/deepmd/pd/train/training.py @@ -29,13 +29,11 @@ from deepmd.common import ( symlink_prefix_files, ) -from deepmd.dpmodel.utils.learning_rate import ( - LearningRateExp, -) from deepmd.loggers.training import ( format_training_message_per_task, ) from deepmd.pd.loss import ( + EnergyHessianStdLoss, EnergyStdLoss, TaskLoss, ) @@ -61,6 +59,9 @@ SAMPLER_RECORD, enable_prim, ) +from deepmd.pd.utils.learning_rate import ( + LearningRateExp, +) from deepmd.pd.utils.stat import ( make_stat_input, ) @@ -185,7 +186,6 @@ def get_dataloader_and_buffer(_data, _params): if dist.is_available() else 0, # setting to 0 diverges the behavior of its iterator; should be >=1 collate_fn=lambda batch: batch[0], # prevent extra conversion - # pin_memory=True, ) _data_buffered = BufferedIterator(iter(_dataloader)) return _dataloader, _data_buffered @@ -274,8 +274,22 @@ def get_lr(lr_params): else: self.opt_type, self.opt_param = get_opt_param(training_params) + # loss_param_tmp for Hessian activation + loss_param_tmp = None + if not self.multi_task: + loss_param_tmp = config["loss"] + else: + loss_param_tmp = { + model_key: config["loss_dict"][model_key] + for model_key in self.model_keys + } + # Model - self.model = get_model_for_wrapper(model_params, resuming=resuming) + self.model = get_model_for_wrapper( + model_params, + resuming=resuming, + _loss_params=loss_param_tmp, + ) # Loss if not self.multi_task: @@ -632,7 +646,7 @@ def warm_up_linear(step, warmup_steps): self.profiling = training_params.get("profiling", False) self.profiling_file = training_params.get("profiling_file", "timeline.json") - def run(self): + def run(self) -> None: if CINN: from paddle import ( jit, @@ -674,11 +688,15 @@ def run(self): core.nvprof_enable_record_event() def step(_step_id, task_key="Default") -> None: + if self.multi_task: + model_index = dp_random.choice( + np.arange(self.num_model, dtype=np.int_), + p=self.model_prob, + ) + task_key = self.model_keys[model_index] # Paddle Profiler if enable_profiling: core.nvprof_nvtx_push(f"Training step {_step_id}") - - self.wrapper.train() if isinstance(self.lr_exp, dict): _lr = self.lr_exp[task_key] else: @@ -722,7 +740,6 @@ def step(_step_id, task_key="Default") -> None: with nvprof_context(enable_profiling, "Adam update"): self.optimizer.step() - self.scheduler.step() else: @@ -905,24 +922,8 @@ def log_loss_valid(_task_key="Default"): self.wrapper.train() self.t0 = time.time() self.total_train_time = 0.0 - for step_id in range(self.num_steps): - if step_id < self.start_step: - continue - if self.multi_task: - chosen_index_list = dp_random.choice( - np.arange( - self.num_model, dtype=np.int32 - ), # int32 should be enough for # models... - p=np.array(self.model_prob), - size=self.world_size, - replace=True, - ) - assert chosen_index_list.size == self.world_size - model_index = chosen_index_list[self.rank] - model_key = self.model_keys[model_index] - else: - model_key = "Default" - step(step_id, model_key) + for step_id in range(self.start_step, self.num_steps): + step(step_id) if JIT: break @@ -1067,9 +1068,11 @@ def get_data(self, is_train=True, task_key="Default"): continue elif not isinstance(batch_data[key], list): if batch_data[key] is not None: - batch_data[key] = batch_data[key].to(DEVICE) + batch_data[key] = batch_data[key].to(DEVICE, blocking=False) else: - batch_data[key] = [item.to(DEVICE) for item in batch_data[key]] + batch_data[key] = [ + item.to(DEVICE, blocking=False) for item in batch_data[key] + ] # we may need a better way to classify which are inputs and which are labels # now wrapper only supports the following inputs: input_keys = [ @@ -1185,8 +1188,16 @@ def get_additional_data_requirement(_model): return additional_data_requirement +def whether_hessian(loss_params): + loss_type = loss_params.get("type", "ener") + return loss_type == "ener" and loss_params.get("start_pref_h", 0.0) > 0.0 + + def get_loss(loss_params, start_lr, _ntypes, _model): loss_type = loss_params.get("type", "ener") + if whether_hessian(loss_params): + loss_params["starter_learning_rate"] = start_lr + return EnergyHessianStdLoss(**loss_params) if loss_type == "ener": loss_params["starter_learning_rate"] = start_lr return EnergyStdLoss(**loss_params) @@ -1202,8 +1213,14 @@ def get_single_model( return model -def get_model_for_wrapper(_model_params, resuming=False): +def get_model_for_wrapper( + _model_params, + resuming=False, + _loss_params=None, +): if "model_dict" not in _model_params: + if _loss_params is not None and whether_hessian(_loss_params): + _model_params["hessian_mode"] = True _model = get_single_model( _model_params, ) @@ -1212,6 +1229,8 @@ def get_model_for_wrapper(_model_params, resuming=False): model_keys = list(_model_params["model_dict"]) do_case_embd, case_embd_index = get_case_embd_config(_model_params) for _model_key in model_keys: + if _loss_params is not None and whether_hessian(_loss_params[_model_key]): + _model_params["model_dict"][_model_key]["hessian_mode"] = True _model[_model_key] = get_single_model( _model_params["model_dict"][_model_key], ) diff --git a/deepmd/pd/utils/auto_batch_size.py b/deepmd/pd/utils/auto_batch_size.py index 0eb5e46d5f..0431fb80ae 100644 --- a/deepmd/pd/utils/auto_batch_size.py +++ b/deepmd/pd/utils/auto_batch_size.py @@ -22,7 +22,7 @@ def __init__( self, initial_batch_size: int = 1024, factor: float = 2.0, - ): + ) -> None: super().__init__( initial_batch_size=initial_batch_size, factor=factor, diff --git a/deepmd/pd/utils/env_mat_stat.py b/deepmd/pd/utils/env_mat_stat.py index a37a9672f9..4ea4850ba7 100644 --- a/deepmd/pd/utils/env_mat_stat.py +++ b/deepmd/pd/utils/env_mat_stat.py @@ -54,10 +54,8 @@ def compute_stat(self, env_mat: dict[str, paddle.Tensor]) -> dict[str, StatItem] for kk, vv in env_mat.items(): stats[kk] = StatItem( number=vv.numel().item(), - sum=vv.sum().item() if vv.numel().item() != 0 else paddle.zeros([]), - squared_sum=paddle.square(vv).sum().item() - if vv.numel().item() != 0 - else paddle.zeros([]), + sum=vv.sum().item(), + squared_sum=paddle.square(vv).sum().item(), ) return stats diff --git a/deepmd/pd/utils/learning_rate.py b/deepmd/pd/utils/learning_rate.py new file mode 100644 index 0000000000..3502434bc0 --- /dev/null +++ b/deepmd/pd/utils/learning_rate.py @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.utils.learning_rate import ( + LearningRateExp, +) + +__all__ = [ + "LearningRateExp", +] diff --git a/deepmd/pd/utils/neighbor_stat.py b/deepmd/pd/utils/neighbor_stat.py index af39161e98..9bc0079475 100644 --- a/deepmd/pd/utils/neighbor_stat.py +++ b/deepmd/pd/utils/neighbor_stat.py @@ -101,7 +101,7 @@ def forward( # count the number of neighbors if not self.mixed_types: mask = rr2 < self.rcut**2 - nnei = paddle.zeros((nframes, nloc, self.ntypes), dtype=paddle.int64) + nnei = paddle.zeros([nframes, nloc, self.ntypes], dtype=paddle.int64) for ii in range(self.ntypes): nnei[:, :, ii] = paddle.sum( mask & ((extend_atype == ii)[:, None, :]), axis=-1 diff --git a/deepmd/pd/utils/nlist.py b/deepmd/pd/utils/nlist.py index aeba4b6ac5..707cbd125b 100644 --- a/deepmd/pd/utils/nlist.py +++ b/deepmd/pd/utils/nlist.py @@ -26,7 +26,7 @@ def extend_input_and_build_neighbor_list( ): nframes, nloc = atype.shape[:2] if box is not None: - box_gpu = box.to(coord.place) + box_gpu = box coord_normalized = normalize_coord( coord.reshape([nframes, nloc, 3]), box_gpu.reshape([nframes, 3, 3]), @@ -475,7 +475,7 @@ def extend_coord_with_ghosts( nbuff.astype(paddle.int64), ) # 3 - nbuff = paddle.amax(nbuff, axis=0) # faster than paddle.max + nbuff = paddle.amax(nbuff, axis=0) nbuff_cpu = nbuff.cpu() xi = ( paddle.arange(-nbuff_cpu[0], nbuff_cpu[0] + 1, 1).to( diff --git a/deepmd/pd/utils/stat.py b/deepmd/pd/utils/stat.py index e57530ca44..e0abb1b289 100644 --- a/deepmd/pd/utils/stat.py +++ b/deepmd/pd/utils/stat.py @@ -562,7 +562,7 @@ def compute_output_stats_atomic( } # reshape merged data to [nf, nloc, ndim] merged_output = { - kk: merged_output[kk].reshape((*merged_natoms[kk].shape, -1)) + kk: merged_output[kk].reshape([*merged_natoms[kk].shape, -1]) for kk in merged_output } diff --git a/deepmd/pd/utils/utils.py b/deepmd/pd/utils/utils.py index 1872d9ca1d..a756491a8d 100644 --- a/deepmd/pd/utils/utils.py +++ b/deepmd/pd/utils/utils.py @@ -37,27 +37,33 @@ def __init__(self, threshold=3.0): super().__init__() def sigmoid(x): - return 1 / (1 + np.exp(-x)) + return paddle.nn.functional.sigmoid(x) def silu(x): - return x * sigmoid(x) + return paddle.nn.functional.silu(x) def silu_grad(x): sig = sigmoid(x) return sig + x * sig * (1 - sig) - self.threshold = threshold - self.slope = float(silu_grad(threshold)) - self.const = float(silu(threshold)) + self.threshold = paddle.to_tensor(threshold) + self.slope = float(silu_grad(self.threshold)) + self.const = float(silu(self.threshold)) def forward(self, x: paddle.Tensor) -> paddle.Tensor: silu_part = F.silu(x) - mask = x >= self.threshold - if paddle.any(mask): - tanh_part = paddle.tanh(self.slope * (x - self.threshold)) + self.const - return paddle.where(x < self.threshold, silu_part, tanh_part) - else: - return silu_part + + # NOTE: control flow to be fixed in to_static + # mask = x >= self.threshold + # if paddle.any(mask).item(): + # tanh_part = paddle.tanh(self.slope * (x - self.threshold)) + self.const + # return paddle.where(x < self.threshold, silu_part, tanh_part) + # else: + # return silu_part + + # NOTE: workaround + tanh_part = paddle.tanh(self.slope * (x - self.threshold)) + self.const + return paddle.where(x < self.threshold, silu_part, tanh_part) class ActivationFn(paddle.nn.Layer): diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index 9ebcb15f85..56df87f931 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -37,8 +37,7 @@ DescrptDPA3JAX = None if INSTALLED_PD: - # not supported yet - DescrptDPA3PD = None + from deepmd.pd.model.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3PD else: DescrptDPA3PD = None @@ -151,8 +150,7 @@ def skip_pd(self) -> bool: n_multi_edge_message, precision, ) = self.param - # return not INSTALLED_PD or precision == "bfloat16" - return True + return not INSTALLED_PD or precision == "bfloat16" @property def skip_dp(self) -> bool: diff --git a/source/tests/pd/model/test_dpa2.py b/source/tests/pd/model/test_dpa2.py index f441007cad..8f1699db53 100644 --- a/source/tests/pd/model/test_dpa2.py +++ b/source/tests/pd/model/test_dpa2.py @@ -199,7 +199,6 @@ def test_consistency( atol=atol, ) - @unittest.skip("skip jit in paddle temporally") def test_jit( self, ) -> None: @@ -264,12 +263,11 @@ def test_jit( [ True, ], # smooth - ["float64"], # precision + ["float32"], # precision [False, True], # use_econf_tebd [True], # new sub-structures (use_sqrt_nnei, g1_out_conv, g1_out_mlp) ): dtype = PRECISION_DICT[prec] - rtol, atol = get_tols(prec) repinit = RepinitArgs( rcut=self.rcut, @@ -330,4 +328,157 @@ def test_jit( dd0.repformers.stddev = paddle.to_tensor(dstd_2, dtype=dtype).to( device=env.DEVICE ) - model = paddle.jit.to_static(dd0) + dd0.forward = paddle.jit.to_static( + dd0.forward, full_graph=True, backend=None + ) + dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype).to(device=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype="int64").to(device=env.DEVICE), + paddle.to_tensor(self.nlist, dtype="int64").to(device=env.DEVICE), + paddle.to_tensor(self.mapping, dtype="int64").to(device=env.DEVICE), + ) + + @unittest.skipIf( + not paddle.device.is_compiled_with_cinn() or not env.CINN, + "disable by default for CINN compiler take minites to compile", + ) + def test_cinn_compiler( + self, + ) -> None: + env.enable_prim(True) + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + davg_2 = rng.normal(size=(self.nt, nnei // 2, 4)) + dstd_2 = rng.normal(size=(self.nt, nnei // 2, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + riti, + riz, + rp1c, + rp1d, + rp1g, + rp1a, + rp2g, + rp2a, + rph, + rp2gate, + rus, + rpz, + sm, + prec, + ect, + ns, + ) in itertools.product( + ["concat", "strip"], # repinit_tebd_input_mode + [ + True, + ], # repinit_set_davg_zero + [ + True, + ], # repformer_update_g1_has_conv + [ + True, + ], # repformer_update_g1_has_drrd + [ + True, + ], # repformer_update_g1_has_grrg + [ + True, + ], # repformer_update_g1_has_attn + [ + True, + ], # repformer_update_g2_has_g1g1 + [ + True, + ], # repformer_update_g2_has_attn + [ + False, + ], # repformer_update_h2 + [ + True, + ], # repformer_attn2_has_gate + ["res_avg", "res_residual"], # repformer_update_style + [ + True, + ], # repformer_set_davg_zero + [ + True, + ], # smooth + ["float32"], # precision + [False, True], # use_econf_tebd + [True], # new sub-structures (use_sqrt_nnei, g1_out_conv, g1_out_mlp) + ): + dtype = PRECISION_DICT[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, + use_sqrt_nnei=ns, + g1_out_conv=ns, + g1_out_mlp=ns, + ) + + # dpa2 new impl + dd0 = DescrptDPA2( + self.nt, + repinit=repinit, + repformer=repformer, + # kwargs for descriptor + smooth=sm, + exclude_types=[], + add_tebd_to_repinit_out=False, + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repinit.mean = paddle.to_tensor(davg, dtype=dtype).to(device=env.DEVICE) + dd0.repinit.stddev = paddle.to_tensor(dstd, dtype=dtype).to( + device=env.DEVICE + ) + dd0.repformers.mean = paddle.to_tensor(davg_2, dtype=dtype).to( + device=env.DEVICE + ) + dd0.repformers.stddev = paddle.to_tensor(dstd_2, dtype=dtype).to( + device=env.DEVICE + ) + dd0.forward = paddle.jit.to_static( + dd0.forward, full_graph=True, backend="CINN" + ) + dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype).to(device=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype="int64").to(device=env.DEVICE), + paddle.to_tensor(self.nlist, dtype="int64").to(device=env.DEVICE), + paddle.to_tensor(self.mapping, dtype="int64").to(device=env.DEVICE), + ) + env.enable_prim(False) diff --git a/source/tests/pd/model/test_dpa3.py b/source/tests/pd/model/test_dpa3.py new file mode 100644 index 0000000000..4125e51ff0 --- /dev/null +++ b/source/tests/pd/model/test_dpa3.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import paddle + +from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DPDescrptDPA3 +from deepmd.dpmodel.descriptor.dpa3 import ( + RepFlowArgs, +) +from deepmd.pd.model.descriptor import ( + DescrptDPA3, +) +from deepmd.pd.utils import ( + env, +) +from deepmd.pd.utils.env import ( + PRECISION_DICT, +) + +from ...seed import ( + GLOBAL_SEED, +) +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) +from .test_mlp import ( + get_tols, +) + +dtype = env.GLOBAL_PD_FLOAT_PRECISION + + +class TestDescrptDPA3(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self) -> None: + TestCaseSingleFrameWithNlist.setUp(self) + + def test_consistency( + self, + ) -> None: + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + ua, + rus, + ruri, + acr, + acer, + acus, + nme, + prec, + ect, + ) in itertools.product( + [True, False], # update_angle + ["res_residual"], # update_style + ["norm", "const"], # update_residual_init + [0, 1], # a_compress_rate + [1, 2], # a_compress_e_rate + [True, False], # a_compress_use_split + [1, 2], # n_multi_edge_message + ["float64"], # precision + [False], # use_econf_tebd + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + if prec == "float64": + atol = 1e-8 # marginal GPU test cases... + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=8, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei - 1, + a_compress_rate=acr, + a_compress_e_rate=acer, + a_compress_use_split=acus, + n_multi_edge_message=nme, + axis_neuron=4, + update_angle=ua, + update_style=rus, + update_residual_init=ruri, + smooth_edge_update=True, + ) + + # dpa3 new impl + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype, place=env.DEVICE) + dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype, place=env.DEVICE) + rd0, _, _, _, _ = dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), + ) + # serialization + dd1 = DescrptDPA3.deserialize(dd0.serialize()) + rd1, _, _, _, _ = dd1( + paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), + ) + np.testing.assert_allclose( + rd0.numpy(), + rd1.numpy(), + rtol=rtol, + atol=atol, + ) + # dp impl + dd2 = DPDescrptDPA3.deserialize(dd0.serialize()) + rd2, _, _, _, _ = dd2.call( + self.coord_ext, self.atype_ext, self.nlist, self.mapping + ) + np.testing.assert_allclose( + rd0.numpy(), + rd2, + rtol=rtol, + atol=atol, + ) + + def test_jit( + self, + ) -> None: + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + ua, + rus, + ruri, + acr, + acer, + acus, + nme, + prec, + ect, + ) in itertools.product( + [True], # update_angle + ["res_residual"], # update_style + ["const"], # update_residual_init + [0, 1], # a_compress_rate + [2], # a_compress_e_rate + [True], # a_compress_use_split + [1, 2], # n_multi_edge_message + ["float64"], # precision + [False], # use_econf_tebd + ): + dtype = PRECISION_DICT[prec] + rtol, atol = get_tols(prec) + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=8, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei - 1, + a_compress_rate=acr, + a_compress_e_rate=acer, + a_compress_use_split=acus, + n_multi_edge_message=nme, + axis_neuron=4, + update_angle=ua, + update_style=rus, + update_residual_init=ruri, + smooth_edge_update=True, + ) + + # dpa3 new impl + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype, place=env.DEVICE) + dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype, place=env.DEVICE) + dd0.forward = paddle.jit.to_static(full_graph=True, backend=None)( + dd0.forward + ) + _ = dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), + ) + + @unittest.skipIf( + not paddle.device.is_compiled_with_cinn() or not env.CINN, + "disable by default for CINN compiler take minites to compile", + ) + def test_cinn_compiler( + self, + ) -> None: + rng = np.random.default_rng(100) + nf, nloc, nnei = self.nlist.shape + davg = rng.normal(size=(self.nt, nnei, 4)) + dstd = rng.normal(size=(self.nt, nnei, 4)) + dstd = 0.1 + np.abs(dstd) + + for ( + ua, + rus, + ruri, + acr, + acer, + acus, + nme, + prec, + ect, + ) in itertools.product( + [True], # update_angle + ["res_residual"], # update_style + ["const"], # update_residual_init + [0, 1], # a_compress_rate + [2], # a_compress_e_rate + [True], # a_compress_use_split + [1, 2], # n_multi_edge_message + ["float32"], # precision + [False], # use_econf_tebd + ): + dtype = PRECISION_DICT[prec] + + repflow = RepFlowArgs( + n_dim=20, + e_dim=10, + a_dim=8, + nlayers=3, + e_rcut=self.rcut, + e_rcut_smth=self.rcut_smth, + e_sel=nnei, + a_rcut=self.rcut - 0.1, + a_rcut_smth=self.rcut_smth, + a_sel=nnei - 1, + a_compress_rate=acr, + a_compress_e_rate=acer, + a_compress_use_split=acus, + n_multi_edge_message=nme, + axis_neuron=4, + update_angle=ua, + update_style=rus, + update_residual_init=ruri, + smooth_edge_update=True, + ) + + # dpa3 new impl + dd0 = DescrptDPA3( + self.nt, + repflow=repflow, + # kwargs for descriptor + exclude_types=[], + precision=prec, + use_econf_tebd=ect, + type_map=["O", "H"] if ect else None, + seed=GLOBAL_SEED, + ).to(env.DEVICE) + + dd0.repflows.mean = paddle.to_tensor(davg, dtype=dtype, place=env.DEVICE) + dd0.repflows.stddev = paddle.to_tensor(dstd, dtype=dtype, place=env.DEVICE) + dd0.forward = paddle.jit.to_static(full_graph=True, backend="CINN")( + dd0.forward + ) + _ = dd0( + paddle.to_tensor(self.coord_ext, dtype=dtype, place=env.DEVICE), + paddle.to_tensor(self.atype_ext, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.nlist, dtype=paddle.int64, place=env.DEVICE), + paddle.to_tensor(self.mapping, dtype=paddle.int64, place=env.DEVICE), + ) diff --git a/source/tests/pd/model/test_permutation.py b/source/tests/pd/model/test_permutation.py index 88672457a9..048cb2ed4b 100644 --- a/source/tests/pd/model/test_permutation.py +++ b/source/tests/pd/model/test_permutation.py @@ -168,6 +168,45 @@ }, } +model_dpa3 = { + "type_map": ["O", "H", "B"], + "descriptor": { + "type": "dpa3", + "repflow": { + "n_dim": 20, + "e_dim": 10, + "a_dim": 8, + "nlayers": 6, + "e_rcut": 6.0, + "e_rcut_smth": 3.0, + "e_sel": 20, + "a_rcut": 4.0, + "a_rcut_smth": 2.0, + "a_sel": 10, + "axis_neuron": 4, + "a_compress_rate": 1, + "a_compress_e_rate": 2, + "a_compress_use_split": True, + "update_angle": True, + "update_style": "res_residual", + "update_residual": 0.1, + "update_residual_init": "const", + "smooth_edge_update": True, + }, + "activation_function": "silut:10.0", + "use_tebd_bias": False, + "precision": "float32", + "concat_output_tebd": False, + }, + "fitting_net": { + "neuron": [24, 24], + "resnet_dt": True, + "precision": "float32", + "activation_function": "silut:10.0", + "seed": 1, + }, +} + model_dpa2tebd = { "type_map": ["O", "H", "B"], "descriptor": { diff --git a/source/tests/pd/model/test_smooth.py b/source/tests/pd/model/test_smooth.py index f907e6f4ee..e67f20bddd 100644 --- a/source/tests/pd/model/test_smooth.py +++ b/source/tests/pd/model/test_smooth.py @@ -21,6 +21,7 @@ from .test_permutation import ( # model_dpau, model_dpa1, model_dpa2, + model_dpa3, model_se_e2_a, ) @@ -220,6 +221,17 @@ def setUp(self) -> None: self.epsilon, self.aprec = None, None +class TestEnergyModelDPA3(unittest.TestCase, SmoothTest): + def setUp(self) -> None: + model_params = copy.deepcopy(model_dpa3) + self.type_split = True + self.model = get_model(model_params).to(env.DEVICE) + # less degree of smoothness, + # error can be systematically removed by reducing epsilon + self.epsilon = 1e-5 + self.aprec = 1e-5 + + # class TestEnergyFoo(unittest.TestCase): # def test(self): # model_params = model_dpau diff --git a/source/tests/pd/test_multitask.py b/source/tests/pd/test_multitask.py index 49c9dbd003..0b85816f7b 100644 --- a/source/tests/pd/test_multitask.py +++ b/source/tests/pd/test_multitask.py @@ -32,6 +32,7 @@ model_dpa1, model_dpa2, model_dpa2tebd, + model_dpa3, model_se_e2_a, ) @@ -393,5 +394,47 @@ def tearDown(self) -> None: MultiTaskTrainTest.tearDown(self) +@unittest.skip( + "Paddle does not support gradient clipping with a mix of float32 and float64 data types." +) +class TestMultiTaskDPA3(unittest.TestCase, MultiTaskTrainTest): + def setUp(self) -> None: + multitask_DPA3 = deepcopy(multitask_template) + multitask_DPA3["model"]["shared_dict"]["my_descriptor"] = model_dpa3[ + "descriptor" + ] + data_file = [str(Path(__file__).parent / "water/data/data_0")] + self.stat_files = "DPA3" + os.makedirs(self.stat_files, exist_ok=True) + self.config = multitask_DPA3 + self.config["training"]["data_dict"]["model_1"]["training_data"]["systems"] = ( + data_file + ) + self.config["training"]["data_dict"]["model_1"]["validation_data"][ + "systems" + ] = data_file + self.config["training"]["data_dict"]["model_1"]["stat_file"] = ( + f"{self.stat_files}/model_1" + ) + self.config["training"]["data_dict"]["model_2"]["training_data"]["systems"] = ( + data_file + ) + self.config["training"]["data_dict"]["model_2"]["validation_data"][ + "systems" + ] = data_file + self.config["training"]["data_dict"]["model_2"]["stat_file"] = ( + f"{self.stat_files}/model_2" + ) + self.config["training"]["numb_steps"] = 1 + self.config["training"]["save_freq"] = 1 + self.origin_config = deepcopy(self.config) + self.config["model"], self.shared_links = preprocess_shared_params( + self.config["model"] + ) + + def tearDown(self) -> None: + MultiTaskTrainTest.tearDown(self) + + if __name__ == "__main__": unittest.main()