From 51643afba3c4457c6715a97f5b990c7126b609ba Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 1 Feb 2024 18:02:50 +0800 Subject: [PATCH 1/7] Feat: add model format for dpa1 --- deepmd/model_format/__init__.py | 6 + deepmd/model_format/dpa1.py | 394 ++++++++++++++ deepmd/model_format/network.py | 266 ++++++++++ deepmd/pt/model/descriptor/dpa1.py | 106 +++- deepmd/pt/model/descriptor/se_a.py | 2 +- deepmd/pt/model/descriptor/se_atten.py | 513 ++++++++++++++++--- deepmd/pt/model/network/mlp.py | 187 ++++++- source/tests/pt/models/{dpa1.pth => dpa1.pt} | Bin 15469 -> 14503 bytes source/tests/pt/models/{dpa2.pth => dpa2.pt} | Bin 179745 -> 171490 bytes source/tests/pt/models/dpa2_tebd.pt | Bin 0 -> 1009 bytes source/tests/pt/models/dpa2_tebd.pth | Bin 1085 -> 0 bytes source/tests/pt/test_descriptor_dpa1.py | 10 +- source/tests/pt/test_descriptor_dpa2.py | 2 +- source/tests/pt/test_dpa1.py | 153 ++++++ source/tests/pt/test_model.py | 2 +- source/tests/pt/test_se_e2_a.py | 2 +- 16 files changed, 1556 insertions(+), 87 deletions(-) create mode 100644 deepmd/model_format/dpa1.py rename source/tests/pt/models/{dpa1.pth => dpa1.pt} (52%) rename source/tests/pt/models/{dpa2.pth => dpa2.pt} (50%) create mode 100644 source/tests/pt/models/dpa2_tebd.pt delete mode 100644 source/tests/pt/models/dpa2_tebd.pth create mode 100644 source/tests/pt/test_dpa1.py diff --git a/deepmd/model_format/__init__.py b/deepmd/model_format/__init__.py index e15f73758e..3aa28ec192 100644 --- a/deepmd/model_format/__init__.py +++ b/deepmd/model_format/__init__.py @@ -14,6 +14,8 @@ EmbeddingNet, FittingNet, NativeLayer, + EmbdLayer, + LayerNorm, NativeNet, NetworkCollection, load_dp_model, @@ -35,10 +37,14 @@ from .se_e2_a import ( DescrptSeA, ) +from .dpa1 import DescrptDPA1 __all__ = [ "InvarFitting", "DescrptSeA", + "DescrptDPA1", + "EmbdLayer", + "LayerNorm", "EnvMat", "make_multilayer_network", "make_embedding_network", diff --git a/deepmd/model_format/dpa1.py b/deepmd/model_format/dpa1.py new file mode 100644 index 0000000000..829339838f --- /dev/null +++ b/deepmd/model_format/dpa1.py @@ -0,0 +1,394 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np + +try: + from deepmd._version import version as __version__ +except ImportError: + __version__ = "unknown" + +import copy +from typing import ( + Any, + List, + Optional, +) + +from .common import ( + DEFAULT_PRECISION, + NativeOP, +) +from .env_mat import ( + EnvMat, +) +from .network import ( + EmbeddingNet, + NetworkCollection, + EmbdLayer, +) + + +class DescrptDPA1(NativeOP): + r"""Attention-based descriptor :math:`\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}`, + which is proposed in pretrainable DPA-1[1] model, is given by + + .. math:: + \mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<, + + where :math:`\hat{\mathcal{G}}^i` represents the embedding matrix:math:`\mathcal{G}^i` + after additional self-attention mechanism and :math:`\mathcal{R}^i` is defined by the full case in the se_e2_a descriptor. + Note that we obtain :math:`\mathcal{G}^i` using the type embedding method by default in this descriptor. + + To perform the self-attention mechanism, the queries :math:`\mathcal{Q}^{i,l} \in \mathbb{R}^{N_c\times d_k}`, + keys :math:`\mathcal{K}^{i,l} \in \mathbb{R}^{N_c\times d_k}`, + and values :math:`\mathcal{V}^{i,l} \in \mathbb{R}^{N_c\times d_v}` are first obtained: + + .. math:: + \left(\mathcal{Q}^{i,l}\right)_{j}=Q_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right), + + .. math:: + \left(\mathcal{K}^{i,l}\right)_{j}=K_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right), + + .. math:: + \left(\mathcal{V}^{i,l}\right)_{j}=V_{l}\left(\left(\mathcal{G}^{i,l-1}\right)_{j}\right), + + where :math:`Q_{l}`, :math:`K_{l}`, :math:`V_{l}` represent three trainable linear transformations + that output the queries and keys of dimension :math:`d_k` and values of dimension :math:`d_v`, and :math:`l` + is the index of the attention layer. + The input embedding matrix to the attention layers, denoted by :math:`\mathcal{G}^{i,0}`, + is chosen as the two-body embedding matrix. + + Then the scaled dot-product attention method is adopted: + + .. math:: + A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})=\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right)\mathcal{V}^{i,l}, + + where :math:`\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) \in \mathbb{R}^{N_c\times N_c}` is attention weights. + In the original attention method, + one typically has :math:`\varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}\right)=\mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right)`, + with :math:`\sqrt{d_{k}}` being the normalization temperature. + This is slightly modified to incorporate the angular information: + + .. math:: + \varphi\left(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l},\mathcal{R}^{i,l}\right) = \mathrm{softmax}\left(\frac{\mathcal{Q}^{i,l} (\mathcal{K}^{i,l})^{T}}{\sqrt{d_{k}}}\right) \odot \hat{\mathcal{R}}^{i}(\hat{\mathcal{R}}^{i})^{T}, + + where :math:`\hat{\mathcal{R}}^{i} \in \mathbb{R}^{N_c\times 3}` denotes normalized relative coordinates, + :math:`\hat{\mathcal{R}}^{i}_{j} = \frac{\boldsymbol{r}_{ij}}{\lVert \boldsymbol{r}_{ij} \lVert}` + and :math:`\odot` means element-wise multiplication. + + Then layer normalization is added in a residual way to finally obtain the self-attention local embedding matrix + :math:`\hat{\mathcal{G}}^{i} = \mathcal{G}^{i,L_a}` after :math:`L_a` attention layers:[^1] + + .. math:: + \mathcal{G}^{i,l} = \mathcal{G}^{i,l-1} + \mathrm{LayerNorm}(A(\mathcal{Q}^{i,l}, \mathcal{K}^{i,l}, \mathcal{V}^{i,l}, \mathcal{R}^{i,l})). + + Parameters + ---------- + rcut + The cut-off radius :math:`r_c` + rcut_smth + From where the environment matrix should be smoothed :math:`r_s` + sel : list[str] + sel[i] specifies the maxmum number of type i atoms in the cut-off radius + ntypes : int + Number of element types + neuron : list[int] + Number of neurons in each hidden layers of the embedding net :math:`\mathcal{N}` + axis_neuron + Number of the axis neuron :math:`M_2` (number of columns of the sub-matrix of the embedding matrix) + tebd_dim: int + Dimension of the type embedding + tebd_input_mode: str + The way to mix the type embeddings. Supported options are `concat`, `dot_residual_s`. + resnet_dt + Time-step `dt` in the resnet construction: + y = x + dt * \phi (Wx + b) + trainable + If the weights of embedding net are trainable. + type_one_side + Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets + attn: int + Hidden dimension of the attention vectors + attn_layer: int + Number of attention layers + attn_dotr: bool + If dot the angular gate to the attention weights + attn_mask: bool + If mask the diagonal of attention weights + exclude_types : List[List[int]] + 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. + set_davg_zero + Set the shift of embedding net input to zero. + activation_function + The activation function in the embedding net. Supported options are |ACTIVATION_FN| + precision + The precision of the embedding net parameters. Supported options are |PRECISION| + scaling_factor: float + The scaling factor of normalization in calculations of attention weights. + If `temperature` is None, the scaling of attention weights is (N_dim * scaling_factor)**0.5 + temperature: Optional[float] + If not None, the scaling of attention weights is `temperature` itself. + spin + The deepspin object. + + Limitations + ----------- + The currently implementation does not support the following features + + 1. type_one_side == False + 2. exclude_types != [] + 3. spin is not None + 4. tebd_input_mode != 'concat' + 5. smooth == True + + References + ---------- + .. [1] Duo Zhang, Hangrui Bi, Fu-Zhi Dai, Wanrun Jiang, Linfeng Zhang, and Han Wang. 2022. + DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. + arXiv preprint arXiv:2208.08236. + """ + def __init__( + self, + rcut: float, + rcut_smth: float, + sel: List[str], + ntypes: int, + neuron: List[int] = [25, 50, 100], + axis_neuron: int = 8, + tebd_dim: int = 8, + tebd_input_mode: str = "concat", + resnet_dt: bool = False, + trainable: bool = True, + type_one_side: bool = True, + attn: int = 128, + attn_layer: int = 2, + attn_dotr: bool = True, + attn_mask: bool = False, + exclude_types: List[List[int]] = [], + set_davg_zero: bool = False, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + scaling_factor=1.0, + normalize=True, + temperature=None, + smooth: bool = True, + concat_output_tebd: bool = True, + spin: Optional[Any] = None, + ) -> None: + ## seed, uniform_seed, multi_task, not included. + if not type_one_side: + raise NotImplementedError("type_one_side == False not implemented") + if exclude_types != []: + raise NotImplementedError("exclude_types is not implemented") + if spin is not None: + raise NotImplementedError("spin is not implemented") + # TODO + if tebd_input_mode != 'concat': + raise NotImplementedError("tebd_input_mode != 'concat' not implemented") + if not smooth: + raise NotImplementedError("smooth == False not implemented") + + self.rcut = rcut + self.rcut_smth = rcut_smth + if isinstance(sel, int): + sel = [sel] + self.sel = sel + self.ntypes = ntypes + self.neuron = neuron + self.axis_neuron = axis_neuron + self.tebd_dim = tebd_dim + self.tebd_input_mode = tebd_input_mode + self.resnet_dt = resnet_dt + self.trainable = trainable + self.type_one_side = type_one_side + self.attn = attn + self.attn_layer = attn_layer + self.attn_dotr = attn_dotr + self.attn_mask = attn_mask + self.exclude_types = exclude_types + self.set_davg_zero = set_davg_zero + self.activation_function = activation_function + self.precision = precision + self.scaling_factor = scaling_factor + self.normalize = normalize + self.temperature = temperature + self.concat_output_tebd = concat_output_tebd + self.spin = spin + + self.type_embedding = EmbdLayer(ntypes, tebd_dim, padding=True, precision=precision) + in_dim = 1 + self.tebd_dim * 2 if self.tebd_input_mode in ['concat'] else 1 + self.embeddings = NetworkCollection( + ndim=0, + ntypes=self.ntypes, + network_type="embedding_network", + ) + self.embeddings[0] = EmbeddingNet( + in_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + ) + # self.dpa1_attention = NeighborGatedAttention + self.env_mat = EnvMat(self.rcut, self.rcut_smth) + self.nnei = np.sum(self.sel) + self.davg = np.zeros([self.ntypes, self.nnei, 4]) + self.dstd = np.ones([self.ntypes, self.nnei, 4]) + self.orig_sel = self.sel + + def __setitem__(self, key, value): + if key in ("avg", "data_avg", "davg"): + self.davg = value + elif key in ("std", "data_std", "dstd"): + self.dstd = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("avg", "data_avg", "davg"): + return self.davg + elif key in ("std", "data_std", "dstd"): + return self.dstd + else: + raise KeyError(key) + + @property + def dim_out(self): + """Returns the output dimension of this descriptor.""" + return self.neuron[-1] * self.axis_neuron + self.tebd_dim * 2 \ + if self.concat_output_tebd else self.neuron[-1] * self.axis_neuron + + def cal_g( + self, + ss, + ll, + ): + nf, nloc, nnei = ss.shape[0:3] + ss = ss.reshape(nf, nloc, nnei, -1) + # nf x nloc x nnei x ng + gg = self.embeddings[ll].call(ss) + return gg + + def call( + self, + coord_ext, + atype_ext, + nlist, + ): + """Compute the descriptor. + + Parameters + ---------- + coord_ext + The extended coordinates of atoms. shape: nf x (nallx3) + atype_ext + The extended aotm types. shape: nf x nall + nlist + The neighbor list. shape: nf x nloc x nnei + + Returns + ------- + descriptor + The descriptor. shape: nf x nloc x (ng x axis_neuron) + gr + The rotationally equivariant and permutationally invariant single particle + representation. shape: nf x nloc x ng x 3 + g2 + The rotationally invariant pair-partical representation. + this descriptor returns None + h2 + The rotationally equivariant pair-partical representation. + this descriptor returns None + sw + The smooth switch function. + """ + + # nf x nloc x nnei x 4 + rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd) + nf, nloc, nnei, _ = rr.shape + + # add type embedding into input + # nf x nall x tebd_dim + atype_embd_ext = self.type_embedding.call(atype_ext) + atype_embd = atype_embd_ext[:, :nloc, :] + # nf x nloc x nnei x tebd_dim + atype_embd_nnei = np.tile(atype_embd[:, :, np.newaxis, :], (1, 1, nnei, 1)) + nlist_mask = nlist != -1 + nlist_masked = np.copy(nlist) + nlist_masked[nlist_masked == -1] = 0 + index = np.tile(nlist_masked.reshape(nf, -1, 1), (1, 1, self.tebd_dim)) + # nf x nloc x nnei x tebd_dim + atype_embd_nlist = np.take_along_axis(atype_embd_ext, index, axis=1).reshape(nf, nloc, nnei, self.tebd_dim) + ng = self.neuron[-1] + ss = rr[..., 0:1] + ss = np.concatenate([ss, atype_embd_nlist, atype_embd_nnei], axis=-1) + + # calculate gg + gg = self.cal_g(ss, 0) + # nf x nloc x ng x 4 + gr = np.einsum("flni,flnj->flij", gg, rr) + # nf x nloc x ng x 4 + gr /= self.nnei + gr1 = gr[:, :, : self.axis_neuron, :] + # nf x nloc x ng x ng1 + grrg = np.einsum("flid,fljd->flij", gr, gr1) + # nf x nloc x (ng x ng1) + grrg = grrg.reshape(nf, nloc, ng * self.axis_neuron) + if self.concat_output_tebd: + grrg = np.concatenate([grrg, atype_embd], axis=-1) + return grrg, gr[..., 1:], None, None, ww + + def serialize(self) -> dict: + """Serialize the descriptor to dict.""" + return { + "rcut": self.rcut, + "rcut_smth": self.rcut_smth, + "sel": self.sel, + "ntypes": self.ntypes, + "neuron": self.neuron, + "axis_neuron": self.axis_neuron, + "tebd_dim": self.tebd_dim, + "tebd_input_mode": self.tebd_input_mode, + "resnet_dt": self.resnet_dt, + "trainable": self.trainable, + "type_one_side": self.type_one_side, + "exclude_types": self.exclude_types, + "set_davg_zero": self.set_davg_zero, + "attn": self.attn, + "attn_layer": self.attn_layer, + "attn_dotr": self.attn_dotr, + "attn_mask": self.attn_mask, + "activation_function": self.activation_function, + "precision": self.precision, + "spin": self.spin, + "scaling_factor": self.scaling_factor, + "normalize": self.normalize, + "temperature": self.temperature, + "concat_output_tebd": self.concat_output_tebd, + "embeddings": self.embeddings.serialize(), + # "attention_layers": self.dpa1_attention.serialize(), + "env_mat": self.env_mat.serialize(), + "type_embedding": self.type_embedding.serialize(), + "@variables": { + "davg": self.davg, + "dstd": self.dstd, + }, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA1": + """Deserialize from dict.""" + data = copy.deepcopy(data) + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + type_embedding = data.pop("type_embedding") + attention_layers = data.pop("attention_layers") + env_mat = data.pop("env_mat") + obj = cls(**data) + obj["davg"] = variables["davg"] + obj["dstd"] = variables["dstd"] + obj.type_embedding = EmbdLayer.deserialize(type_embedding) + obj.embeddings = NetworkCollection.deserialize(embeddings) + obj.env_mat = EnvMat.deserialize(env_mat) + # obj.dpa1_attention = NeighborGatedAttention.deserialize(attention_layers) + return obj diff --git a/deepmd/model_format/network.py b/deepmd/model_format/network.py index f2056c0b95..508eb07d56 100644 --- a/deepmd/model_format/network.py +++ b/deepmd/model_format/network.py @@ -322,6 +322,272 @@ def fn(x): return y +class EmbdLayer(NativeLayer): + """Implementation of embedding layer. + + Parameters + ---------- + w : np.ndarray, optional + The embedding weights of the layer. + padding : bool, optional + Whether the embedding layer need to add one padding in the last channel. + """ + + def __init__( + self, + num_channel, + num_out, + padding: bool = True, + precision: str = DEFAULT_PRECISION, + ) -> None: + self.padding = padding + self.num_channel = num_channel + 1 if self.padding else num_channel + super().__init__(num_in=self.num_channel, + num_out=num_out, + bias=False, + use_timestep=False, + activation_function=None, + resnet=False, + precision=precision, + ) + if self.padding: + self.w[-1] = 0. + + def serialize(self) -> dict: + """Serialize the layer to a dict. + + Returns + ------- + dict + The serialized layer. + """ + data = { + "w": self.w + } + return { + "padding": self.padding, + "precision": self.precision, + "@variables": data, + } + + @classmethod + def deserialize(cls, data: dict) -> "EmbdLayer": + """Deserialize the layer from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = copy.deepcopy(data) + variables = data.pop("@variables") + padding = data.pop("padding") + assert variables["w"] is not None and len(variables["w"].shape) == 2 + num_channel, num_out = variables["w"].shape + obj = cls( + num_channel, + num_out, + padding=False, + **data, + ) + obj.w, = ( + variables["w"], + ) + obj.padding = padding + obj.check_shape_consistency() + return obj + + def __setitem__(self, key, value): + if key in ("w", "matrix"): + self.w = value + elif key == "precision": + self.precision = value + elif key == "padding": + self.padding = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("w", "matrix"): + return self.w + elif key == "precision": + return self.precision + elif key == "padding": + return self.padding + else: + raise KeyError(key) + + def dim_channel(self) -> int: + return self.w.shape[0] + + def call(self, x: np.ndarray) -> np.ndarray: + """Forward pass. + + Parameters + ---------- + x : np.ndarray + The input. + + Returns + ------- + np.ndarray + The output. + """ + if self.w is None: + raise ValueError("w must be set") + y = np.take(self.w, x, axis=0) + return y + + +class LayerNorm(NativeLayer): + """Implementation of Layer Normalization layer. + + Parameters + ---------- + w : np.ndarray, optional + The learnable weights of the normalization scale in the layer. + b : np.ndarray, optional + The learnable biases of the normalization shift in the layer. + eps : float, optional + A small value added to prevent division by zero in calculations. + uni_init : bool, optional + If initialize the weights to be zeros and ones. + """ + + def __init__( + self, + num_in, + eps: float = 1e-5, + uni_init: bool = True, + precision: str = DEFAULT_PRECISION, + ) -> None: + self.eps = eps + self.uni_init = uni_init + self.num_in = num_in + super().__init__(num_in=1, + num_out=num_in, + bias=True, + use_timestep=False, + activation_function=None, + resnet=False, + precision=precision, + ) + self.w = self.w.squeeze(0) # keep the weight shape to be [num_in] + if self.uni_init: + self.w = 1. + self.b = 0. + + def serialize(self) -> dict: + """Serialize the layer to a dict. + + Returns + ------- + dict + The serialized layer. + """ + data = { + "w": self.w, + "b": self.b, + } + return { + "eps": self.eps, + "precision": self.precision, + "@variables": data, + } + + @classmethod + def deserialize(cls, data: dict) -> "LayerNorm": + """Deserialize the layer from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + data = copy.deepcopy(data) + variables = data.pop("@variables") + if variables["w"] is not None: + assert len(variables["w"].shape) == 1 + if variables["b"] is not None: + assert len(variables["b"].shape) == 1 + num_in, = variables["w"].shape + obj = cls( + num_in, + **data, + ) + obj.w, = ( + variables["w"], + ) + obj.b, = ( + variables["b"], + ) + obj._check_shape_consistency() + return obj + + def _check_shape_consistency(self): + if self.b is not None and self.w.shape[0] != self.b.shape[0]: + raise ValueError( + f"dim 1 of w {self.w.shape[0]} is not equal to shape " + f"of b {self.b.shape[0]}", + ) + + def __setitem__(self, key, value): + if key in ("w", "matrix"): + self.w = value + elif key in ("b", "bias"): + self.b = value + elif key == "precision": + self.precision = value + elif key == "eps": + self.eps = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ("w", "matrix"): + return self.w + elif key in ("b", "bias"): + return self.b + elif key == "precision": + return self.precision + elif key == "eps": + return self.eps + else: + raise KeyError(key) + + def dim_out(self) -> int: + return self.w.shape[0] + + def call(self, x: np.ndarray) -> np.ndarray: + """Forward pass. + + Parameters + ---------- + x : np.ndarray + The input. + + Returns + ------- + np.ndarray + The output. + """ + if self.w is None or self.b is None: + raise ValueError("w/b must be set") + y = self.layer_norm_numpy(x, tuple((self.num_in,)), self.w, self.b, self.eps) + return y + + @staticmethod + def layer_norm_numpy(x, shape, weight, bias, eps): + # mean and variance + mean = np.mean(x, axis=tuple(range(-len(shape), 0)), keepdims=True) + var = np.var(x, axis=tuple(range(-len(shape), 0)), keepdims=True) + # normalize + x_normalized = (x - mean) / np.sqrt(var + eps) + # shift and scale + x_ln = x_normalized * weight + bias + return x_ln + + def make_multilayer_network(T_NetworkLayer, ModuleBase): class NN(ModuleBase): """Native representation of a neural network. diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 23f521b6d8..611c9b1179 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -13,9 +13,12 @@ TypeEmbedNet, ) -from .se_atten import ( - DescrptBlockSeAtten, +from .se_atten import DescrptBlockSeAtten, NeighborGatedAttention +from deepmd.pt.model.network.mlp import EmbdLayer, NetworkCollection +from deepmd.model_format import ( + EnvMat as DPEnvMat, ) +from deepmd.pt.utils import env @Descriptor.register("dpa1") @@ -37,17 +40,16 @@ def __init__( attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, - post_ln=True, - ffn=False, - ffn_embed_dim=1024, - activation="tanh", + activation_function="tanh", + precision: str = "float64", + resnet_dt: bool = False, scaling_factor=1.0, - head_num=1, normalize=True, temperature=None, - return_rot=False, concat_output_tebd: bool = True, type: Optional[str] = None, + old_impl: bool = False, + **kwargs, ): super().__init__() del type @@ -65,17 +67,22 @@ def __init__( attn_layer=attn_layer, attn_dotr=attn_dotr, attn_mask=attn_mask, - post_ln=post_ln, - ffn=ffn, - ffn_embed_dim=ffn_embed_dim, - activation=activation, + activation_function=activation_function, + precision=precision, + resnet_dt=resnet_dt, scaling_factor=scaling_factor, - head_num=head_num, normalize=normalize, temperature=temperature, - return_rot=return_rot, + old_impl=old_impl, + **kwargs ) - self.type_embedding = TypeEmbedNet(ntypes, tebd_dim) + self.type_embedding_old = None + self.type_embedding = None + self.old_impl = old_impl + if self.old_impl: + self.type_embedding_old = TypeEmbedNet(ntypes, tebd_dim) + else: + self.type_embedding = EmbdLayer(ntypes, tebd_dim, padding=True, precision=precision) self.tebd_dim = tebd_dim self.concat_output_tebd = concat_output_tebd @@ -168,7 +175,12 @@ def forward( del mapping nframes, nloc, nnei = nlist.shape nall = extended_coord.view(nframes, -1).shape[1] // 3 - g1_ext = self.type_embedding(extended_atype) + if self.old_impl: + assert self.type_embedding_old is not None + g1_ext = self.type_embedding_old(extended_atype) + else: + assert self.type_embedding is not None + g1_ext = self.type_embedding(extended_atype) g1_inp = g1_ext[:, :nloc, :] g1, g2, h2, rot_mat, sw = self.se_atten( nlist, @@ -181,3 +193,65 @@ def forward( g1 = torch.cat([g1, g1_inp], dim=-1) return g1, rot_mat, g2, h2, sw + + def set_stat_mean_and_stddev( + self, + mean: torch.Tensor, + stddev: torch.Tensor, + ) -> None: + self.se_atten.mean = mean + self.se_atten.stddev = stddev + + def serialize(self) -> dict: + obj = self.se_atten + return { + "rcut": obj.rcut, + "rcut_smth": obj.rcut_smth, + "sel": obj.sel, + "ntypes": obj.ntypes, + "neuron": obj.neuron, + "axis_neuron": obj.axis_neuron, + "tebd_dim": obj.tebd_dim, + "tebd_input_mode": obj.tebd_input_mode, + "set_davg_zero": obj.set_davg_zero, + "attn": obj.attn_dim, + "attn_layer": obj.attn_layer, + "attn_dotr": obj.attn_dotr, + "attn_mask": obj.attn_mask, + "activation_function": obj.activation_function, + "precision": obj.precision, + "resnet_dt": obj.resnet_dt, + "scaling_factor": obj.scaling_factor, + "normalize": obj.normalize, + "temperature": obj.temperature, + "concat_output_tebd": self.concat_output_tebd, + "embeddings": obj.filter_layers.serialize(), + "attention_layers": obj.dpa1_attention.serialize(), + "env_mat": DPEnvMat(obj.rcut, obj.rcut_smth).serialize(), + "type_embedding": self.type_embedding.serialize(), + "@variables": { + "davg": obj["davg"].detach().cpu().numpy(), + "dstd": obj["dstd"].detach().cpu().numpy(), + }, + ## to be updated when the options are supported. + "trainable": True, + "type_one_side": True, + "exclude_types": [], + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "DescrptDPA1": + variables = data.pop("@variables") + embeddings = data.pop("embeddings") + type_embedding = data.pop("type_embedding") + attention_layers = data.pop("attention_layers") + env_mat = data.pop("env_mat") + obj = cls(**data) + t_cvt = lambda xx: torch.tensor(xx, dtype=obj.se_atten.prec, device=env.DEVICE) + obj.type_embedding = EmbdLayer.deserialize(type_embedding) + obj.se_atten["davg"] = t_cvt(variables["davg"]) + obj.se_atten["dstd"] = t_cvt(variables["dstd"]) + obj.se_atten.filter_layers = NetworkCollection.deserialize(embeddings) + obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize(attention_layers) + return obj diff --git a/deepmd/pt/model/descriptor/se_a.py b/deepmd/pt/model/descriptor/se_a.py index 3f42736dca..0ece6a61c6 100644 --- a/deepmd/pt/model/descriptor/se_a.py +++ b/deepmd/pt/model/descriptor/se_a.py @@ -451,7 +451,7 @@ def forward( xyz_scatter = torch.zeros( [nfnl, 4, self.filter_neuron[-1]], dtype=self.prec, device=env.DEVICE ) - for ii, ll in enumerate(self.filter_layers.networks): + for ii, ll in enumerate(self.filter_layers._networks): # nfnl x nt x 4 rr = dmatrix[:, self.sec[ii] : self.sec[ii + 1], :] ss = rr[:, :, :1] diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 78cba59da7..ba7c533ff7 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -6,6 +6,13 @@ import numpy as np import torch +import torch.nn as nn +import torch.nn.functional as torch_func +from deepmd.pt.utils.env import ( + PRECISION_DICT, + DEFAULT_PRECISION, +) +from deepmd.pt.utils.utils import ActivationFn from deepmd.pt.model.descriptor.descriptor import ( DescriptorBlock, @@ -18,6 +25,11 @@ NeighborWiseAttention, TypeFilter, ) +from deepmd.pt.model.network.mlp import EmbeddingNet, NetworkCollection, MLPLayer, LayerNorm + +from deepmd.model_format import ( + EnvMat as DPEnvMat, +) from deepmd.pt.utils import ( env, ) @@ -34,23 +46,22 @@ def __init__( neuron: list = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, - tebd_input_mode: str = "concat", + tebd_input_mode: str = 'concat', # set_davg_zero: bool = False, set_davg_zero: bool = True, # TODO attn: int = 128, attn_layer: int = 2, attn_dotr: bool = True, attn_mask: bool = False, - post_ln=True, - ffn=False, - ffn_embed_dim=1024, - activation="tanh", + activation_function="tanh", + precision: str = "float64", + resnet_dt: bool = False, scaling_factor=1.0, - head_num=1, normalize=True, temperature=None, - return_rot=False, type: Optional[str] = None, + old_impl: bool = False, + **kwargs, ): """Construct an embedding net of type `se_atten`. @@ -65,6 +76,8 @@ def __init__( del type self.rcut = rcut self.rcut_smth = rcut_smth + self.neuron = neuron + self.filter_neuron = self.neuron self.filter_neuron = neuron self.axis_neuron = axis_neuron self.tebd_dim = tebd_dim @@ -74,15 +87,14 @@ def __init__( self.attn_layer = attn_layer self.attn_dotr = attn_dotr self.attn_mask = attn_mask - self.post_ln = post_ln - self.ffn = ffn - self.ffn_embed_dim = ffn_embed_dim - self.activation = activation + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + self.resnet_dt = resnet_dt self.scaling_factor = scaling_factor - self.head_num = head_num self.normalize = normalize self.temperature = temperature - self.return_rot = return_rot + self.old_impl = old_impl if isinstance(sel, int): sel = [sel] @@ -93,22 +105,24 @@ def __init__( self.split_sel = self.sel self.nnei = sum(sel) self.ndescrpt = self.nnei * 4 - self.dpa1_attention = NeighborWiseAttention( - self.attn_layer, - self.nnei, - self.filter_neuron[-1], - self.attn_dim, - dotr=self.attn_dotr, - do_mask=self.attn_mask, - post_ln=self.post_ln, - ffn=self.ffn, - ffn_embed_dim=self.ffn_embed_dim, - activation=self.activation, - scaling_factor=self.scaling_factor, - head_num=self.head_num, - normalize=self.normalize, - temperature=self.temperature, - ) + if self.old_impl: + self.dpa1_attention = NeighborWiseAttention(self.attn_layer, self.nnei, self.filter_neuron[-1], + self.attn_dim, + dotr=self.attn_dotr, do_mask=self.attn_mask, + activation=self.activation_function, + scaling_factor=self.scaling_factor, + normalize=self.normalize, + temperature=self.temperature) + else: + self.dpa1_attention = NeighborGatedAttention(self.attn_layer, + self.nnei, + self.filter_neuron[-1], + self.attn_dim, + dotr=self.attn_dotr, + do_mask=self.attn_mask, + scaling_factor=self.scaling_factor, + normalize=self.normalize, + temperature=self.temperature) wanted_shape = (self.ntypes, self.nnei, 4) mean = torch.zeros( @@ -119,19 +133,26 @@ def __init__( ) self.register_buffer("mean", mean) self.register_buffer("stddev", stddev) + self.embd_input_dim = 1 + self.tebd_dim * 2 if self.tebd_input_mode in ['concat'] else 1 + self.filter_layers_old = None + self.filter_layers = None - filter_layers = [] - one = TypeFilter( - 0, - self.nnei, - self.filter_neuron, - return_G=True, - tebd_dim=self.tebd_dim, - use_tebd=True, - tebd_mode=self.tebd_input_mode, - ) - filter_layers.append(one) - self.filter_layers = torch.nn.ModuleList(filter_layers) + if self.old_impl: + filter_layers = [] + one = TypeFilter(0, self.nnei, self.filter_neuron, return_G=True, tebd_dim=self.tebd_dim, use_tebd=True, + tebd_mode=self.tebd_input_mode) + filter_layers.append(one) + self.filter_layers_old = torch.nn.ModuleList(filter_layers) + else: + filter_layers = NetworkCollection(ndim=0, ntypes=len(sel), network_type="embedding_network") + filter_layers[0] = EmbeddingNet( + self.embd_input_dim, + self.filter_neuron, + activation_function=self.activation_function, + precision=self.precision, + resnet_dt=self.resnet_dt, + ) + self.filter_layers = filter_layers def get_rcut(self) -> float: """Returns the cut-off radius.""" @@ -172,6 +193,22 @@ def dim_emb(self): """Returns the output dimension of embedding.""" return self.filter_neuron[-1] + def __setitem__(self, key, value): + 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 compute_input_stats(self, merged): """Update mean and stddev for descriptor elements.""" sumr = [] @@ -282,7 +319,6 @@ def forward( self.rcut_smth, ) # [nfxnlocxnnei, self.ndescrpt] - dmatrix = dmatrix.view(-1, self.ndescrpt) nlist_mask = nlist != -1 nlist[nlist == -1] = 0 sw = torch.squeeze(sw, -1) @@ -300,23 +336,39 @@ def forward( atype_tebd_nlist = torch.gather(atype_tebd_ext, dim=1, index=index) # nb x nloc x nnei x nt atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) - ret = self.filter_layers[0]( - dmatrix, - atype_tebd=atype_tebd_nnei, - nlist_tebd=atype_tebd_nlist, - ) # shape is [nframes*nall, self.neei, out_size] - input_r = torch.nn.functional.normalize( - dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 - ) - ret = self.dpa1_attention( - ret, nlist_mask, input_r=input_r, sw=sw - ) # shape is [nframes*nloc, self.neei, out_size] - inputs_reshape = dmatrix.view(-1, self.nnei, 4).permute( - 0, 2, 1 - ) # shape is [nframes*natoms[0], 4, self.neei] - xyz_scatter = torch.matmul( - inputs_reshape, ret - ) # shape is [nframes*natoms[0], 4, out_size] + if self.old_impl: + assert self.filter_layers_old is not None + dmatrix = dmatrix.view(-1, self.ndescrpt) # shape is [nframes*nall, self.ndescrpt] + gg = self.filter_layers_old[0]( + dmatrix, + atype_tebd=atype_tebd_nnei, + nlist_tebd=atype_tebd_nlist, + ) # shape is [nframes*nall, self.neei, out_size] + input_r = torch.nn.functional.normalize(dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1) + gg = self.dpa1_attention(gg, nlist_mask, input_r=input_r, + sw=sw) # shape is [nframes*nloc, self.neei, out_size] + inputs_reshape = dmatrix.view(-1, self.nnei, 4).permute(0, 2, + 1) # shape is [nframes*natoms[0], 4, self.neei] + xyz_scatter = torch.matmul(inputs_reshape, gg) # shape is [nframes*natoms[0], 4, out_size] + else: + assert self.filter_layers is not None + dmatrix = dmatrix.view(-1, self.nnei, 4) + nfnl = dmatrix.shape[0] + # nfnl x nnei x 4 + rr = dmatrix + ss = rr[:, :, :1] + if self.tebd_input_mode in ['concat']: + nlist_tebd = atype_tebd_nlist.reshape(nfnl, nnei, self.tebd_dim) + atype_tebd = atype_tebd_nnei.reshape(nfnl, nnei, self.tebd_dim) + # nfnl x nnei x (1 + tebd_dim * 2) + ss = torch.concat([ss, nlist_tebd, atype_tebd], dim=2) + # nfnl x nnei x ng + gg = self.filter_layers._networks[0](ss) + input_r = torch.nn.functional.normalize(dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1) + gg = self.dpa1_attention(gg, nlist_mask, input_r=input_r, + sw=sw) # shape is [nframes*nloc, self.neei, out_size] + # nfnl x 4 x ng + xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) xyz_scatter = xyz_scatter / self.nnei xyz_scatter_1 = xyz_scatter.permute(0, 2, 1) rot_mat = xyz_scatter_1[:, :, 1:4] @@ -326,13 +378,354 @@ def forward( ) # shape is [nframes*nloc, self.filter_neuron[-1], self.axis_neuron] return ( result.view(-1, nloc, self.filter_neuron[-1] * self.axis_neuron), - ret.view(-1, nloc, self.nnei, self.filter_neuron[-1]), + gg.view(-1, nloc, self.nnei, self.filter_neuron[-1]), dmatrix.view(-1, nloc, self.nnei, 4)[..., 1:], rot_mat.view(-1, self.filter_neuron[-1], 3), sw, ) +class NeighborGatedAttention(nn.Module): + def __init__(self, + layer_num: int, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention net. + """ + super(NeighborGatedAttention, self).__init__() + self.layer_num = layer_num + self.nnei = nnei + self.embed_dim = embed_dim + self.hidden_dim = hidden_dim + self.dotr = dotr + self.do_mask = do_mask + self.scaling_factor = scaling_factor + self.normalize = normalize + self.temperature = temperature + self.precision = precision + self.network_type = NeighborGatedAttentionLayer + attention_layers = [] + for i in range(self.layer_num): + attention_layers.append(NeighborGatedAttentionLayer(nnei, + embed_dim, + hidden_dim, + dotr=dotr, + do_mask=do_mask, + scaling_factor=scaling_factor, + normalize=normalize, + temperature=temperature, + precision=precision)) + self.attention_layers = nn.ModuleList(attention_layers) + + def forward( + self, + input_G, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + ): + """ + Args: + input_G: Input G, [nframes * nloc, nnei, embed_dim] + nei_mask: neighbor mask, [nframes * nloc, nnei] + input_r: normalized radial, [nframes, nloc, nei, 3] + Returns: + out: Output G, [nframes * nloc, nnei, embed_dim] + """ + out = input_G + # https://github.com/pytorch/pytorch/issues/39165#issuecomment-635472592 + for layer in self.attention_layers: + out = layer(out, nei_mask, input_r=input_r, sw=sw) + return out + + def _convert_key(self, key): + if isinstance(key, int): + idx = key + else: + if isinstance(key, tuple): + pass + elif isinstance(key, str): + key = tuple([int(tt) for tt in key.split("_")[1:]]) + else: + raise TypeError(key) + assert isinstance(key, tuple) + assert len(key) == self.ndim + idx = sum([tt * self.ntypes**ii for ii, tt in enumerate(key)]) + return idx + + def __getitem__(self, key): + return self.attention_layers[self._convert_key(key)] + + def __setitem__(self, key, value): + if isinstance(value, self.network_type): + pass + elif isinstance(value, dict): + value = self.network_type.deserialize(value) + else: + raise TypeError(value) + self.attention_layers[self._convert_key(key)] = value + + def serialize(self) -> dict: + """Serialize the networks to a dict. + Returns + ------- + dict + The serialized networks. + """ + # network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()} + # network_type_name = network_type_map_inv[self.network_type] + return { + "layer_num": self.layer_num, + "nnei": self.nnei, + "embed_dim": self.embed_dim, + "hidden_dim": self.hidden_dim, + "dotr": self.dotr, + "do_mask": self.do_mask, + "scaling_factor": self.scaling_factor, + "normalize": self.normalize, + "temperature": self.temperature, + "precision": self.precision, + "attention_layers": [layer.serialize() for layer in self.attention_layers] + } + + @classmethod + def deserialize(cls, data: dict) -> "NeighborGatedAttention": + """Deserialize the networks from a dict. + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + attention_layers = data.pop("attention_layers") + obj = cls(**data) + for ii, network in enumerate(attention_layers): + obj[ii] = network + return obj + + +class NeighborGatedAttentionLayer(nn.Module): + def __init__(self, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention layer. + """ + super(NeighborGatedAttentionLayer, self).__init__() + self.nnei = nnei + self.embed_dim = embed_dim + self.hidden_dim = hidden_dim + self.dotr = dotr + self.do_mask = do_mask + self.scaling_factor = scaling_factor + self.normalize = normalize + self.temperature = temperature + self.precision = precision + self.attention_layer = GatedAttentionLayer(nnei, + embed_dim, + hidden_dim, + dotr=dotr, + do_mask=do_mask, + scaling_factor=scaling_factor, + normalize=normalize, + temperature=temperature, + precision=precision, + ) + self.attn_layer_norm = LayerNorm(self.embed_dim, precision=precision) + + def forward( + self, + x, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + ): + residual = x + x = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) + x = residual + x + x = self.attn_layer_norm(x) + return x + + def serialize(self) -> dict: + """Serialize the networks to a dict. + Returns + ------- + dict + The serialized networks. + """ + return { + "nnei": self.nnei, + "embed_dim": self.embed_dim, + "hidden_dim": self.hidden_dim, + "dotr": self.dotr, + "do_mask": self.do_mask, + "scaling_factor": self.scaling_factor, + "normalize": self.normalize, + "temperature": self.temperature, + "precision": self.precision, + "attention_layer": self.attention_layer.serialize(), + "attn_layer_norm": self.attn_layer_norm.serialize() + } + + @classmethod + def deserialize(cls, data: dict) -> "NeighborGatedAttentionLayer": + """Deserialize the networks from a dict. + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + attention_layer = data.pop("attention_layer") + attn_layer_norm = data.pop("attn_layer_norm") + obj = cls(**data) + obj.attention_layer = GatedAttentionLayer.deserialize(attention_layer) + obj.attn_layer_norm = LayerNorm.deserialize(attn_layer_norm) + return obj + + +class GatedAttentionLayer(nn.Module): + def __init__(self, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + bias: bool = True, + smooth: bool = True, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention net. + """ + super(GatedAttentionLayer, self).__init__() + self.nnei = nnei + self.embed_dim = embed_dim + self.hidden_dim = hidden_dim + self.dotr = dotr + self.do_mask = do_mask + self.bias = bias + self.smooth = smooth + self.scaling_factor = scaling_factor + self.temperature = temperature + self.precision = precision + if temperature is None: + self.scaling = (self.hidden_dim * scaling_factor) ** -0.5 + else: + self.scaling = temperature + self.normalize = normalize + self.in_proj = MLPLayer(embed_dim, hidden_dim * 3, bias=bias, use_timestep=False, bavg=0., stddev=1., + precision=precision) + self.out_proj = MLPLayer(hidden_dim, embed_dim, bias=bias, use_timestep=False, bavg=0., stddev=1., + precision=precision) + + def forward( + self, + query, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + attnw_shift: float = 20.0, + ): + """ + Args: + query: input G, [nframes * nloc, nnei, embed_dim] + nei_mask: neighbor mask, [nframes * nloc, nnei] + input_r: normalized radial, [nframes, nloc, nei, 3] + Returns: + type_embedding: + """ + q, k, v = self.in_proj(query).chunk(3, dim=-1) + # [nframes * nloc, nnei, hidden_dim] + q = q.view(-1, self.nnei, self.hidden_dim) + k = k.view(-1, self.nnei, self.hidden_dim) + v = v.view(-1, self.nnei, self.hidden_dim) + if self.normalize: + q = torch_func.normalize(q, dim=-1) + k = torch_func.normalize(k, dim=-1) + v = torch_func.normalize(v, dim=-1) + q = q * self.scaling + k = k.transpose(1, 2) + # [nframes * nloc, nnei, nnei] + attn_weights = torch.bmm(q, k) + # [nframes * nloc, nnei] + nei_mask = nei_mask.view(-1, self.nnei) + if self.smooth: + # [nframes * nloc, nnei] + assert sw is not None + sw = sw.view([-1, self.nnei]) + attn_weights = (attn_weights + attnw_shift) * sw[:, :, None] * sw[:, None, :] - attnw_shift + else: + attn_weights = attn_weights.masked_fill(~nei_mask.unsqueeze(1), float("-inf")) + attn_weights = torch_func.softmax(attn_weights, dim=-1) + attn_weights = attn_weights.masked_fill(~nei_mask.unsqueeze(-1), float(0.0)) + if self.smooth: + assert sw is not None + attn_weights = attn_weights * sw[:, :, None] * sw[:, None, :] + if self.dotr: + assert input_r is not None, "input_r must be provided when dotr is True!" + angular_weight = torch.bmm(input_r, input_r.transpose(1, 2)) + attn_weights = attn_weights * angular_weight + o = torch.bmm(attn_weights, v) + output = self.out_proj(o) + return output + + def serialize(self) -> dict: + """Serialize the networks to a dict. + Returns + ------- + dict + The serialized networks. + """ + # network_type_map_inv = {v: k for k, v in self.NETWORK_TYPE_MAP.items()} + # network_type_name = network_type_map_inv[self.network_type] + return { + "nnei": self.nnei, + "embed_dim": self.embed_dim, + "hidden_dim": self.hidden_dim, + "dotr": self.dotr, + "do_mask": self.do_mask, + "scaling_factor": self.scaling_factor, + "normalize": self.normalize, + "temperature": self.temperature, + "bias": self.bias, + "smooth": self.smooth, + "precision": self.precision, + "in_proj": self.in_proj.serialize(), + "out_proj": self.out_proj.serialize() + } + + @classmethod + def deserialize(cls, data: dict) -> "GatedAttentionLayer": + """Deserialize the networks from a dict. + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + in_proj = data.pop("in_proj") + out_proj = data.pop("out_proj") + obj = cls(**data) + obj.in_proj = MLPLayer.deserialize(in_proj) + obj.out_proj = MLPLayer.deserialize(out_proj) + return obj + + def analyze_descrpt(matrix, ndescrpt, natoms, mixed_type=False, real_atype=None): """Collect avg, square avg and count of descriptors in a batch.""" ntypes = natoms.shape[1] - 2 diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index d76abd82f9..5bd9c1a23d 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -8,6 +8,7 @@ import numpy as np import torch import torch.nn as nn +import torch.nn.functional as torch_func from deepmd.pt.utils import ( env, @@ -18,6 +19,8 @@ from deepmd.model_format import ( NativeLayer, ) +from deepmd.model_format import EmbdLayer as DPEmbdLayer +from deepmd.model_format import LayerNorm as DPLayerNorm from deepmd.model_format import NetworkCollection as DPNetworkCollection from deepmd.model_format import ( make_embedding_network, @@ -188,6 +191,188 @@ def check_load_param(ss): return obj +class EmbdLayer(MLPLayer): + def __init__( + self, + num_channel, + num_out, + padding: bool = True, + stddev: float = 1., + precision: str = DEFAULT_PRECISION, + ): + self.padding = padding + self.num_channel = num_channel + 1 if self.padding else num_channel + super().__init__(num_in=self.num_channel, + num_out=num_out, + bias=False, + use_timestep=False, + activation_function=None, + resnet=False, + stddev=stddev, + precision=precision, + ) + if self.padding: + nn.init.zeros_(self.matrix.data[-1]) + + def dim_channel(self) -> int: + return self.matrix.shape[0] + + def forward( + self, + xx: torch.Tensor, + ) -> torch.Tensor: + """One Embedding layer used by DP model. + + Parameters + ---------- + xx: torch.Tensor + The input of index. + + Returns + ------- + yy: torch.Tensor + The output. + """ + yy = torch_func.embedding(xx, self.matrix) + return yy + + def serialize(self) -> dict: + """Serialize the layer to a dict. + + Returns + ------- + dict + The serialized layer. + """ + nl = DPEmbdLayer( + self.matrix.shape[0], + self.matrix.shape[1], + padding=False, + precision=self.precision, + ) + nl.w = self.matrix.detach().cpu().numpy() + data = nl.serialize() + data["padding"] = self.padding + return data + + @classmethod + def deserialize(cls, data: dict) -> "EmbdLayer": + """Deserialize the layer from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + padding = data["padding"] + nl = DPEmbdLayer.deserialize(data) + obj = cls( + nl["matrix"].shape[0], + nl["matrix"].shape[1], + padding=False, + precision=nl["precision"], + ) + obj.padding = padding + prec = PRECISION_DICT[obj.precision] + check_load_param = \ + lambda ss: nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) \ + if nl[ss] is not None else None + obj.matrix = check_load_param("matrix") + return obj + + +class LayerNorm(MLPLayer): + def __init__( + self, + num_in, + eps: float = 1e-5, + uni_init: bool = True, + bavg: float = 0., + stddev: float = 1., + precision: str = DEFAULT_PRECISION, + ): + self.eps = eps + self.uni_init = uni_init + self.num_in = num_in + super().__init__(num_in=1, + num_out=num_in, + bias=True, + use_timestep=False, + activation_function=None, + resnet=False, + bavg=bavg, + stddev=stddev, + precision=precision, + ) + self.matrix = torch.nn.Parameter(self.matrix.squeeze(0)) + if self.uni_init: + nn.init.ones_(self.matrix.data) + nn.init.zeros_(self.bias.data) + + def dim_out(self) -> int: + return self.matrix.shape[0] + + def forward( + self, + xx: torch.Tensor, + ) -> torch.Tensor: + """One Layer Norm used by DP model. + + Parameters + ---------- + xx: torch.Tensor + The input of index. + + Returns + ------- + yy: torch.Tensor + The output. + """ + yy = torch_func.layer_norm(xx, tuple((self.num_in,)), self.matrix, self.bias, self.eps) + return yy + + def serialize(self) -> dict: + """Serialize the layer to a dict. + + Returns + ------- + dict + The serialized layer. + """ + nl = DPLayerNorm( + self.matrix.shape[0], + eps=self.eps, + precision=self.precision, + ) + nl.w = self.matrix.detach().cpu().numpy() + nl.b = self.bias.detach().cpu().numpy() + data = nl.serialize() + return data + + @classmethod + def deserialize(cls, data: dict) -> "LayerNorm": + """Deserialize the layer from a dict. + + Parameters + ---------- + data : dict + The dict to deserialize from. + """ + nl = DPLayerNorm.deserialize(data) + obj = cls( + nl["matrix"].shape[0], + eps=nl["eps"], + precision=nl["precision"], + ) + prec = PRECISION_DICT[obj.precision] + check_load_param = \ + lambda ss: nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) \ + if nl[ss] is not None else None + obj.matrix = check_load_param("matrix") + obj.bias = check_load_param("bias") + return obj + + MLP_ = make_multilayer_network(MLPLayer, nn.Module) @@ -217,4 +402,4 @@ def __init__(self, *args, **kwargs): # init both two base classes DPNetworkCollection.__init__(self, *args, **kwargs) nn.Module.__init__(self) - self.networks = self._networks = torch.nn.ModuleList(self._networks) + self._networks = torch.nn.ModuleList(self._networks) diff --git a/source/tests/pt/models/dpa1.pth b/source/tests/pt/models/dpa1.pt similarity index 52% rename from source/tests/pt/models/dpa1.pth rename to source/tests/pt/models/dpa1.pt index 75acf2fa15d874dc7c63a0fe1bbc3653267b2c4b..74f69b4b6e00cbdd7e7650fd15ef6dad931f3cd6 100644 GIT binary patch delta 2100 zcmb7_%TE(Q9LJ{=y46C#AT=s&pe-ORY#-S2kOFC9LsCtaCMKj8vI+*m78qC_9@hHm z!Gj67ns_ppKs=}k+A}8n0c!Cr_(=3%OduXO5))^-yR@B^BGa^N=JWf0fAjmz?65iW z5Y8U6lz|p3J^?-Wxk4M##aMoP zXeP;Og>+00p8r+&8^^Ht`A3lwBf4}8mU;ld#UYL#81LiQO9SH*bl=DTKRr4=JVCo@ zaYNT_smmX7mU8+E67VWAhu3Zng8B%&&choKnCRWZ{4uqsd5PZXD@?#SY<^oZ7kPLm z0`Kl+&gu6wf1!8cwmd-QN$r8j!uXNhD}^rw_-gP`SToUkboVIQ4O0SKmEx-M#nDo? z;cEfDk;1JYo^LDU-EqUW0(>Vaod1ykZ6bStADUwM0r*~kACz4GN)=bOUC!l!9|ibn zAGti4Tx$Yc&zFlx<%ORG*z;u{xxATNUj_Irt(=@VPcvakfB!qPv26rkTfAa4A2)#F z%**6*hl2W@S1Gw-j8kH#$ysgyR928D6x4#`vgrwRxEScN0g@Mhx3Vn;_Vqlb3 zIE8M%q0y00@L*<%VyE9`-8oL^NR^;>E_dL2TD;}2K?}0+peACao%O zwxKhTJyI4Koolueh^dm%JrznRXkCSvYFTAzAq`UsVpXWhCaV@yXkI}-Rfw_6Dn~1s zYe_*d6)LNhRZ9rv_Hb4~TdiadXC1Q2*G8glbT6jcJKbh8=NhIOB$WQY>hg*6&6Pw- zrcoAet9(-m3Y;QSGPJCUtC3emD+uLow!TS5+IDx_R-tqbB4`p1Cak0I-_tGl7XS_V QGT|Z}em0>Gj+bKp0{EsNSO5S3 delta 2336 zcma*oOH30%7zgm#(soxXv@uXqqwieM8D zCbgsSs>Fmu6QlIxNp2oYw2JS;?KE)Muet)QemenN-_&U&un2C`kLY zH6gA1$HHh$de;ui_N)50Y-6(B(v&P})jlv;0RZQw5~=t?G!eTJUr5@W_KA3WK8k*5 zPoeR;gM4ZKByaO_qAClQMHL{R`UpE{omU3c3Q~|d3vY>PDJU-o zp%lCwfp_)@niZ5mTY^n4?fy<_1g|$}TCpxVz_q@6zx+g{xXO{E&tp|I0>G@KYLo-m#0xaeG+gYlW_RcN3mE z0Hk-W(NYg|RZcC3GZf^(S$Qig>z8=+*Jx|12gj>nS11?&XxQRFV-{Uu=+HyHxd9AT zl{j~xME7sBPIF9&L(39n+NPP(T`BWsHuO|Nw5%UOo8}&7thZ8cv{-K5TWyNCfhthAYD;jeDBQSmRKAxvIo!GmXJ$Iwy|ODUFl%t!z%njm z?Cgqh6U!21+-rDIG(^Q0@x5_s?m{TY0d&G+u~lAXMtrYOZi<^ZTJ?0a)+@jYzBrZU zxyM47FpgE4KiD)Va>6a7ae1wTF|?6p#%t|T<`hY@^8j+V+-R)7iO9>D$-HmWnBjKP znEJ?heYcS23}wdfV@276o0Kthkmdzto=}GKk@GQhkvT`PQLbB&GK)TP-3&*_oX3GC4bG+WN_{!Vc1mQpVsUO`bAE%DkkEp^wZ_`^ocB@R4TF?+iTC$R}+CACT9h1r`5P>a_VQ#88$dNXp-m8e)|muczISF%rlt7r_P=3 zx_h27dwx!C4y-T?hw`ETjCMos1mkjN*9yaOXSY?gLD`kg8(hxrmChdU(yVxKpjzW~ z2Kva@JE8M(r-cn<4YH@&IQMV<4!rG zuEwRzTnOhoCv*>XIjhE31XUO+f>$^vRs>fVRyZf|S#NVWCs#UeuS)mGb-rU}&3qUX zTm(C3WO+qpKl$`}+BWU)GD)JkJyqsv*#1~0FWS2$-@$ZYUg zeDK*W=bTDsU9;4)hb4!N>zu0}BF{OmR;rL)&U#W0(n6B8Y2^}a(gk`6oAYjzNOCzB zGKpZephbK^_wYIIbvYMTI+rv{{R$iOqWR876h&J!d2O+ji{8gY?{_(uRXUf$g>mUH zw>A_WmQ3j@^kWn_SK=61F6RSm3>hkm|W@sR(}Vp9Z-zW1ucdf~!(O zw{(|tBUd(Lh4ayh5N4oFe8$Z#=ax$6WAIa#^k$5oZQ~`obE|&FBIh=~;O*W8@8DYR zmX&Pw|S%0r?XsgA!j~k@LYmohX=mUJ{ zgD&UumChHy5Sj~bn9>r0mphNJh&{^L$6U@AE1k!|oQyNe(wsW-G318k!{^iDAhR|M z=FOJjhtLRU4vp`IRy#vhV@`XWPxJ;K=}njOtxD(HUBq4< zzt*cW&H2_!L)M&Oa-ifAJ~gEGwVms&0-vi}M=Af)Zp13gXkK z5!tYNcpCgMH&nK23xw|?Qp2-cimahzh>gsKA-x!i?Eo3(L^@~(T$>jQb8!d!7@3lg z?otvpf#8ae6-p9zDncM5Dg$1P8US}ir6eS|luR!qc!koHL4x54gmku~LE}6LvYat+ zK2wHkQ4W}v-9al$3W3$;l<;7elB*42falHG(0u~hXdW`7!yqZTXK3(prHBoSsd7wo zDtr;$4>m@RhA(=j!PS^Z$f!0!bWBPTE^mZ3j*OwIj3iYs;-EYx6>gfG4%cG3z`B?e zORY$y% zQmYA+Lq!}4fmB}ZU@EA`&t-eBCg9v!Yuk56T8S*RgcWGRcVgKlxeXJ{kI zAv~U7OMEKxK$Ca8zsK|NS&1#+U~QLJNa;cm zv(Y;*M$DsZ-Vhkwg+)vnL?x!d&MqldRByW%hU)ELFm}8tJF&$bunSu*WjB6TDvuk$ zV97*rPo9*ZZz7K-jmD3^BZ;(_78K%p;bLMo#NEUumf=`ZC=5%Ya5{nj7Ebphv2Y@` zHx2$hX$%Xk=0#>`8Yn?YGDX%IZAgl&*O@p|7v?6j$V!7#3nO94=ukM3oRWai)#gQG zbe&`92)LA-0gz&cAt@;D*AF69&au2eoXK-Zy}(2^2I zA@xQIh13P)@n!y;!eSr|PNYV|_|)V`cDgIyXzX%Wn92j_L~1Cv)7jJ%ilWQj5&p*@ zIKtR_Lp>EpMH zZBa%B{8Gl?GMvpcL1tn&yn}F-ro3>LrWic9Ig@88j!uG!nLMU=Kc+7`e3cx8>NPJ? z`ZF4%$Nwp_{}{H|1rjc_qK6IBi4 zHvFk81s+s4g~9o5;V`=!&uIwed4}aQ?C{9j+>J*bLtkhNhQ`q{eAkVl&*7C|xs6Fg zK_F&UeL1u=JUC04f^?k)3iFkzjK+>yrCJ-p z02$qR3~gU3%bwyBmUK_y;aIET%HeT_!+eIg5ME!xmF_$s>%D_@@8KED1ebcqu&4)l zaiIp_Sq(cFjZgRB(YSr7sUIDZN&_|iuL(Thi8?lM**%J3&bS?cpbS?XiZ zVCa{{Los(*M~YjK#lmq=ue}J-t;U|O>|-=qO^Y^~0S0FC2+Un3o3N<6d-}{7wJz9} z-HXNI5jGFT;!*sY$Ko+;xs(_2yHYvMV-Y36-CKr?-t2g_*-j%jz#a%Z&_* zyBVQ(^)2Z28FqCl6Lz`~6tN_U(05kjz)E zaY@WAL7hC=v?7PcX+-~6*q%dYvaB&>cnz7Xl3U$B9y0s0Gr16U_D`|0>P-Xj6p0;r zicylW#bY1^TP`IPzblorPH?#l4;f1|rOoVIo;H!koJ{I>#SW@JDVGIn2^`EFPZeE@ zB?AUq=mf~s;0Bmw;nCP)@r}l_mMqGQY#^2y@oz0NVhb}Peq(0rM41sKy{Gb6>sTuI za#lw|!$1?fZcU|v5Aog*9Elwsf`3?f2! z3bf^=vZJ}m3u8%~fIvS^TBJ+59M&ny)26ukZtm%TDEdxud@OBtMPnI}OD%A$HI)Q3cw$ zobRpgZYxNm`p!G%vH~7U!wX~K!vcCLctB&yU?}9Vbfz#KCKZmSCxR(O5;VHQ@VQuq zPYYAoQ^7{>+^oi9bB92bohQ3B>k?r5cnPN4seIV(g|cMFpuwOblsz3xDe4NB?L3C2 z6eaUxu%U<)Zqein@MIAWq$x$dffP0%i#-?|VIyKCS3ZDOatMN}17vu306$FY2X+~I z8Zq9Kfo*&RlB>MNh&+v9%Ns+lV~1i$c>`PA-EU&crM!jTmCDId?AZ{#={)5; zBl1M24Pk&2gLsTIJmPcerWU7aMWBW&hoQwh(h%o4bvGBM@-omn*sBbL0oU+|w+u`w zp>S*0Fg%T62mLs@gohiVx$icXkne&zd!qlP5*}<#8$w}CsVOT2Iex3MD-t8}I#3go zp`R?+w^;Dc$kAQt8o|9yHA59j9==avq=>jW6_AIW067h4?oAqMQe*xm<>m zgF|8D;9m6HiU1a#4TE`jBHNqN&J5-$&AllM?j0CDW*B1iw#h3{hBG4BW`s5uf3{`g z(mFsTP|fZepf1EAo|84`1=r^_K@U_&hgm6==^O6k52a{UJs7Ckxt$!FO;P| z29;s)jVv+?vZi3yPro8%DkJhbP!l!4l^e6+%du(DwpoU@p(dC#)RXo!NIBd$lm{b& zIEG2^;ZQC0c?U5L^9*8wv%_TT0ydH7Jtl@xT2Wc33G!@+9SqM^!+3bM zZ8qg-b)h&&p3=w%VGqIDAO^UKv+o|Bo3%m|uqi9K0ILc$0U7Qc&f?Vy?+xd1>VAy9 zNR(m12%a!CAg>OwgK@fQ1W%aNTX}VO9-*u{^g>y}WYA!U9?27C_12EKoHdfirAFg% z**=mdPi+i*Y#hbovc4h=u8opl#VCr)BWzrZ%VVQ>T-H~ZA*4ck{@6kB2u&4q#GcW{ z;!nP$4}0?6Uy%y&H{w$QqVE|U0fy13Ry<>=UmL0K!F?eJ&&yEbrUaCZmp=0)f1~q_v3=5YY5HZ%2{S{L6 zp^{&qe9egb&B`~LR5_ePscGYS!0NG5){ml8p>m0&*kjIRO~?SHW7$+W*55<{yYfFS z!0J9tfG14spdU|<&A|EkW?dCUi9M7MP4xe{=);oBD)IdBzVh027w1WrU`CSsER)$iwmpYz;FVhQW$&$f?VeEJw zd>oz(Pmbq-$NRDU_){`s^)@faSNby|e@xcKHNccg+$guE0ZNBX?RZLgxss0aA~q$) zWqT!$OGL>G#pZy)H}Qly0m#fZEtRF|q;Z)y$7K@fvW-le`8jxtS+JFPhbW z3>^XWH)lZnEq3@jLdvlI^Fmk>WDpsCLkN4bH|5DND87Zi*&BTeuk|#z0q(kmt*w9d zok-CeqtL3%U_|ZJo+eQaf1W6RUNveA8Gc8` z?y5YvscJmh5sTT0ScE;zU)H=|#S@?g<_Qox7>A~b8L)riK&YOWTEG&Y7s3)BgU|uD z6I(oW?7|kl5yWqNBZ%Ks1EA*i49Ix43%y_de&Tq8-AiF_mlsdsuLBXwPwCy0_;xxw z(d6@lY`Tq?k_cn5S9)72)j$Y>tGCJU&~0Pz9bf*~)5zC*)oxQ>V?$;Kk} z-@q1E|4nRR8HwLmMh3H|hUUF7^my{XWL`xg%X?0r!w!b#(a9O``Mh*E`CJ%m*~_cQ z$lH0I^uk%5WN;bE(H>voy_S1#JmtxaxAQ!yK@ITI?JRK8!967!x=!&t4r}ak7&#>y zUx+fcR?%&rqLql=5#B&ZH@tUTJ2e(Er>0u*7;M*=wDjMx!_aGSQ?nO}>yKto z87?A}-Q$_7!x752slA45kf^~8P*al)F?V!>mugaMV}cA=B6>&4#H`0g%4F|iyJ3s# ze*?C#M8t0_5%B_~Dhc*Y$GW6*H00fp+NC1VrDP$KIjt%g?wSxF(KZz)}p|314KC75au8fzxq~QwYpyMkrcrY`oFitWbSL#^K-4@2j)2=2hf1|Xn+d0;cZ19% z2`{C?Wn~r-XX6Zg;8xjeQsxk$ju|qG?yVnN5Gu_j;ylP)l-)ebEGzZYcfzFw*iVoY z>duF>StBGDweNx&24!U>KK;**gO5CS&oNLu+iI?Hx#n8&1Xk`L_`T3|VRmy}sH`le z{t{@MT?B?XvGC^X0vL>c_g$3MTo)EmS5rTGfpQ;#?x#8G!eGgqC}|mymP5vT94A&* zR#1N>TtYSK(xeBdyNdi_Q68k;LuBU0x=`t1>aRv4WV376=nk*1A#6-nb8R-eDTR~B-asT3ol_`nr1qoem@;V-wKp?|G*6N7ZZcfj zLZruN^lS5COv+XwY~${@HZN7$PJ|uI73><+G_|gFYE6CpoJFvsJ`S!d;1{9aMn%G9 z>}Sv7{SWICr6)-8NoF3C08Riy-4P&l)mT1or?P?7A zYA+G?p|7&pjj740v=C)Koa7YP))FQiAlgAv_-YIW&+|lhf%LelBuj^=--@m)mJU<< z2y>hFwsZS@lXR5m$IyR$q!+1u9DQQJopgeFC&^;w_)r+3*K6ZUMZb7)*MznL>f&5N09JVenN$(QlA8408P7I+Ac>ud2H3weTf6NGlNBQoJHd~OhsvF)9V|iDMOoSfpKMY>h$=~Nd4Y$@ zuKha8icB06ytHJ56iV$dSO;uyb~hMpQo@NCA;CL$dqiX4u?7=dY_LK}W=JURB{O00 zR<$`B@6JL^N(>QVC1{$Rg%rwkP6gFR<>Mtd#i@9K7Yyb_COLr^UAWTxt}aAMBw`Y2 z#i)~S8=EIhsx;$$dF+F{H#!$l1W{>8Lf6@vD-A< zW!;E$16t7%>vLXpCvp!DGF0kGgkFdnBxg}OTY~X((GES)?7fN72XgN#k@`}*9~f4l z1$sr8lpG@Tm%uy^wd_UPOwJ{uMS_>6ccnwC`hGmLtVGSjS^7X!Z@j6y9|yM)F&}-3 z1Ky7oUl1x65Ur4#DR&tjZ+4>Kt7qod++_*U03r-TcV=Th9`=qEFw3ntH}0Rg@J%w36T@kSyHg?-D5z_J1@y}SSv{2Pxj;FnjJDn=2g zf=@v=TKKAazDXWU^fBC5xhrE$%2=X|Lya*$b62Lx4kF;IaOSOirINZgaUbkk87kdO zgj*=$@tP{8<21XKI1@ZjcAQlaX`%*}VeSKA(j=nY22Bs-%af^ny99>XEQshIZuXXR zQ;0Ja9atz;Q@aLASCz?kQ2TG(W$u0^xcpFX#EhwUax2q_HJv7Fu}HB*#apy`=WOLp z;>{rUArC_8&B$ttcr(oKH@}`yjR@T|~NzNF_Ghk_(7%H{X(YK@{p#77}R@-;eChC~{Gl zd=C-tORskYA=Vi zYuNF&f_f{-WMlI~p>d4~US1Omo7eP`9wgpF8c&8l)`UU%S~}k13Ve>YHH2Hsedt^p zFS&`Zj`rS}Lfm`JL;$|`&J^;!x1NZPkRzMdMdD=*#_a|osx&9R!3mc(675lPo#$4E zRgIg7z8M^VCsq-ja9fD*7-Xy~m9|oQ8|}L_>oTP6)Zc+NW8f~bvXcnA;Oe^3(r#)$ zj=r%;Pf+_wZ1;gBgD_P;MT9-vPnVm`^3z0khI=U&lB8#e@Ej^KK-x>~eK^k$xrN&M z`DUXVlJzF(01*#D#`CmT?zppC=t)d&C*tqidDWvLwe@E@$-!LaolZIoHQg1Ghry+W7B$*ztvL zyCGv&nADSqy}+`oP|BiqHq6^qCiSLvA1L3Gjqi~oq`uVehXeP4%qN1uG(HT=6e9MQ z>G^o-?o>IK2o{;@NB;8K1ct}?LB;#*U3h13AdNkbbB^Qu2`egNk_Hi=7;>Mm$tBb-#Yadi6Y(`Ro>XN-!dpXDIkGp~?4%k@#35YO zDNkbD+(?w6eDWzzCQHMJFdQ16ES5)5dn9+1=M6PGu0|1~0?beKl}A&13=5WQsCz0@ z9!rFA%;AjSTYor+h5urSW93k*Qc0wnJZKsE?+KG`CgLsVhywXmYER&iU0M~9;w?a` zh%*ttJ+GtKziBS53QJ(mbd!jCn-3TNIXSUm(&7lX@99`5-bL?}BYAN$g#@N@P2QP^ z|5*QNbaypTYq$;g3+Q<14kG>y_P6wgcP644rx9VgjJ0bvQsU#4JBcy_zdg!T@t4T& zJd-BX5_2ZJ|6E^u&73IBqW)|sd3KCEhuU>AR?TcEnQh3qM4QKHC!fQlT2B-w_hkKK zR6!xieD2Bm$iU5bJ5ZQ>fenM-zYDn_9DJD+a$))JEe$5A24e_MZg%*6viLG0sBjtvXy+40|s`sWTC+${*sz0&RAsRK_-$dr7TmSKZC z1syzo^f@UmJSfDMfnvrWXt;QDd_+*1j)9ycHeeEzz4|k9i*;Nav50Z|!-ZEuq3Lv` zAaC`^H@_w_^E^yw3c$5_xChS6GT^|$x~cPr4K@TFJo)~wv%o6iLKkQIn+=zbW;^k{U^YvK*-m}_bCjRip!{NK^B00SEM8k3Z2Cf6 zu?4M#D=z(NwqOhknhIIhzD_ad|C?EViv{iE^WV(C{_P*QKAjgWXfIzc{zp()eW~71 z3a7pkRYt4v@Ce{u3))SK9ds|+F5TCGxY%SCI?1g~VTM_P(XD7f=HDO0)mqW)$ZEVK zy55STG8_FOy536WgOmRiTyN#pYy3@gy_KR&EtL$S6h)bOO)~uA?_c$$HbbdeW-S@l5<&OZ9SP~5PhCU=2Wls!8-uIjEBlHYYsR3Dthm-p|tf>-vB{rqsUVmi^XuT zQRJzmC4s`>Hw=fSQqg2KugMOT1z5I?qObL2IYxuO&utWbt&Kwh_}qr(^7Z-Q8%3Yn z&}(Bofg33 zhE!wehC4-r;Y{3MhnxY{R>1dh3N3qv>ZDxR2C_sDp#)5se7z+hx zH(yWB5rd%s&F2F-RVM~>fp-%J@Oc5+@%lb5=tzrY^8y)$r^S;BI^BOG-IgZxg*{=TCdleteW35RVV!;PKJ$T7cujFw0na=wJZti})5k z`FsFp7V(&CJQQGMMLgw|9u2UvA|7-nUldnX#EsQ>+|Wn%KRt@bSa9;?@BAJSP>oJO zt$)30C}wd};1f5jAhS~3sK?J3{u`x_Dl05KTDsv4A>*-vjAuLy1{~411Y{K%=dtqN zpNR|O%+{ti(U{2gePUsS8E*TMw+(arGlCV@hODb?hQ3CBDM}VV?Y^gFA*aF1UNHTnk&tG$qOBFy-!pT$_^ z)nTy>5k~*8W@Q&QN^4VuQ3z{RbD=sXqKtD4e!+&?V6YWO8OPw9zK{I*Xf6KC-`W&y zO!f}|R&D75Ajaq)0Qu@muNq1>#2J6~iwtCZ!JcH|g>|q}>w1o7rB;XIQCr_l&x$S0 z0H+c}1F(7v4bYHi6!vw#wnbBo!oJSuTeKw8=)XmAHq2+o(~bT7_Y5nxSh%2!|Bh$1 z7W1z~@O%NqnA%`9{_1C{0&c2Pd4j1}b)~Ji%qFgwRaQ)od~~_rK429Ui-tm@u&r4= zJ4C`2Kq zUY#Jgw@@?HWWiJ@!cEm$HqrRB|B1s&Ddxh-qIX#(#Wb7}fW=c;Nww&BRzUeW{_YyV z@vMsSar~(}1jpM|LlkEGX+o^BB8odn%;&6#x~_Lw8Krr*X}aiLR!QMnOJ<>W{qh=e zaEBJpFh1vRyCQ9e)`@;+r4+ibp~oa;6IU_Id87c zr)8^*-Maa{2#&!2JEZ@^erifxy~@}lAxy`x91Nm!SUtDOc*CuFp7|K((N9$Mw>T;c z_}_Mef@FLP&Hl7C9^RRSf3r_GBY6L`y{r>dw?Al18?V>N{)!Hgo9jP%&5dPOAS zG>6_>_HT6qS^XDr%M11Y3~c7M-!&cQx>3hBF%bN8bX?I7ZBehejj4@#9ow$zL4pF` z6{V$>;i}k2FKRXzd;%`y7H^VO(k zL3y_ysK8%?ce@PfpM{FEk5>b12L)}6)sJ83FWS3i)-8hJ`55*%E~~C-M&)zB818S> zc>Ul-{(~3%u7;GM7VM*s2S*p3s9@6-I~N_hj8g2=v7+JF3u^Nd6%z1rz2L8^6usW~%(+v@Z|!~{8|@musTQVcQ$4pH zrwV;U(52l2_=y3=B4Cn!HSTo1E`o0#f9Pm@kP(-SAHAr#n_kqenl^|TxVxs|@kfkF z-EKXa{)@+|dSAAAaE1+-iErN^80~`|y4h^G>N=dwFB7{9^dKP6dV6fTL9jc(AV}M>mArptS-{Ta6mc?D$Kuv*9 zm;yx`1r39WH9+P@1S~omz1vhbJJ-a8$*G%*j+=D&jT zt5ur>3%ZACfYMC}*n$8*AGeLrPoA%~Z4xwe;P29Tf84acC0|Ij`;ulIDMJ&pkKW8} z6Ey|GHlqTyn*|LWlQh6M2DpR(zkFJCn|^X3pN{zVu+yGot>1&~dQm~`|GwQY%_G^N zHsw?I?YdbDgnZiauKh7l7S|tSo~o^A$`)MFu`Pnl-PQWBt?H#M!tr*hMh_AWgMlAi z*b{i?e{hFhRH$`d-&6dPuq(I$5~pd(&wUJ)4^aFZ)&N>a7Zgw$_e;X|>3R)>Bz)x1 zx$6{=hy6tD>Kj|dB#b{gNc*l~t1+qfx_;qSk1O;6&^Ak_N1+(o3Jt5jT+z_1LSNMHio%LKfB=O zX_u#G3vOn52PH1miwa45&h(r9v?GzMj!XJmE=()TD6&!;AzofWv zm2MWhB~U$Me^3vy2C8QcYasQ-oj6s}E|D81aOlq1yxK{{#YhKPUj&4rwFIf5O-!P5;)W zt<_gCt7>>sFmC1%9YjcAoA%Xh_T&M5eylpGO+4sHPl2Ncw0i)J2w?AC()Ij>$8=NY zsS!`%ldvAsuE*F^2Vxra8254AI6^)OA6@XTAS0R{KYB^?OFBs*^uwM{&161fA0J6( zWXCJo>d!uftIytpe(^ydpv?n_-h+U>BEa#UHp0z&jNSbmr}6kfnWtXdBRcM%zPjfL zj{A1-C!bu$asSdLPJP;#*zI2j{_&}2_vly9?!iq(oO<+W!GUi4QP+;$o0q5l@U-AS z^G7<2P0f8qaA4WTI*j1J8%xKPds4by#VY*0SRV}p2flyosx^QSm;*O{%0<<^&)~*7 z(P9kqbD;e*4e--5JbeB5hdq4Tvx5KZ=Y93EsV_Y%_^(Y55&ZYb_8&TX3U~dY9Tzpl z{(Khw_j1{f$$_K0?F$W8`W*Ui({qCV%wOxY%v0ZfPVk>wkFlxdy@LPD_?x*MUa$%N zdp&%^qQL%h>m>#Mxt{L4{d)ex-{xhLs%!S5|NgUA^xq#E;M`tZ{@i`y4!xoQ;`Rv+ zbo}MJq_Z2ZtpCm&7?+5Cr6p#x;0o@+mx+Ga zu?pY9cNjjOzdG^5wCgea&52(IU1-#eBRFtoo6CKj9HEyK9JnfPk2ydTvkmC(tgUYD zeq7z_cs}@TFte-y_A|iP0|LMissVo8FF3GGL1@Up#;g95{U7TY(*DPSbGL5T{N$C^#@NU8h2xy85870}tylHnj~gehzeG z>c$Zq*l=mwHqUz${ZiVx>LdjR?rGjWFu>xd&bUEa-Pq@Gb$gx{{FB&S1FU0!oEHQ@ zRSymD-t&S3+ch3{U?Ek}exdE|>8pETE-#|>>uc9A>XsMKRlDHxD}e)fVJ}UU^N3Ri z91>h*&elzqr`8=3T;=Yo!`Rdo#Q3?Y4S!&}W6u@}uG%5Om^Sg^!|1@6-KS?7({ delta 31079 zcmai7349b)($9<}-7`tZSwECesWMy1=8J1n{>+bjUDEIXQ z&*Vgnl10{FO1>c_9cq2fvJOm10axPk`wGi_MFu5Pt_!B*z3ci!TFS&)-;gp9mB4i! zUV&#Pc7nclB)fgLR$iBtRoS63e3CUREX zzNzKD+o7dHFBlT;wxq+k@WkP>3o2$UxMODS9dl+Z^4+CXS>T(ED#QK0N~Y3M7GCR{ zQx?wNW-eDb&+nUG?ppxP)O5HUZimKjA!DIdhs#%0DU1=VeqS|HG6ecWB*VefG{nb% z6p`3X^!sYKdbD^MxA<|-OLJ;`tF(O-`BtNyUHrZPQxwkZ zT*K|Whbvm^_uX6WTL(>KrwAt^g!Edij$+?>{;nIy7?>&~!+Q(UAV-LY#|1b1IF6O= zut$i5&SGM>EWfXTAw|~PgS@w$ytiF`-|lkX9yl0ChpTf1xHu;gVnh)(VzYF=uaO%W zQR~}R7QqJdA>Qm^zweQ9-=nbYDb}nLII<$(idYKoSrVXWZZcdD9qIcuMi)4J2l$IO zkxtlVONP(q<|5SAfjJaDvZW)H(Zc7sg)i`~UiAB3D)+q%mT9aj5n63_=x;BChq|P( z=l=5-a6Y`Q?XJ-G1|Oq0sk?Q)x5_N*e1GRUkNbUZm;2s<%%{^~v%NDoB59CX+Xm?* zH{m@Ff8X!>pxpN%_~ID6BkYN@!Rp9P-75XQQ@lxZt?#3zVFI?-xC8bEEJ(E%H}?=)lurZpb-2bQ-YCN>-&K#`H?I6uiy7mx$kFUA1aA( z(%?dRbSFmsOKn~M%3J-$TmA0${Za0_taNS9QmB7L)4R**yGjg@!YpB7Oj9^K5}ghA zCZxfIg?3JYKcZ8k!u^t{AS}?MV>V3al?Dr|qM)YZMEHQpcIe`bgov1wPFa2_NdYr@ zOI`2*6_>d1&Il~Ug<`Av>Zh39Tfhb!29 zMywqkiX9D$W3Pv9aVfd@dLva`>!eYvjb)v56aGnqDV1T#FO9~}a%qeO&i2cMt#KXU zK!pt+i=!9zVjUYldg1Z-fG>P1E(88GB^~PGIzwK3N~cP{G^q_1T`NsyuxOYbp8-8m z(okb0c>OjwkR$Y-hE*C*6}qGfCS&ZBrmI@ubbK~EH@v-VQ+#;uy9{j$rP;i#=$9&$ zwicL@kPYDpxzLo5+>vEisamNT0{_ISdlULXZemI&Y*VAQL2^pV@S#kEio^_T-BVPP zDm;;x%6t+~u;FkrktyjJQNa0{-HB+lLDdE0J8>EfUSfxbJ544U4PFursa_FWol^>s zMhyy>M)g6~F`GW5omk?I*@Y#)v>QLmr9Bo1mH?i$` zrk_lF0ZZJ(7qR4*Uc%3E>17Ly@5)RRVRRQeoan+Rn8uz3C1=9ec#O*i(2eHHOwSp?3tNizF)TEU+q_ z)AaZXJB&@Y!AI#{qU!M#oT}3^c*av2@QjBKIO5wfc*bj88Bf{n>kOXl6lgeQI|hw_ zluVxOyekEmmT9MqSCh%<$z{pCa}lbwxKzqxGJ4Z|RV(Cl&4!<|@J&fqJW+IQ*IvDN zLS+3BMf+ii^Qk|U{L%pYESCm?x2p^O=sFq-Cz1-WLZ#;NZmEeZQzA~G<|r&8HKD0n zDlG1k1}6d?p<{m=Sh`dAdKDy3iVRB~mt%={d^48((ggf0mu|7b*(sT*B9@rBu{%$g zSjUOX5+*(%GhgV=6XuCPBrNOhfS4Y{(CI1^ihFQ|GC-Fx9brok&dy2&VS&FOq5xNs zD|*C&y{DT0lv=o$#RjFF1)joo6ME%CL~lmV z0@%`q~$BUIW6-4O?VNzv2T@>_69Y zie+rH!~W|9aORRD!md+O-lSZfToJ&Ra!t8>DW@PTa0(G@Dfg$`cyQ(M#Fn0y2ruOE zq^1CQQo{#i&$oGuJtZ(apXS$0wKqy;3=#p0^EqQYPQG+|FMm8P-3CF0la>hiY7nJTT6+7di;cHhat>dt+JuTq_ zGAX5yGwEGd3RD-SvPpBg3WJ71&ZKu;oJki7Q)t4hR1g+OD&kD~tS}Z96e$yBwNgC< zwqrG0mQO2+f$xfFimXxFQrf`>Ab(Iy=lPxsB$$syaUB{aljDR{0GIm?sxzKsfYv^W@U9O@QUw_$q(3JFS2% z#nk?3DaPPz|0+F%B|h#CW63W)f}iEmqt`*8b0(av=M;)5WfV%oXIiSk?oy&s6Jq)F z{%R>-lQ95Wlcn}kDD^x-Xx@P_{Wzsswp~l9?fp2VPzv|=6P>RkzAbpM9Of*@m)>AP zS{S^EC1v(|3yX|qqx&e=_>_bfm;T0E(Nx?FTS1_n5rzk z6xa}W4`GZ`2vZhcS%cDO@ug~_Y=aNTtSN&yvjW>i$Q)$Lx`JZeIF(3OnUF8MI;fs` zcaYkP2zv(!SrJ&HYv`5)CgdwGrJ)7J4dx8|eTN;=2ixGu!QSbwJ2(pqJQU7J zN&}vx@PV%~lCi{Tm4YR|l!~9_Qd$SLNvLn&3$m9yX|(;sD)K4iFnquqcEN)@_Vk0% zLpbpoJ$$zNXb8_?Y7?*w<-}tEJUAipO&+S|Gp39!df>jHoP3QQ5ylO*!}~+|oQKb) z$^uIofoSc>xQeA4nUF?iq}nHDIHnh=w{aNHa94Noh1bW!h=N`P$uk_ohC|fw40v={ z5!^YPXRtO1TUaqjG;Bu*n+vb*jD*v}Ir+XB&NG<`hqxOU`HErA4N3-6YAjHXHH*il zLDX(r+f=v^YqhUb=#;9Mkk5tHO6wsIGlJ6uTf1o|AT2SrEs&N{TUvD~O=(J8h7U;6 zh7lQx2_oD+!j`qlsKO7E0@RTb3!6wieioZ3}PA zhc9nr{49W~kvuyp0G=K30iCdKBxT3?G6B4Q;w!ITM-m4c6*OOYF?0m<9YynKJ-%4= zC>!kwHjW}Xa#5G8{Rq?zZ>e;E3Hg>)RSJV{B9`)YaYWJ`H}zI_t*ke`YsC_enHRBy zyH@dxlFl$8ou{0|5?KJ_#_(yHH?*-uU)Fz&jfH2`jn;r1(%|zLeh*T&$3X6=gH{qXLSs9k_v=2Y=v@e_}yX7$l zIQk1PVFJ&uSl33`@%VsaR6Bw1ImSH3=lD+%inN{724#7cL2+zuL3U4YXHCOL?Rb

%e#z?3#8e+G19V9X>5V$TeIOWx8!9nG&XQaRn&kV z|0ykoz$R=_II&kLKWUbh7#kK!OQ|8{VWl0QMh^tbCwiK>zR<|ybVeN3q(%J#>*=w%$_tJ3MaYX;-pk(7Op`R0MD}c zz^SnlOWd2gu!Kj1_=!h^ux~i&ex0T9)DivU&gH47}~@4^g*` zxAcS9r+UDrPsc%ZrU*N4;|ou<5&ZSGY>T@M;a10OvGiZK;4P zaN9ISr*z1i77I<&)B`=GdI+4t>cWYvTG?|tE7F|E$4jaLT2kTz-Bw7yVTt?ccP#m( zKk&0$x(xd#c7u$HcsNoq9@bVA!Qkoi;$p-#foH=93_Ra-dU0s1;BWpGLNV|hZBTsk zSO$%T-w|3jy%fgJpf^uaA+T&lHr-8RKAXT^yC5BFw5%*Gk}{Z(4tz3|hD74KV#V}X_4fi+97au1=V)e)aK?(L#ck4Nxf4VPW3mK)E7u2sQ$*_^%Z#3NA;s% z%U#{Xo2ZPjQFWjz-XPl`|E>h++r(~@<3(vK)sBP1cl8mxR32Yn*$r&7V}){hycxRA z&K4$6`4$*IJ4d*c$`k7=oi?X5i5@1yjOE#M6?}2FO}LF@Q|c>S!c;2Wu2hRKuQF1Y zMzRX%x*{8|RwG4eItgarSTaEroUgROvz75+pOYiZq`F!7rmW5mRdb@nyGS-W0=hP3 zDL1X*qEt!JIn))uevK05l58GqtVsX2ej11p?;+XR2>2n8MK_s& z1vcSc60U=~1tWx7Dz8VUWTPiG(9_2H%51y>6j5*=2{zGyuuIIl7e>O+h4Bzo$u0?P zqI5q2wkm*lp^l!nkz*ECMS`U&0X|*K?h6xz2dHj4TtW2*+ z(QvFP9?8XS>mp!qb%d~!FuOPb=mHTPzMF)5&@DFzf1>hUbOjP9#vw|LB-n?pV1iii zJ{>MRM6!p;B~0d+*HL(cWRK#I4HN!M<;Tz+tW7cul5E1`BzpqgQV2In1#m$UVU3gk zQ*rq5642(8_7ma&)S+|OwcvpLkwO#64w5-r5I^{3)CyUIN=C^o@d^%6u=CB0uB5XKMwp$RP{0q$`yaS@CrR2Wv;v87ZL4PoB#`! zr?IEFg4`K*`sXZ};h!b_m2iI}2jT6Ttw4H}B(GuU3=&?a@*Cu$!Am*{Z_@KyW=b^)I+%RyiZ6Rn6|W*tJ{~44!KfS_yZO2B&UmtneN^ zzfZ1e?3@UVOB3L?r2~Y2ko=#-=BAo1qI8M`ACZ5WYPt&_)AJ`3GI)t*pT9slO`=ap zG-z3*@EJXy!N@5U&QkdtI%=r!IhFrKHsDQG1fJbi&szc~mnFc;v zB*A~N(g`suB88vm`Db)Zk?=n%|3dkU-9FV}{Uwt9O3pdHqNDH|J^zl*@rZv=`7(Ek zcV)D2g#=eoc@At{87qbfY_hQkP*x zkp%=+|!}uN9@rAk7Qq|mw*(|ALkvtnLb=eTNHW5C*Cm#B(%@O;O98c@d zt;m9jYk4^3k}yx8nTTJ@U^wNI)JffOVU$=vG8cEDdV#~HqC%=D!g0xgz&f0Yib+rc zE%y!<-Bd2+p37U8A@n0be;PD)kHhAo0VEp;d)JK?22pu1dMcaU)i6(ev(5#3a?n#l zscsne)RgrdEk2lD%cmoJf-|*|u&*{AF4yJ?BMCZ+x|y;*S-6Rw%g{I3c;SM}hS4M# z11;-E2xF-{4*imCFOt0UG+qEpWfr?#LHCrC=w=!Zl)#f4B0<`~C!&mv7?!sZV4?zu z7bem3Wab;Z;fMt1#srA2!`qIHiQ-hMyIp{*s6Ra~+9pjSNd@;zdR|A!yUzwcY-IBg zrmwsxQM!W=car|byf{&sNrG8$^u8N}yQn;yi0N@UY*Hl&=I}UTS0OBh=8|w8_XHEN z_?b__1vpH@#D!F@5}U|4yq=hkKBP*A1u#FBw9qGcU^WkxS3B$cu^6Jsmv!W zA;3}|J?#3Tbbm*ohJ?$=NCi1>iF7wfmt+4q(CV_m$AK7O1<6*DtKN0R*Wb@)C3b%? z?SM@PQ1u$-EK4zL!&<7SdigS8>Cg217zChHc$~^lKnn~No}}_q z%o4mYu;R32c_0D4-;~Dg4{%8*>?gnh<{W+rfMMB0vV)ZC)Y}8}6r6k@9=5DWvzJKE z5#lf05rOS?;V=oBaZm<`N2vU~TFK_77fAFXRy$$vj!5ApdVU$*QzX1X<)f6|*!g_R zj!3bEWXHIJe&5ki_$vwi1`gRHzDniSxL*QtwD39!-azd+a8!;J-=ycaD6`c!qzZqh z=i}s#Mi0i*+a!30`=im5Xp`P0$q61zjh;?$eGM)yNO+Pw!_MR}wZV^g-p?i@^^BfP zMjsIFLpb>$PL2&U8QuS&3tD#KSe~M~k9dm1vvr(|J|@v8*gH-}?6{q+@=lZRQ?l{u zPMnNBBf%N!{OZnR;VeC$gUnq6#m}kyFLF@Z5j%^ee^bR5aBx?S@Edg@A<hU(qgjSMY2eE&~F4(DetSEIL8Q!7o9W6vqS%7^9 zu-*>IjV^GEL5JB%9x2l50ne9lPO59rk+w+e?DyGE_vom`Q@JBCNvs2ajUe+Wz4Ln6rWoX5+6}t@j{gPmyM- z7%$FLy-3nqq{J8F?Fbc**r4kpe5M+BvjBS^i4gk`D2MmRgc*jEuZIhd+#uvq zIS;z-%jUX(%ofv78CXnD3K919$>?GVu zf{C#5Nw+YG%9G)xCx?i)Q5jGF*g^!4COh&8YbuFv=biBrOVm1zq!p+Y7qnc;w^Y-q zVg`)e-xu9L@yn;{7DwN^b?3&n60WB-n5WzIo`)=j&)y) zg}^&=&7q~&LZ5wiuA{Sw?_kZ@!S#PDzus)vm~S0nw`qT1^=I=;`*xz++FcBRGv*$4 zUgNs6=B{&+>kj3g3o+TXRn5!d4a~0vd3aE=Xrs~lg%xNb(j*f zKogg;CoYs?O@>IO?`*&+U`3!&^ZLZ3N^ly z13f}vHNI>1^a@4tLQPQ0EybY}Yl2etl!PF8VY?$&yF)m#kQ~{3(H+VKg_=lg4)hB} z-9k+$9v%>ix`pUlM%{rSs9T6GRjE73@?Z33QP3=lFXE)thkT`nGQ4R>sNpTrWYA{Muuu#y(xuQ5$Z*rxDbhG@$EXnGD$=;E zIe=L6)Lf*A$sMm8Zxf{Kz;I3{AF-UI8 zuOUh9MnC>PBzL3xFiTf&3N_SjO)fpWB@~0*x@1}xiotG;59OBn5JYopT-fXZWU%S_ z2C?nyo94iJ^mk9QeN!nqQB)l+AAQ(jzF{oYh3BK8I9I9(&%=*}#A?Gc@OY>pF4ctR z;ip55aH%dl_lFqaQjMFNTlQP}h-ODF+lN>OA3kE~$F>L82MZ&%k;`U_l5E-Yi220F zHgn3vx9WMzPiFHY+rTNXS6;A8Kv>ZICfmEkhC$wIcbXz_>=1Dv1Oi7aHmo`CCChTt zy${=;UBj=Im(g8jQ-pjh-hy+)G4v7|74%e!ZOhu#$$`IF=9)q^%s%`o27#H23rALediTQWAp{|qt2Y&+K8Yjevf z3?H*WVOz3mu%3@G+z}fzy|cZUElgheiRE{M1YJ+Fy_xzJr_DLZc4rFg%BLY-7CnZH zLr$6P(Af6tnn4Sk3&r6f zZp+}qUxc7A+jO;SlRaOW^NH=d+M&%Y2sKUQZ2Q%Yl|7e2s70?`18uovPE)q`Qa$&p zx#!sS3u#*Q8`9J~$jYpP<-ldjP3XI{V2M!P=dolQ{v!l^*?z0t^21j`u$S$-ki8{W z&DqQLUHs*?udnJ`)XQvIZ&K(ibUg`i2@u?Z`7H65B0Z({TWy zrumd@n%c2DJR!7d78-KEWBt+0WvG+8ta;#2>ocZr#5POqeakIFtftc$+bac`&W2mf zr!%%~YS-Qz7!iWlZ0ppHJv=f5vDp@i5qp$5vDp?%CH76`8JBIM6xb_e=0VK%Poe3` zHcg@F%C<@@2=7z`*rJB$? zJU;|KOEt0A99R&XpiA&?)y%0{NCf?*%9@kYcAbt_>Yea!pJ{(I-v=oVueSEgNYEf^ zu$dOLy4u=(h8E;~K?kC5mdtZ{%Pjas9${f3Uh}h`hxYz1x8R@bHxYl^f4qOyX^?-Y zwkFA`d#q6zTC3~+rbp)ZtX*WA&l+{77J2Gb9a6C>-(c5_;S08Ap)u@Nt0U=koep^) z>uojmrnBpa0ePAsJ@wY8`C6UsH*`8xy9)K=HTb{3t@hwIJ+WV{jw5C|isVemnyRzS z`j#^|*&C(ii#0;D-uR*kafIio$md`UI-d<#lTCLUF@Gl)wHKIU2 zKx;mWUJ%@Q?th!;aOrgn^dHJHQPTQ_PKmtNk0!75TcdPCQTL?|shV6QKkm1t>Y^s> zYZG*_-niF4It_24_t~#@kdnU9>5#pP(YD;hhNzkJtsc3FAz6Is?7E~!s&7%Mx2yYe zzWr)6YZ+AYoxV5A63j9zMsvW)j<9a6Q+kY#3``ykolK)di)XV?EsbQHMVZOfM$Y^xJ=NYyq&mU*x9usqe5TH8*ECOX{mIiqbgHXT|XTZ6Wp zK}PDcOr#E}*yfT`YYbVYHN^y+FK`aiG?O=j__%QBbrNY!NIO1&Y=)O6G7Rndkl6WM?I zf!*hauT~arJ|&*Mm=k zIFV}&k&iNy_pCA`8^<^uQZ>0qesh%}+0^0J6$gDmL$Vog?#b}rQ=zw9r$eq@jkb+l zZHOAjoqFU`47m}Jx@6O;N2+gONH*WRI-`%8ptZ>+X_mgXt^rIoJ?`2*Q*EI2_7Ova zDFHOEG+;2#jo-5zY}kqo$)@Q&_nzRqXvObk*65H=vfd70Z#ur#&Cw&H)*$j5MCy`F z*gPFl8Mb0Wvg!Zv*(ZVrbkcki9YwNtjUkpMEzqImmzizBWCy7Sd(eujNGx4islzY2){MHMw?rO*YIR0{jL z{-(YfohjFtpv8|bjZ$im{u--`(5AQ0p$RFgL&)Fyiw~6=1qD)hg8gS8`8>= zm0z6>p2Qlq>vYKHSZ_g&W!yXT$kbXy204~FqDQKB8PZC_gFoZ{kcUe@!?6s;ZBWj; zi{$CGhJbFxZ-3WBi2M$-?WtNr%B#aKfY%^Z+YD=%@BW#2uSpo}HPPXgW7ivP)1&28 z>(RCt^rt@AGzHgNDvw6Bx?~gftO+<@K4vtcPLGyzH=q$w8w|d%S z{&2jkL(8{pG#u8}y`~4%ZM1e?z2u8OL^W!)VV?Pht^v29AYXr`vh|;P+GbL%ib-z< zSJ58L^2Zyk@fq4y4R7gM$pi1hlNc7jI(oVBUj+yQu-To6xd>n~=&{q~m=ZQt^pf?~|P` zJ%2%6tZG%Xny4^L0O{||y6BKEZn6rxWBx_`2H(NU|{Gi*cuhi_K?K1NVT?1~l z%fO@?Hk^GXI2l{@D&)Xs9HMEP4bDsYGI(bN@`lX@=N&Nt8=^Y($)pKtTGzgu_iJ5$ za_$ziENY82YOWUP_*REhhM`FByb~Ecx&$AFz<0q_6wB3HFfC~NYxrIVmQQcNbo1gC zYosoYPMy~w70cXw^Qcu@b3q4|z4u#t^wi$^T+)NMMmewEnx^Z#?xIc~hZf>7+5Luh za{r(M%a`xR(f;IqYm`1Nf7BrrwT7hsSI_xmleqk!PKE5-ik*$yY9O5h|Ba3yiKwtp zUbfXhx>gggfd}tjnp+t>qU!8cZGUn`9a?r7Ez^vGo!(dn(_@C_1s4l3q4sVDG zuEH(Ps53e!S_hWjsk8RzrDaEJM;(Y8RV42+Dy)vtVP(%YeAV=AMkmDTkcw8PJaU`C z34yraU;{k{J@Bm~crZ4^>r}|+nPqPqEjyKXC13xAxGjF_Q3K<C7#kk`3=rH&F+chc_75aB7ksw5kDDoS)CwcEp4YQ*;fu32wuR^X%GR{|Fvx zHB)sevg}ZP7M(;AX#!;<9 zDux->A8QA8UKBhcTTN8Bm7>0INzM~74_Gmxfq+VRUKq*<#| zA=k0a7Vb90j{9C6Qt7NjZrE*f(7NDY18F{d#xW=OoK}bbdYC5j%Y*h{y3tO2VH@?J tnmyM1ehUh5io~CCtY$ynvIzWEuduMAFY{+S@k3d71e3ji|1mtr{vSkbI9LDx diff --git a/source/tests/pt/models/dpa2_tebd.pt b/source/tests/pt/models/dpa2_tebd.pt new file mode 100644 index 0000000000000000000000000000000000000000..fa84a9b5fa508ae96276361a02fb83dc5175a3e8 GIT binary patch literal 1009 zcmWIWW@cev;NW1u0747`3@HVPM)4)7Nh$g%i6x181=%@nPCN`zNT7i!IXOQkCpEbw zGe57G%fBckwJ0^kB{R8%tB^r6sE{#&4QO(1Vo6bEMIlpiNq$jshF*MWNoGzlSA0=w zQfX#R3ebkU;{2lcG9#`+W(|!9cA$=8pvuJb)It`B7A}|k(xjZ!V2CJJA!`I9P?13) zTLd$ZPA(`dWcQXRDdY(BX7FbA=4dPA^yctpY%ApI%mBH9yQGjOsF1g~wvaD^6R06R zH?<@Yu)IJ7P>`Weu(nVLYVcKEFoA!-4Eqx7l( literal 0 HcmV?d00001 diff --git a/source/tests/pt/models/dpa2_tebd.pth b/source/tests/pt/models/dpa2_tebd.pth deleted file mode 100644 index 3d4fc5511c93036a18be1b290fcdecc482215bf6..0000000000000000000000000000000000000000 GIT binary patch literal 0 HcmV?d00001 literal 1085 zcmWIWW@cev;NW1u0AdV047vF!sX6iGshQ~+CB^zFi6x181=%@nP8Rmh+jRLB@105mx@Hz_qGB{MHw4`Nm!Q*uduQF4Y} zd}&E$PBB+}QEF0YW==|cNornkeo=gx5mzCzhDHQCP;W6%Wny}2AqzwcmrH(WQch|x zM3k$LHG&bS$e@rdf*D9B7nByVdrOoQas+xacr$x*v=wrCb9ghh6>@cEfZW4fQpgij z$Xi@n$QQv0)DWMWT9OFzSRp@HUO*!PD9TVMSX(Frw!SR2s2J$ALSb*lVz3UlUkXKP z3q>73%-(Y`P z?Y_N0-0}TyYI%RGUVnUlUGB=6m&4xKzj|P0d}!wz`|9Kr2MOK>`?tipFM7!OWB;uA z(`2Tk*4yct0hJWn7;su)ND1z^u1=w=}MPy|JPBrvI9nxSui zZUV9kMNv$c0CX*?3D9r}@MdGvfhv__)`e>V Date: Thu, 1 Feb 2024 10:07:57 +0000 Subject: [PATCH 2/7] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/model_format/__init__.py | 8 +- deepmd/model_format/dpa1.py | 23 +- deepmd/model_format/network.py | 58 ++--- deepmd/pt/model/descriptor/dpa1.py | 32 ++- deepmd/pt/model/descriptor/se_atten.py | 328 ++++++++++++++---------- deepmd/pt/model/network/mlp.py | 106 ++++---- source/tests/pt/test_descriptor_dpa1.py | 4 +- source/tests/pt/test_dpa1.py | 91 +++++-- 8 files changed, 388 insertions(+), 262 deletions(-) diff --git a/deepmd/model_format/__init__.py b/deepmd/model_format/__init__.py index 3aa28ec192..b1814c6cb5 100644 --- a/deepmd/model_format/__init__.py +++ b/deepmd/model_format/__init__.py @@ -4,6 +4,9 @@ PRECISION_DICT, NativeOP, ) +from .dpa1 import ( + DescrptDPA1, +) from .env_mat import ( EnvMat, ) @@ -11,11 +14,11 @@ InvarFitting, ) from .network import ( + EmbdLayer, EmbeddingNet, FittingNet, - NativeLayer, - EmbdLayer, LayerNorm, + NativeLayer, NativeNet, NetworkCollection, load_dp_model, @@ -37,7 +40,6 @@ from .se_e2_a import ( DescrptSeA, ) -from .dpa1 import DescrptDPA1 __all__ = [ "InvarFitting", diff --git a/deepmd/model_format/dpa1.py b/deepmd/model_format/dpa1.py index 829339838f..cd367ab5a7 100644 --- a/deepmd/model_format/dpa1.py +++ b/deepmd/model_format/dpa1.py @@ -21,9 +21,9 @@ EnvMat, ) from .network import ( + EmbdLayer, EmbeddingNet, NetworkCollection, - EmbdLayer, ) @@ -147,6 +147,7 @@ class DescrptDPA1(NativeOP): DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation. arXiv preprint arXiv:2208.08236. """ + def __init__( self, rcut: float, @@ -183,7 +184,7 @@ def __init__( if spin is not None: raise NotImplementedError("spin is not implemented") # TODO - if tebd_input_mode != 'concat': + if tebd_input_mode != "concat": raise NotImplementedError("tebd_input_mode != 'concat' not implemented") if not smooth: raise NotImplementedError("smooth == False not implemented") @@ -215,8 +216,10 @@ def __init__( self.concat_output_tebd = concat_output_tebd self.spin = spin - self.type_embedding = EmbdLayer(ntypes, tebd_dim, padding=True, precision=precision) - in_dim = 1 + self.tebd_dim * 2 if self.tebd_input_mode in ['concat'] else 1 + self.type_embedding = EmbdLayer( + ntypes, tebd_dim, padding=True, precision=precision + ) + in_dim = 1 + self.tebd_dim * 2 if self.tebd_input_mode in ["concat"] else 1 self.embeddings = NetworkCollection( ndim=0, ntypes=self.ntypes, @@ -255,8 +258,11 @@ def __getitem__(self, key): @property def dim_out(self): """Returns the output dimension of this descriptor.""" - return self.neuron[-1] * self.axis_neuron + self.tebd_dim * 2 \ - if self.concat_output_tebd else self.neuron[-1] * self.axis_neuron + return ( + self.neuron[-1] * self.axis_neuron + self.tebd_dim * 2 + if self.concat_output_tebd + else self.neuron[-1] * self.axis_neuron + ) def cal_g( self, @@ -302,7 +308,6 @@ def call( sw The smooth switch function. """ - # nf x nloc x nnei x 4 rr, ww = self.env_mat.call(coord_ext, atype_ext, nlist, self.davg, self.dstd) nf, nloc, nnei, _ = rr.shape @@ -318,7 +323,9 @@ def call( nlist_masked[nlist_masked == -1] = 0 index = np.tile(nlist_masked.reshape(nf, -1, 1), (1, 1, self.tebd_dim)) # nf x nloc x nnei x tebd_dim - atype_embd_nlist = np.take_along_axis(atype_embd_ext, index, axis=1).reshape(nf, nloc, nnei, self.tebd_dim) + atype_embd_nlist = np.take_along_axis(atype_embd_ext, index, axis=1).reshape( + nf, nloc, nnei, self.tebd_dim + ) ng = self.neuron[-1] ss = rr[..., 0:1] ss = np.concatenate([ss, atype_embd_nlist, atype_embd_nnei], axis=-1) diff --git a/deepmd/model_format/network.py b/deepmd/model_format/network.py index 508eb07d56..2fd9727a59 100644 --- a/deepmd/model_format/network.py +++ b/deepmd/model_format/network.py @@ -342,16 +342,17 @@ def __init__( ) -> None: self.padding = padding self.num_channel = num_channel + 1 if self.padding else num_channel - super().__init__(num_in=self.num_channel, - num_out=num_out, - bias=False, - use_timestep=False, - activation_function=None, - resnet=False, - precision=precision, - ) + super().__init__( + num_in=self.num_channel, + num_out=num_out, + bias=False, + use_timestep=False, + activation_function=None, + resnet=False, + precision=precision, + ) if self.padding: - self.w[-1] = 0. + self.w[-1] = 0.0 def serialize(self) -> dict: """Serialize the layer to a dict. @@ -361,9 +362,7 @@ def serialize(self) -> dict: dict The serialized layer. """ - data = { - "w": self.w - } + data = {"w": self.w} return { "padding": self.padding, "precision": self.precision, @@ -390,9 +389,7 @@ def deserialize(cls, data: dict) -> "EmbdLayer": padding=False, **data, ) - obj.w, = ( - variables["w"], - ) + (obj.w,) = (variables["w"],) obj.padding = padding obj.check_shape_consistency() return obj @@ -464,18 +461,19 @@ def __init__( self.eps = eps self.uni_init = uni_init self.num_in = num_in - super().__init__(num_in=1, - num_out=num_in, - bias=True, - use_timestep=False, - activation_function=None, - resnet=False, - precision=precision, - ) + super().__init__( + num_in=1, + num_out=num_in, + bias=True, + use_timestep=False, + activation_function=None, + resnet=False, + precision=precision, + ) self.w = self.w.squeeze(0) # keep the weight shape to be [num_in] if self.uni_init: - self.w = 1. - self.b = 0. + self.w = 1.0 + self.b = 0.0 def serialize(self) -> dict: """Serialize the layer to a dict. @@ -510,17 +508,13 @@ def deserialize(cls, data: dict) -> "LayerNorm": assert len(variables["w"].shape) == 1 if variables["b"] is not None: assert len(variables["b"].shape) == 1 - num_in, = variables["w"].shape + (num_in,) = variables["w"].shape obj = cls( num_in, **data, ) - obj.w, = ( - variables["w"], - ) - obj.b, = ( - variables["b"], - ) + (obj.w,) = (variables["w"],) + (obj.b,) = (variables["b"],) obj._check_shape_consistency() return obj diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index 611c9b1179..db54132576 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -6,19 +6,25 @@ import torch +from deepmd.model_format import EnvMat as DPEnvMat from deepmd.pt.model.descriptor import ( Descriptor, ) +from deepmd.pt.model.network.mlp import ( + EmbdLayer, + NetworkCollection, +) from deepmd.pt.model.network.network import ( TypeEmbedNet, ) +from deepmd.pt.utils import ( + env, +) -from .se_atten import DescrptBlockSeAtten, NeighborGatedAttention -from deepmd.pt.model.network.mlp import EmbdLayer, NetworkCollection -from deepmd.model_format import ( - EnvMat as DPEnvMat, +from .se_atten import ( + DescrptBlockSeAtten, + NeighborGatedAttention, ) -from deepmd.pt.utils import env @Descriptor.register("dpa1") @@ -74,7 +80,7 @@ def __init__( normalize=normalize, temperature=temperature, old_impl=old_impl, - **kwargs + **kwargs, ) self.type_embedding_old = None self.type_embedding = None @@ -82,7 +88,9 @@ def __init__( if self.old_impl: self.type_embedding_old = TypeEmbedNet(ntypes, tebd_dim) else: - self.type_embedding = EmbdLayer(ntypes, tebd_dim, padding=True, precision=precision) + self.type_embedding = EmbdLayer( + ntypes, tebd_dim, padding=True, precision=precision + ) self.tebd_dim = tebd_dim self.concat_output_tebd = concat_output_tebd @@ -195,9 +203,9 @@ def forward( return g1, rot_mat, g2, h2, sw def set_stat_mean_and_stddev( - self, - mean: torch.Tensor, - stddev: torch.Tensor, + self, + mean: torch.Tensor, + stddev: torch.Tensor, ) -> None: self.se_atten.mean = mean self.se_atten.stddev = stddev @@ -253,5 +261,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": obj.se_atten["davg"] = t_cvt(variables["davg"]) obj.se_atten["dstd"] = t_cvt(variables["dstd"]) obj.se_atten.filter_layers = NetworkCollection.deserialize(embeddings) - obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize(attention_layers) + obj.se_atten.dpa1_attention = NeighborGatedAttention.deserialize( + attention_layers + ) return obj diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index ba7c533ff7..82681ed345 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -8,11 +8,6 @@ import torch import torch.nn as nn import torch.nn.functional as torch_func -from deepmd.pt.utils.env import ( - PRECISION_DICT, - DEFAULT_PRECISION, -) -from deepmd.pt.utils.utils import ActivationFn from deepmd.pt.model.descriptor.descriptor import ( DescriptorBlock, @@ -21,18 +16,23 @@ from deepmd.pt.model.descriptor.env_mat import ( prod_env_mat_se_a, ) +from deepmd.pt.model.network.mlp import ( + EmbeddingNet, + LayerNorm, + MLPLayer, + NetworkCollection, +) from deepmd.pt.model.network.network import ( NeighborWiseAttention, TypeFilter, ) -from deepmd.pt.model.network.mlp import EmbeddingNet, NetworkCollection, MLPLayer, LayerNorm - -from deepmd.model_format import ( - EnvMat as DPEnvMat, -) from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) @DescriptorBlock.register("se_atten") @@ -46,7 +46,7 @@ def __init__( neuron: list = [25, 50, 100], axis_neuron: int = 16, tebd_dim: int = 8, - tebd_input_mode: str = 'concat', + tebd_input_mode: str = "concat", # set_davg_zero: bool = False, set_davg_zero: bool = True, # TODO attn: int = 128, @@ -106,23 +106,30 @@ def __init__( self.nnei = sum(sel) self.ndescrpt = self.nnei * 4 if self.old_impl: - self.dpa1_attention = NeighborWiseAttention(self.attn_layer, self.nnei, self.filter_neuron[-1], - self.attn_dim, - dotr=self.attn_dotr, do_mask=self.attn_mask, - activation=self.activation_function, - scaling_factor=self.scaling_factor, - normalize=self.normalize, - temperature=self.temperature) + self.dpa1_attention = NeighborWiseAttention( + self.attn_layer, + self.nnei, + self.filter_neuron[-1], + self.attn_dim, + dotr=self.attn_dotr, + do_mask=self.attn_mask, + activation=self.activation_function, + scaling_factor=self.scaling_factor, + normalize=self.normalize, + temperature=self.temperature, + ) else: - self.dpa1_attention = NeighborGatedAttention(self.attn_layer, - self.nnei, - self.filter_neuron[-1], - self.attn_dim, - dotr=self.attn_dotr, - do_mask=self.attn_mask, - scaling_factor=self.scaling_factor, - normalize=self.normalize, - temperature=self.temperature) + self.dpa1_attention = NeighborGatedAttention( + self.attn_layer, + self.nnei, + self.filter_neuron[-1], + self.attn_dim, + dotr=self.attn_dotr, + do_mask=self.attn_mask, + scaling_factor=self.scaling_factor, + normalize=self.normalize, + temperature=self.temperature, + ) wanted_shape = (self.ntypes, self.nnei, 4) mean = torch.zeros( @@ -133,18 +140,29 @@ def __init__( ) self.register_buffer("mean", mean) self.register_buffer("stddev", stddev) - self.embd_input_dim = 1 + self.tebd_dim * 2 if self.tebd_input_mode in ['concat'] else 1 + self.embd_input_dim = ( + 1 + self.tebd_dim * 2 if self.tebd_input_mode in ["concat"] else 1 + ) self.filter_layers_old = None self.filter_layers = None if self.old_impl: filter_layers = [] - one = TypeFilter(0, self.nnei, self.filter_neuron, return_G=True, tebd_dim=self.tebd_dim, use_tebd=True, - tebd_mode=self.tebd_input_mode) + one = TypeFilter( + 0, + self.nnei, + self.filter_neuron, + return_G=True, + tebd_dim=self.tebd_dim, + use_tebd=True, + tebd_mode=self.tebd_input_mode, + ) filter_layers.append(one) self.filter_layers_old = torch.nn.ModuleList(filter_layers) else: - filter_layers = NetworkCollection(ndim=0, ntypes=len(sel), network_type="embedding_network") + filter_layers = NetworkCollection( + ndim=0, ntypes=len(sel), network_type="embedding_network" + ) filter_layers[0] = EmbeddingNet( self.embd_input_dim, self.filter_neuron, @@ -338,18 +356,26 @@ def forward( atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt) if self.old_impl: assert self.filter_layers_old is not None - dmatrix = dmatrix.view(-1, self.ndescrpt) # shape is [nframes*nall, self.ndescrpt] + dmatrix = dmatrix.view( + -1, self.ndescrpt + ) # shape is [nframes*nall, self.ndescrpt] gg = self.filter_layers_old[0]( dmatrix, atype_tebd=atype_tebd_nnei, nlist_tebd=atype_tebd_nlist, ) # shape is [nframes*nall, self.neei, out_size] - input_r = torch.nn.functional.normalize(dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1) - gg = self.dpa1_attention(gg, nlist_mask, input_r=input_r, - sw=sw) # shape is [nframes*nloc, self.neei, out_size] - inputs_reshape = dmatrix.view(-1, self.nnei, 4).permute(0, 2, - 1) # shape is [nframes*natoms[0], 4, self.neei] - xyz_scatter = torch.matmul(inputs_reshape, gg) # shape is [nframes*natoms[0], 4, out_size] + input_r = torch.nn.functional.normalize( + dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 + ) + gg = self.dpa1_attention( + gg, nlist_mask, input_r=input_r, sw=sw + ) # shape is [nframes*nloc, self.neei, out_size] + inputs_reshape = dmatrix.view(-1, self.nnei, 4).permute( + 0, 2, 1 + ) # shape is [nframes*natoms[0], 4, self.neei] + xyz_scatter = torch.matmul( + inputs_reshape, gg + ) # shape is [nframes*natoms[0], 4, out_size] else: assert self.filter_layers is not None dmatrix = dmatrix.view(-1, self.nnei, 4) @@ -357,16 +383,19 @@ def forward( # nfnl x nnei x 4 rr = dmatrix ss = rr[:, :, :1] - if self.tebd_input_mode in ['concat']: + if self.tebd_input_mode in ["concat"]: nlist_tebd = atype_tebd_nlist.reshape(nfnl, nnei, self.tebd_dim) atype_tebd = atype_tebd_nnei.reshape(nfnl, nnei, self.tebd_dim) # nfnl x nnei x (1 + tebd_dim * 2) ss = torch.concat([ss, nlist_tebd, atype_tebd], dim=2) # nfnl x nnei x ng gg = self.filter_layers._networks[0](ss) - input_r = torch.nn.functional.normalize(dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1) - gg = self.dpa1_attention(gg, nlist_mask, input_r=input_r, - sw=sw) # shape is [nframes*nloc, self.neei, out_size] + input_r = torch.nn.functional.normalize( + dmatrix.reshape(-1, self.nnei, 4)[:, :, 1:4], dim=-1 + ) + gg = self.dpa1_attention( + gg, nlist_mask, input_r=input_r, sw=sw + ) # shape is [nframes*nloc, self.neei, out_size] # nfnl x 4 x ng xyz_scatter = torch.matmul(rr.permute(0, 2, 1), gg) xyz_scatter = xyz_scatter / self.nnei @@ -386,20 +415,20 @@ def forward( class NeighborGatedAttention(nn.Module): - def __init__(self, - layer_num: int, - nnei: int, - embed_dim: int, - hidden_dim: int, - dotr: bool = False, - do_mask: bool = False, - scaling_factor: float = 1.0, - normalize: bool = True, - temperature: float = None, - precision: str = DEFAULT_PRECISION, - ): - """Construct a neighbor-wise attention net. - """ + def __init__( + self, + layer_num: int, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention net.""" super(NeighborGatedAttention, self).__init__() self.layer_num = layer_num self.nnei = nnei @@ -414,31 +443,37 @@ def __init__(self, self.network_type = NeighborGatedAttentionLayer attention_layers = [] for i in range(self.layer_num): - attention_layers.append(NeighborGatedAttentionLayer(nnei, - embed_dim, - hidden_dim, - dotr=dotr, - do_mask=do_mask, - scaling_factor=scaling_factor, - normalize=normalize, - temperature=temperature, - precision=precision)) + attention_layers.append( + NeighborGatedAttentionLayer( + nnei, + embed_dim, + hidden_dim, + dotr=dotr, + do_mask=do_mask, + scaling_factor=scaling_factor, + normalize=normalize, + temperature=temperature, + precision=precision, + ) + ) self.attention_layers = nn.ModuleList(attention_layers) def forward( - self, - input_G, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + self, + input_G, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, ): """ Args: input_G: Input G, [nframes * nloc, nnei, embed_dim] nei_mask: neighbor mask, [nframes * nloc, nnei] input_r: normalized radial, [nframes, nloc, nei, 3] - Returns: - out: Output G, [nframes * nloc, nnei, embed_dim] + + Returns + ------- + out: Output G, [nframes * nloc, nnei, embed_dim] """ out = input_G # https://github.com/pytorch/pytorch/issues/39165#issuecomment-635472592 @@ -475,6 +510,7 @@ def __setitem__(self, key, value): def serialize(self) -> dict: """Serialize the networks to a dict. + Returns ------- dict @@ -493,12 +529,13 @@ def serialize(self) -> dict: "normalize": self.normalize, "temperature": self.temperature, "precision": self.precision, - "attention_layers": [layer.serialize() for layer in self.attention_layers] + "attention_layers": [layer.serialize() for layer in self.attention_layers], } @classmethod def deserialize(cls, data: dict) -> "NeighborGatedAttention": """Deserialize the networks from a dict. + Parameters ---------- data : dict @@ -512,19 +549,19 @@ def deserialize(cls, data: dict) -> "NeighborGatedAttention": class NeighborGatedAttentionLayer(nn.Module): - def __init__(self, - nnei: int, - embed_dim: int, - hidden_dim: int, - dotr: bool = False, - do_mask: bool = False, - scaling_factor: float = 1.0, - normalize: bool = True, - temperature: float = None, - precision: str = DEFAULT_PRECISION, - ): - """Construct a neighbor-wise attention layer. - """ + def __init__( + self, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention layer.""" super(NeighborGatedAttentionLayer, self).__init__() self.nnei = nnei self.embed_dim = embed_dim @@ -535,24 +572,25 @@ def __init__(self, self.normalize = normalize self.temperature = temperature self.precision = precision - self.attention_layer = GatedAttentionLayer(nnei, - embed_dim, - hidden_dim, - dotr=dotr, - do_mask=do_mask, - scaling_factor=scaling_factor, - normalize=normalize, - temperature=temperature, - precision=precision, - ) + self.attention_layer = GatedAttentionLayer( + nnei, + embed_dim, + hidden_dim, + dotr=dotr, + do_mask=do_mask, + scaling_factor=scaling_factor, + normalize=normalize, + temperature=temperature, + precision=precision, + ) self.attn_layer_norm = LayerNorm(self.embed_dim, precision=precision) def forward( - self, - x, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, + self, + x, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, ): residual = x x = self.attention_layer(x, nei_mask, input_r=input_r, sw=sw) @@ -562,6 +600,7 @@ def forward( def serialize(self) -> dict: """Serialize the networks to a dict. + Returns ------- dict @@ -578,12 +617,13 @@ def serialize(self) -> dict: "temperature": self.temperature, "precision": self.precision, "attention_layer": self.attention_layer.serialize(), - "attn_layer_norm": self.attn_layer_norm.serialize() + "attn_layer_norm": self.attn_layer_norm.serialize(), } @classmethod def deserialize(cls, data: dict) -> "NeighborGatedAttentionLayer": """Deserialize the networks from a dict. + Parameters ---------- data : dict @@ -598,21 +638,21 @@ def deserialize(cls, data: dict) -> "NeighborGatedAttentionLayer": class GatedAttentionLayer(nn.Module): - def __init__(self, - nnei: int, - embed_dim: int, - hidden_dim: int, - dotr: bool = False, - do_mask: bool = False, - scaling_factor: float = 1.0, - normalize: bool = True, - temperature: float = None, - bias: bool = True, - smooth: bool = True, - precision: str = DEFAULT_PRECISION, - ): - """Construct a neighbor-wise attention net. - """ + def __init__( + self, + nnei: int, + embed_dim: int, + hidden_dim: int, + dotr: bool = False, + do_mask: bool = False, + scaling_factor: float = 1.0, + normalize: bool = True, + temperature: float = None, + bias: bool = True, + smooth: bool = True, + precision: str = DEFAULT_PRECISION, + ): + """Construct a neighbor-wise attention net.""" super(GatedAttentionLayer, self).__init__() self.nnei = nnei self.embed_dim = embed_dim @@ -629,26 +669,42 @@ def __init__(self, else: self.scaling = temperature self.normalize = normalize - self.in_proj = MLPLayer(embed_dim, hidden_dim * 3, bias=bias, use_timestep=False, bavg=0., stddev=1., - precision=precision) - self.out_proj = MLPLayer(hidden_dim, embed_dim, bias=bias, use_timestep=False, bavg=0., stddev=1., - precision=precision) + self.in_proj = MLPLayer( + embed_dim, + hidden_dim * 3, + bias=bias, + use_timestep=False, + bavg=0.0, + stddev=1.0, + precision=precision, + ) + self.out_proj = MLPLayer( + hidden_dim, + embed_dim, + bias=bias, + use_timestep=False, + bavg=0.0, + stddev=1.0, + precision=precision, + ) def forward( - self, - query, - nei_mask, - input_r: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - attnw_shift: float = 20.0, + self, + query, + nei_mask, + input_r: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + attnw_shift: float = 20.0, ): """ Args: query: input G, [nframes * nloc, nnei, embed_dim] nei_mask: neighbor mask, [nframes * nloc, nnei] input_r: normalized radial, [nframes, nloc, nei, 3] - Returns: - type_embedding: + + Returns + ------- + type_embedding: """ q, k, v = self.in_proj(query).chunk(3, dim=-1) # [nframes * nloc, nnei, hidden_dim] @@ -669,11 +725,15 @@ def forward( # [nframes * nloc, nnei] assert sw is not None sw = sw.view([-1, self.nnei]) - attn_weights = (attn_weights + attnw_shift) * sw[:, :, None] * sw[:, None, :] - attnw_shift + attn_weights = (attn_weights + attnw_shift) * sw[:, :, None] * sw[ + :, None, : + ] - attnw_shift else: - attn_weights = attn_weights.masked_fill(~nei_mask.unsqueeze(1), float("-inf")) + attn_weights = attn_weights.masked_fill( + ~nei_mask.unsqueeze(1), float("-inf") + ) attn_weights = torch_func.softmax(attn_weights, dim=-1) - attn_weights = attn_weights.masked_fill(~nei_mask.unsqueeze(-1), float(0.0)) + attn_weights = attn_weights.masked_fill(~nei_mask.unsqueeze(-1), 0.0) if self.smooth: assert sw is not None attn_weights = attn_weights * sw[:, :, None] * sw[:, None, :] @@ -687,6 +747,7 @@ def forward( def serialize(self) -> dict: """Serialize the networks to a dict. + Returns ------- dict @@ -707,12 +768,13 @@ def serialize(self) -> dict: "smooth": self.smooth, "precision": self.precision, "in_proj": self.in_proj.serialize(), - "out_proj": self.out_proj.serialize() + "out_proj": self.out_proj.serialize(), } @classmethod def deserialize(cls, data: dict) -> "GatedAttentionLayer": """Deserialize the networks from a dict. + Parameters ---------- data : dict diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 5bd9c1a23d..2d79fedba0 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -16,11 +16,11 @@ device = env.DEVICE +from deepmd.model_format import EmbdLayer as DPEmbdLayer +from deepmd.model_format import LayerNorm as DPLayerNorm from deepmd.model_format import ( NativeLayer, ) -from deepmd.model_format import EmbdLayer as DPEmbdLayer -from deepmd.model_format import LayerNorm as DPLayerNorm from deepmd.model_format import NetworkCollection as DPNetworkCollection from deepmd.model_format import ( make_embedding_network, @@ -193,24 +193,25 @@ def check_load_param(ss): class EmbdLayer(MLPLayer): def __init__( - self, - num_channel, - num_out, - padding: bool = True, - stddev: float = 1., - precision: str = DEFAULT_PRECISION, + self, + num_channel, + num_out, + padding: bool = True, + stddev: float = 1.0, + precision: str = DEFAULT_PRECISION, ): self.padding = padding self.num_channel = num_channel + 1 if self.padding else num_channel - super().__init__(num_in=self.num_channel, - num_out=num_out, - bias=False, - use_timestep=False, - activation_function=None, - resnet=False, - stddev=stddev, - precision=precision, - ) + super().__init__( + num_in=self.num_channel, + num_out=num_out, + bias=False, + use_timestep=False, + activation_function=None, + resnet=False, + stddev=stddev, + precision=precision, + ) if self.padding: nn.init.zeros_(self.matrix.data[-1]) @@ -218,14 +219,14 @@ def dim_channel(self) -> int: return self.matrix.shape[0] def forward( - self, - xx: torch.Tensor, + self, + xx: torch.Tensor, ) -> torch.Tensor: """One Embedding layer used by DP model. Parameters ---------- - xx: torch.Tensor + xx : torch.Tensor The input of index. Returns @@ -274,36 +275,41 @@ def deserialize(cls, data: dict) -> "EmbdLayer": ) obj.padding = padding prec = PRECISION_DICT[obj.precision] - check_load_param = \ - lambda ss: nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) \ - if nl[ss] is not None else None + check_load_param = ( + lambda ss: nn.Parameter( + data=torch.tensor(nl[ss], dtype=prec, device=device) + ) + if nl[ss] is not None + else None + ) obj.matrix = check_load_param("matrix") return obj class LayerNorm(MLPLayer): def __init__( - self, - num_in, - eps: float = 1e-5, - uni_init: bool = True, - bavg: float = 0., - stddev: float = 1., - precision: str = DEFAULT_PRECISION, + self, + num_in, + eps: float = 1e-5, + uni_init: bool = True, + bavg: float = 0.0, + stddev: float = 1.0, + precision: str = DEFAULT_PRECISION, ): self.eps = eps self.uni_init = uni_init self.num_in = num_in - super().__init__(num_in=1, - num_out=num_in, - bias=True, - use_timestep=False, - activation_function=None, - resnet=False, - bavg=bavg, - stddev=stddev, - precision=precision, - ) + super().__init__( + num_in=1, + num_out=num_in, + bias=True, + use_timestep=False, + activation_function=None, + resnet=False, + bavg=bavg, + stddev=stddev, + precision=precision, + ) self.matrix = torch.nn.Parameter(self.matrix.squeeze(0)) if self.uni_init: nn.init.ones_(self.matrix.data) @@ -313,14 +319,14 @@ def dim_out(self) -> int: return self.matrix.shape[0] def forward( - self, - xx: torch.Tensor, + self, + xx: torch.Tensor, ) -> torch.Tensor: """One Layer Norm used by DP model. Parameters ---------- - xx: torch.Tensor + xx : torch.Tensor The input of index. Returns @@ -328,7 +334,9 @@ def forward( yy: torch.Tensor The output. """ - yy = torch_func.layer_norm(xx, tuple((self.num_in,)), self.matrix, self.bias, self.eps) + yy = torch_func.layer_norm( + xx, tuple((self.num_in,)), self.matrix, self.bias, self.eps + ) return yy def serialize(self) -> dict: @@ -365,9 +373,13 @@ def deserialize(cls, data: dict) -> "LayerNorm": precision=nl["precision"], ) prec = PRECISION_DICT[obj.precision] - check_load_param = \ - lambda ss: nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) \ - if nl[ss] is not None else None + check_load_param = ( + lambda ss: nn.Parameter( + data=torch.tensor(nl[ss], dtype=prec, device=device) + ) + if nl[ss] is not None + else None + ) obj.matrix = check_load_param("matrix") obj.bias = check_load_param("bias") return obj diff --git a/source/tests/pt/test_descriptor_dpa1.py b/source/tests/pt/test_descriptor_dpa1.py index 2caeb5890e..a7a0f9a4d6 100644 --- a/source/tests/pt/test_descriptor_dpa1.py +++ b/source/tests/pt/test_descriptor_dpa1.py @@ -12,7 +12,9 @@ DescrptBlockSeAtten, DescrptDPA1, ) -from deepmd.pt.model.network.mlp import EmbdLayer +from deepmd.pt.model.network.mlp import ( + EmbdLayer, +) from deepmd.pt.utils import ( env, ) diff --git a/source/tests/pt/test_dpa1.py b/source/tests/pt/test_dpa1.py index 400a611c05..ffd5a25bd5 100644 --- a/source/tests/pt/test_dpa1.py +++ b/source/tests/pt/test_dpa1.py @@ -1,12 +1,13 @@ -import torch, copy -import unittest +# SPDX-License-Identifier: LGPL-3.0-or-later import itertools +import unittest + import numpy as np +import torch try: - from deepmd.model_format import ( - DescrptDPA1 as DPDescrptDPA1 - ) + from deepmd.model_format import DescrptDPA1 as DPDescrptDPA1 + support_se_atten = True except ModuleNotFoundError: support_se_atten = False @@ -14,7 +15,7 @@ support_se_atten = False from deepmd.pt.model.descriptor.dpa1 import ( - DescrptDPA1 + DescrptDPA1, ) from deepmd.pt.utils import ( env, @@ -32,13 +33,14 @@ dtype = env.GLOBAL_PT_FLOAT_PRECISION + @unittest.skipIf(not support_se_atten, "EnvMat not supported") class TestDescrptSeAtten(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self): TestCaseSingleFrameWithNlist.setUp(self) def test_consistency( - self, + self, ): rng = np.random.default_rng(100) nf, nloc, nnei = self.nlist.shape @@ -47,15 +49,19 @@ def test_consistency( dstd = 0.1 + np.abs(dstd) for idt, prec in itertools.product( - [False, True], - ["float64", "float32"], + [False, True], + ["float64", "float32"], ): dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) err_msg = f"idt={idt} prec={prec}" # dpa1 new impl dd0 = DescrptDPA1( - self.rcut, self.rcut_smth, self.sel, self.nt, attn_layer=0, # TODO add support for non-zero layer + self.rcut, + self.rcut_smth, + self.sel, + self.nt, + attn_layer=0, # TODO add support for non-zero layer # precision=prec, # resnet_dt=idt, old_impl=False, @@ -75,22 +81,34 @@ def test_consistency( torch.tensor(self.nlist, dtype=int, device=env.DEVICE), ) np.testing.assert_allclose( - rd0.detach().cpu().numpy(), rd1.detach().cpu().numpy(), - rtol=rtol, atol=atol, err_msg=err_msg, + rd0.detach().cpu().numpy(), + rd1.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, ) # dp impl dd2 = DPDescrptDPA1.deserialize(dd0.serialize()) rd2, _, _, _, _ = dd2.call( - self.coord_ext, self.atype_ext, self.nlist, + self.coord_ext, + self.atype_ext, + self.nlist, ) np.testing.assert_allclose( - rd0.detach().cpu().numpy(), rd2, - rtol=rtol, atol=atol, err_msg=err_msg, + rd0.detach().cpu().numpy(), + rd2, + rtol=rtol, + atol=atol, + err_msg=err_msg, ) # old impl if idt is False and prec == "float64": dd3 = DescrptDPA1( - self.rcut, self.rcut_smth, self.sel, self.nt, attn_layer=0, # TODO add support for non-zero layer + self.rcut, + self.rcut_smth, + self.sel, + self.nt, + attn_layer=0, # TODO add support for non-zero layer # precision=prec, # resnet_dt=idt, old_impl=True, @@ -101,16 +119,29 @@ def test_consistency( dd0_state_dict_attn = dd0.se_atten.dpa1_attention.state_dict() dd3_state_dict_attn = dd3.se_atten.dpa1_attention.state_dict() for i in dd3_state_dict: - dd3_state_dict[i] = dd0_state_dict[i.replace('.deep_layers.', '.layers.') - .replace('filter_layers_old.', 'filter_layers._networks.').replace('.attn_layer_norm.weight', '.attn_layer_norm.matrix')].detach().clone() - if '.bias' in i and 'attn_layer_norm' not in i: + dd3_state_dict[i] = ( + dd0_state_dict[ + i.replace(".deep_layers.", ".layers.") + .replace("filter_layers_old.", "filter_layers._networks.") + .replace( + ".attn_layer_norm.weight", ".attn_layer_norm.matrix" + ) + ] + .detach() + .clone() + ) + if ".bias" in i and "attn_layer_norm" not in i: dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0) dd3.se_atten.load_state_dict(dd3_state_dict) dd0_state_dict_tebd = dd0.type_embedding.state_dict() dd3_state_dict_tebd = dd3.type_embedding_old.state_dict() for i in dd3_state_dict_tebd: - dd3_state_dict_tebd[i] = dd0_state_dict_tebd[i.replace('embedding.weight', 'matrix')].detach().clone() + dd3_state_dict_tebd[i] = ( + dd0_state_dict_tebd[i.replace("embedding.weight", "matrix")] + .detach() + .clone() + ) dd3.type_embedding_old.load_state_dict(dd3_state_dict_tebd) rd3, _, _, _, _ = dd3( @@ -119,12 +150,15 @@ def test_consistency( torch.tensor(self.nlist, dtype=int, device=env.DEVICE), ) np.testing.assert_allclose( - rd0.detach().cpu().numpy(), rd3.detach().cpu().numpy(), - rtol=rtol, atol=atol, err_msg=err_msg, + rd0.detach().cpu().numpy(), + rd3.detach().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=err_msg, ) def test_jit( - self, + self, ): rng = np.random.default_rng() nf, nloc, nnei = self.nlist.shape @@ -133,15 +167,18 @@ def test_jit( dstd = 0.1 + np.abs(dstd) for idt, prec in itertools.product( - [False, True], - ["float64", "float32"], + [False, True], + ["float64", "float32"], ): dtype = PRECISION_DICT[prec] rtol, atol = get_tols(prec) err_msg = f"idt={idt} prec={prec}" # sea new impl dd0 = DescrptDPA1( - self.rcut, self.rcut_smth, self.sel, self.nt, + self.rcut, + self.rcut_smth, + self.sel, + self.nt, # precision=prec, # resnet_dt=idt, old_impl=False, @@ -150,4 +187,4 @@ def test_jit( dd0.se_atten.dstd = torch.tensor(dstd, dtype=dtype, device=env.DEVICE) # dd1 = DescrptDPA1.deserialize(dd0.serialize()) model = torch.jit.script(dd0) - # model = torch.jit.script(dd1) \ No newline at end of file + # model = torch.jit.script(dd1) From b4770de7fca6ce81e79c4312342ce65efb7194e2 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 1 Feb 2024 18:25:35 +0800 Subject: [PATCH 3/7] fix pre-commit --- deepmd/model_format/dpa1.py | 5 +++-- deepmd/model_format/network.py | 2 +- deepmd/pt/model/descriptor/dpa1.py | 5 ++++- deepmd/pt/model/descriptor/se_atten.py | 24 ++++++++++----------- deepmd/pt/model/network/mlp.py | 30 +++++++++++++------------- 5 files changed, 35 insertions(+), 31 deletions(-) diff --git a/deepmd/model_format/dpa1.py b/deepmd/model_format/dpa1.py index cd367ab5a7..e0cf193ce6 100644 --- a/deepmd/model_format/dpa1.py +++ b/deepmd/model_format/dpa1.py @@ -28,8 +28,9 @@ class DescrptDPA1(NativeOP): - r"""Attention-based descriptor :math:`\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}`, - which is proposed in pretrainable DPA-1[1] model, is given by + r"""Attention-based descriptor which is proposed in the pretrainable DPA-1[1] model. + + This descriptor, :math:`\mathcal{D}^i \in \mathbb{R}^{M \times M_{<}}`, is given by .. math:: \mathcal{D}^i = \frac{1}{N_c^2}(\hat{\mathcal{G}}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \hat{\mathcal{G}}^i_<, diff --git a/deepmd/model_format/network.py b/deepmd/model_format/network.py index 2fd9727a59..695ebafdc1 100644 --- a/deepmd/model_format/network.py +++ b/deepmd/model_format/network.py @@ -567,7 +567,7 @@ def call(self, x: np.ndarray) -> np.ndarray: """ if self.w is None or self.b is None: raise ValueError("w/b must be set") - y = self.layer_norm_numpy(x, tuple((self.num_in,)), self.w, self.b, self.eps) + y = self.layer_norm_numpy(x, (self.num_in,), self.w, self.b, self.eps) return y @staticmethod diff --git a/deepmd/pt/model/descriptor/dpa1.py b/deepmd/pt/model/descriptor/dpa1.py index db54132576..0427a101ae 100644 --- a/deepmd/pt/model/descriptor/dpa1.py +++ b/deepmd/pt/model/descriptor/dpa1.py @@ -256,7 +256,10 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": attention_layers = data.pop("attention_layers") env_mat = data.pop("env_mat") obj = cls(**data) - t_cvt = lambda xx: torch.tensor(xx, dtype=obj.se_atten.prec, device=env.DEVICE) + + def t_cvt(xx): + return torch.tensor(xx, dtype=obj.se_atten.prec, device=env.DEVICE) + obj.type_embedding = EmbdLayer.deserialize(type_embedding) obj.se_atten["davg"] = t_cvt(variables["davg"]) obj.se_atten["dstd"] = t_cvt(variables["dstd"]) diff --git a/deepmd/pt/model/descriptor/se_atten.py b/deepmd/pt/model/descriptor/se_atten.py index 82681ed345..08660768a5 100644 --- a/deepmd/pt/model/descriptor/se_atten.py +++ b/deepmd/pt/model/descriptor/se_atten.py @@ -425,11 +425,11 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: float = None, + temperature: Optional[float] = None, precision: str = DEFAULT_PRECISION, ): """Construct a neighbor-wise attention net.""" - super(NeighborGatedAttention, self).__init__() + super().__init__() self.layer_num = layer_num self.nnei = nnei self.embed_dim = embed_dim @@ -467,9 +467,9 @@ def forward( ): """ Args: - input_G: Input G, [nframes * nloc, nnei, embed_dim] - nei_mask: neighbor mask, [nframes * nloc, nnei] - input_r: normalized radial, [nframes, nloc, nei, 3] + input_G: Input G, [nframes * nloc, nnei, embed_dim]. + nei_mask: neighbor mask, [nframes * nloc, nnei]. + input_r: normalized radial, [nframes, nloc, nei, 3]. Returns ------- @@ -558,11 +558,11 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: float = None, + temperature: Optional[float] = None, precision: str = DEFAULT_PRECISION, ): """Construct a neighbor-wise attention layer.""" - super(NeighborGatedAttentionLayer, self).__init__() + super().__init__() self.nnei = nnei self.embed_dim = embed_dim self.hidden_dim = hidden_dim @@ -647,13 +647,13 @@ def __init__( do_mask: bool = False, scaling_factor: float = 1.0, normalize: bool = True, - temperature: float = None, + temperature: Optional[float] = None, bias: bool = True, smooth: bool = True, precision: str = DEFAULT_PRECISION, ): """Construct a neighbor-wise attention net.""" - super(GatedAttentionLayer, self).__init__() + super().__init__() self.nnei = nnei self.embed_dim = embed_dim self.hidden_dim = hidden_dim @@ -698,9 +698,9 @@ def forward( ): """ Args: - query: input G, [nframes * nloc, nnei, embed_dim] - nei_mask: neighbor mask, [nframes * nloc, nnei] - input_r: normalized radial, [nframes, nloc, nei, 3] + query: input G, [nframes * nloc, nnei, embed_dim]. + nei_mask: neighbor mask, [nframes * nloc, nnei]. + input_r: normalized radial, [nframes, nloc, nei, 3]. Returns ------- diff --git a/deepmd/pt/model/network/mlp.py b/deepmd/pt/model/network/mlp.py index 2d79fedba0..0cfb407f6e 100644 --- a/deepmd/pt/model/network/mlp.py +++ b/deepmd/pt/model/network/mlp.py @@ -275,13 +275,14 @@ def deserialize(cls, data: dict) -> "EmbdLayer": ) obj.padding = padding prec = PRECISION_DICT[obj.precision] - check_load_param = ( - lambda ss: nn.Parameter( - data=torch.tensor(nl[ss], dtype=prec, device=device) + + def check_load_param(ss): + return ( + nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) + if nl[ss] is not None + else None ) - if nl[ss] is not None - else None - ) + obj.matrix = check_load_param("matrix") return obj @@ -334,9 +335,7 @@ def forward( yy: torch.Tensor The output. """ - yy = torch_func.layer_norm( - xx, tuple((self.num_in,)), self.matrix, self.bias, self.eps - ) + yy = torch_func.layer_norm(xx, (self.num_in,), self.matrix, self.bias, self.eps) return yy def serialize(self) -> dict: @@ -373,13 +372,14 @@ def deserialize(cls, data: dict) -> "LayerNorm": precision=nl["precision"], ) prec = PRECISION_DICT[obj.precision] - check_load_param = ( - lambda ss: nn.Parameter( - data=torch.tensor(nl[ss], dtype=prec, device=device) + + def check_load_param(ss): + return ( + nn.Parameter(data=torch.tensor(nl[ss], dtype=prec, device=device)) + if nl[ss] is not None + else None ) - if nl[ss] is not None - else None - ) + obj.matrix = check_load_param("matrix") obj.bias = check_load_param("bias") return obj From 8fec4cb3639cd30de465a78ab41039a9a6f18377 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 1 Feb 2024 21:58:35 +0800 Subject: [PATCH 4/7] fix uts --- deepmd/pt/model/task/ener.py | 4 ++-- source/tests/pt/test_ener_fitting.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index e40a6bda44..526a6cc4f0 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -340,11 +340,11 @@ def forward( else: if self.use_tebd: atom_energy = ( - self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] + self.filter_layers._networks[0](xx) + self.bias_atom_e[atype] ) outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] else: - for type_i, ll in enumerate(self.filter_layers.networks): + for type_i, ll in enumerate(self.filter_layers._networks): mask = (atype == type_i).unsqueeze(-1) mask = torch.tile(mask, (1, 1, self.dim_out)) atom_energy = ll(xx) diff --git a/source/tests/pt/test_ener_fitting.py b/source/tests/pt/test_ener_fitting.py index eece8447df..4af40a22ad 100644 --- a/source/tests/pt/test_ener_fitting.py +++ b/source/tests/pt/test_ener_fitting.py @@ -126,7 +126,7 @@ def test_new_old( dd1 = ft1.state_dict() for kk, vv in dd1.items(): new_kk = kk - new_kk = new_kk.replace("filter_layers_old", "filter_layers.networks") + new_kk = new_kk.replace("filter_layers_old", "filter_layers._networks") new_kk = new_kk.replace("deep_layers", "layers") new_kk = new_kk.replace("final_layer", "layers.3") dd1[kk] = dd0[new_kk] From 42c75def2de59ef5fc869ca98edbd7778de4d262 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 1 Feb 2024 22:38:17 +0800 Subject: [PATCH 5/7] fix uts --- source/tests/pt/test_descriptor_dpa2.py | 2 +- source/tests/pt/test_fitting_net.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/pt/test_descriptor_dpa2.py b/source/tests/pt/test_descriptor_dpa2.py index 22371f612d..093f53bc32 100644 --- a/source/tests/pt/test_descriptor_dpa2.py +++ b/source/tests/pt/test_descriptor_dpa2.py @@ -112,7 +112,7 @@ def setUp(self): with open(Path(CUR_DIR) / "models" / "dpa2_hyb.json") as fp: self.model_json = json.load(fp) self.file_model_param = Path(CUR_DIR) / "models" / "dpa2.pt" - self.file_type_embed = Path(CUR_DIR) / "models" / "dpa2_tebd.pth" + self.file_type_embed = Path(CUR_DIR) / "models" / "dpa2_tebd.pt" def test_descriptor_hyb(self): # torch.manual_seed(0) diff --git a/source/tests/pt/test_fitting_net.py b/source/tests/pt/test_fitting_net.py index e12a397347..9bd0246ce9 100644 --- a/source/tests/pt/test_fitting_net.py +++ b/source/tests/pt/test_fitting_net.py @@ -111,7 +111,7 @@ def test_consistency(self): ).to(env.DEVICE) for name, param in my_fn.named_parameters(): matched = re.match( - "filter_layers\.networks\.(\d).layers\.(\d)\.([a-z]+)", name + "filter_layers\._networks\.(\d).layers\.(\d)\.([a-z]+)", name ) key = None if matched: From 603128add336511b961ebcf6882a3213cb25e375 Mon Sep 17 00:00:00 2001 From: Duo <50307526+iProzd@users.noreply.github.com> Date: Thu, 1 Feb 2024 23:05:52 +0800 Subject: [PATCH 6/7] fix uts --- source/tests/pt/models/dpa2_tebd.pth | Bin 0 -> 1085 bytes source/tests/pt/test_descriptor_dpa2.py | 2 +- 2 files changed, 1 insertion(+), 1 deletion(-) create mode 100644 source/tests/pt/models/dpa2_tebd.pth diff --git a/source/tests/pt/models/dpa2_tebd.pth b/source/tests/pt/models/dpa2_tebd.pth new file mode 100644 index 0000000000000000000000000000000000000000..3d4fc5511c93036a18be1b290fcdecc482215bf6 GIT binary patch literal 1085 zcmWIWW@cev;NW1u0AdV047vF!sX6iGshQ~+CB^zFi6x181=%@nP8Rmh+jRLB@105mx@Hz_qGB{MHw4`Nm!Q*uduQF4Y} zd}&E$PBB+}QEF0YW==|cNornkeo=gx5mzCzhDHQCP;W6%Wny}2AqzwcmrH(WQch|x zM3k$LHG&bS$e@rdf*D9B7nByVdrOoQas+xacr$x*v=wrCb9ghh6>@cEfZW4fQpgij z$Xi@n$QQv0)DWMWT9OFzSRp@HUO*!PD9TVMSX(Frw!SR2s2J$ALSb*lVz3UlUkXKP z3q>73%-(Y`P z?Y_N0-0}TyYI%RGUVnUlUGB=6m&4xKzj|P0d}!wz`|9Kr2MOK>`?tipFM7!OWB;uA z(`2Tk*4yct0hJWn7;su)ND1z^u1=w=}MPy|JPBrvI9nxSui zZUV9kMNv$c0CX*?3D9r}@MdGvfhv__)`e>V Date: Fri, 2 Feb 2024 00:31:40 +0800 Subject: [PATCH 7/7] Add dp impl serialization test --- deepmd/model_format/dpa1.py | 2 +- source/tests/pt/test_dpa1.py | 40 ++++++++++++++++++++++++------------ 2 files changed, 28 insertions(+), 14 deletions(-) diff --git a/deepmd/model_format/dpa1.py b/deepmd/model_format/dpa1.py index e0cf193ce6..3ca28a4fae 100644 --- a/deepmd/model_format/dpa1.py +++ b/deepmd/model_format/dpa1.py @@ -390,7 +390,7 @@ def deserialize(cls, data: dict) -> "DescrptDPA1": variables = data.pop("@variables") embeddings = data.pop("embeddings") type_embedding = data.pop("type_embedding") - attention_layers = data.pop("attention_layers") + attention_layers = data.pop("attention_layers", None) env_mat = data.pop("env_mat") obj = cls(**data) obj["davg"] = variables["davg"] diff --git a/source/tests/pt/test_dpa1.py b/source/tests/pt/test_dpa1.py index ffd5a25bd5..ae34305064 100644 --- a/source/tests/pt/test_dpa1.py +++ b/source/tests/pt/test_dpa1.py @@ -101,9 +101,23 @@ def test_consistency( atol=atol, err_msg=err_msg, ) + # dp impl serialization + dd3 = DPDescrptDPA1.deserialize(dd2.serialize()) + rd3, _, _, _, _ = dd3.call( + self.coord_ext, + self.atype_ext, + self.nlist, + ) + np.testing.assert_allclose( + rd0.detach().cpu().numpy(), + rd3, + rtol=rtol, + atol=atol, + err_msg=err_msg, + ) # old impl if idt is False and prec == "float64": - dd3 = DescrptDPA1( + dd4 = DescrptDPA1( self.rcut, self.rcut_smth, self.sel, @@ -114,12 +128,12 @@ def test_consistency( old_impl=True, ).to(env.DEVICE) dd0_state_dict = dd0.se_atten.state_dict() - dd3_state_dict = dd3.se_atten.state_dict() + dd4_state_dict = dd4.se_atten.state_dict() dd0_state_dict_attn = dd0.se_atten.dpa1_attention.state_dict() - dd3_state_dict_attn = dd3.se_atten.dpa1_attention.state_dict() - for i in dd3_state_dict: - dd3_state_dict[i] = ( + dd4_state_dict_attn = dd4.se_atten.dpa1_attention.state_dict() + for i in dd4_state_dict: + dd4_state_dict[i] = ( dd0_state_dict[ i.replace(".deep_layers.", ".layers.") .replace("filter_layers_old.", "filter_layers._networks.") @@ -131,27 +145,27 @@ def test_consistency( .clone() ) if ".bias" in i and "attn_layer_norm" not in i: - dd3_state_dict[i] = dd3_state_dict[i].unsqueeze(0) - dd3.se_atten.load_state_dict(dd3_state_dict) + dd4_state_dict[i] = dd4_state_dict[i].unsqueeze(0) + dd4.se_atten.load_state_dict(dd4_state_dict) dd0_state_dict_tebd = dd0.type_embedding.state_dict() - dd3_state_dict_tebd = dd3.type_embedding_old.state_dict() - for i in dd3_state_dict_tebd: - dd3_state_dict_tebd[i] = ( + dd4_state_dict_tebd = dd4.type_embedding_old.state_dict() + for i in dd4_state_dict_tebd: + dd4_state_dict_tebd[i] = ( dd0_state_dict_tebd[i.replace("embedding.weight", "matrix")] .detach() .clone() ) - dd3.type_embedding_old.load_state_dict(dd3_state_dict_tebd) + dd4.type_embedding_old.load_state_dict(dd4_state_dict_tebd) - rd3, _, _, _, _ = dd3( + rd4, _, _, _, _ = dd4( torch.tensor(self.coord_ext, dtype=dtype, device=env.DEVICE), torch.tensor(self.atype_ext, dtype=int, device=env.DEVICE), torch.tensor(self.nlist, dtype=int, device=env.DEVICE), ) np.testing.assert_allclose( rd0.detach().cpu().numpy(), - rd3.detach().cpu().numpy(), + rd4.detach().cpu().numpy(), rtol=rtol, atol=atol, err_msg=err_msg,