From b54e10963e52b1beeee930a956d325bb6fbf0277 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Tue, 13 Feb 2024 18:37:16 +0800 Subject: [PATCH 01/37] feat: redo pt dipole --- deepmd/pt/model/task/dipole.py | 183 +++++++++++++++++++++++++++++---- 1 file changed, 162 insertions(+), 21 deletions(-) diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index d911613a5b..e24cb8600d 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -2,20 +2,51 @@ import logging import torch +from typing import ( + List, + Optional, + Tuple, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) +import numpy as np -from deepmd.pt.model.network.network import ( - ResidualDeep, +from deepmd.pt.utils import ( + env, +) + +from deepmd.pt.model.network.mlp import ( + NetworkCollection, + FittingNet, ) from deepmd.pt.model.task.fitting import ( Fitting, ) +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + log = logging.getLogger(__name__) -class DipoleFittingNetType(Fitting): +class DipoleFittingNet(Fitting): def __init__( - self, ntypes, embedding_width, neuron, out_dim, resnet_dt=True, **kwargs + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, + **kwargs, ): """Construct a fitting net for dipole. @@ -27,22 +58,79 @@ def __init__( - resnet_dt: Using time-step in the ResNet construction. """ super().__init__() + self.var_name = var_name self.ntypes = ntypes - self.embedding_width = embedding_width - self.out_dim = out_dim + self.dim_descrpt = dim_descrpt + self.dim_out = dim_out + self.neuron = neuron + self.distinguish_types = distinguish_types + self.use_tebd = not self.distinguish_types + self.resnet_dt = resnet_dt + self.numb_fparam = numb_fparam + self.numb_aparam = numb_aparam + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] - filter_layers = [] - one = ResidualDeep( - 0, embedding_width, neuron, 0.0, out_dim=self.out_dim, resnet_dt=resnet_dt + # init constants + if self.numb_fparam > 0: + self.register_buffer( + "fparam_avg", + torch.zeros(self.numb_fparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "fparam_inv_std", + torch.ones(self.numb_fparam, dtype=self.prec, device=device), + ) + else: + self.fparam_avg, self.fparam_inv_std = None, None + if self.numb_aparam > 0: + self.register_buffer( + "aparam_avg", + torch.zeros(self.numb_aparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "aparam_inv_std", + torch.ones(self.numb_aparam, dtype=self.prec, device=device), + ) + else: + self.aparam_avg, self.aparam_inv_std = None, None + + in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam + out_dim = 3 + + self.filter_layers = NetworkCollection( + 1 if self.distinguish_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim, + out_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if self.distinguish_types else 1) + ], ) - filter_layers.append(one) - self.filter_layers = torch.nn.ModuleList(filter_layers) if "seed" in kwargs: log.info("Set seed to %d in fitting net.", kwargs["seed"]) torch.manual_seed(kwargs["seed"]) - def forward(self, inputs, atype, atype_tebd, rot_mat): + def forward( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): """Based on embedding net output, alculate total energy. Args: @@ -55,13 +143,66 @@ def forward(self, inputs, atype, atype_tebd, rot_mat): ------- - vec_out: output vector. Its shape is [nframes, nloc, 3]. """ - nframes, nloc, _ = inputs.size() - if atype_tebd is not None: - inputs = torch.concat([inputs, atype_tebd], dim=-1) - vec_out = self.filter_layers[0](inputs) # Shape is [nframes, nloc, m1] - assert list(vec_out.size()) == [nframes, nloc, self.out_dim] - vec_out = vec_out.view(-1, 1, self.out_dim) - vec_out = ( - torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3) + xx = descriptor + nframes, nloc, nd = xx.shape + # check input dim + if nd != self.dim_descrpt: + raise ValueError( + "get an input descriptor of dim {nd}," + "which is not consistent with {self.dim_descrpt}." + ) + # check fparam dim, concate to input descriptor + if self.numb_fparam > 0: + assert fparam is not None, "fparam should not be None" + assert self.fparam_avg is not None + assert self.fparam_inv_std is not None + if fparam.shape[-1] != self.numb_fparam: + raise ValueError( + "get an input fparam of dim {fparam.shape[-1]}, ", + "which is not consistent with {self.numb_fparam}.", + ) + fparam = fparam.view([nframes, self.numb_fparam]) + nb, _ = fparam.shape + t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) + t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) + fparam = (fparam - t_fparam_avg) * t_fparam_inv_std + fparam = torch.tile(fparam.reshape([nframes, 1, -1]), [1, nloc, 1]) + xx = torch.cat( + [xx, fparam], + dim=-1, + ) + # check aparam dim, concate to input descriptor + if self.numb_aparam > 0: + assert aparam is not None, "aparam should not be None" + assert self.aparam_avg is not None + assert self.aparam_inv_std is not None + if aparam.shape[-1] != self.numb_aparam: + raise ValueError( + "get an input aparam of dim {aparam.shape[-1]}, ", + "which is not consistent with {self.numb_aparam}.", + ) + aparam = aparam.view([nframes, nloc, self.numb_aparam]) + nb, nloc, _ = aparam.shape + t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) + t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) + aparam = (aparam - t_aparam_avg) * t_aparam_inv_std + xx = torch.cat( + [xx, aparam], + dim=-1, + ) + + outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion + if self.use_tebd: + atom_dipole = self.filter_layers.networks[0](xx) + outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] + else: + 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_dipole = ll(xx) + atom_dipole = atom_dipole * mask + outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] + outs = ( + torch.bmm(outs, gr).squeeze(-2).view(nframes, nloc, 3) ) # Shape is [nframes, nloc, 3] - return vec_out + return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} From ca75551631e860cde69a5241bc5f33ccc9f2226f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 13 Feb 2024 10:38:58 +0000 Subject: [PATCH 02/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/dipole.py | 25 +++++++++++-------------- 1 file changed, 11 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index e24cb8600d..07a1531f33 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -1,29 +1,26 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging - -import torch from typing import ( List, Optional, - Tuple, -) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, ) -import numpy as np -from deepmd.pt.utils import ( - env, -) +import torch from deepmd.pt.model.network.mlp import ( - NetworkCollection, FittingNet, + NetworkCollection, ) from deepmd.pt.model.task.fitting import ( Fitting, ) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) dtype = env.GLOBAL_PT_FLOAT_PRECISION device = env.DEVICE @@ -193,8 +190,8 @@ def forward( outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion if self.use_tebd: - atom_dipole = self.filter_layers.networks[0](xx) - outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] + atom_dipole = self.filter_layers.networks[0](xx) + outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] else: for type_i, ll in enumerate(self.filter_layers.networks): mask = (atype == type_i).unsqueeze(-1) From bdd2b5c325620b3df879a6c18e3ff387489bc529 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 14 Feb 2024 19:38:58 +0800 Subject: [PATCH 03/37] fix: numpy warning --- deepmd/dpmodel/model/linear_atomic_model.py | 6 ++++-- deepmd/pt/model/task/__init__.py | 4 ++-- deepmd/pt/model/task/dipole.py | 4 ++-- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/deepmd/dpmodel/model/linear_atomic_model.py b/deepmd/dpmodel/model/linear_atomic_model.py index 8a53cb9229..89df234932 100644 --- a/deepmd/dpmodel/model/linear_atomic_model.py +++ b/deepmd/dpmodel/model/linear_atomic_model.py @@ -318,14 +318,16 @@ def _compute_weight( ), axis=-1, ) # handle masked nnei. - sigma = numerator / denominator + with np.errstate(divide='ignore', invalid='ignore'): + sigma = numerator / denominator u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin) coef = np.zeros_like(u) left_mask = sigma < self.sw_rmin mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax) right_mask = sigma >= self.sw_rmax coef[left_mask] = 1 - smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 + with np.errstate(invalid='ignore'): + smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 coef[mid_mask] = smooth[mid_mask] coef[right_mask] = 0 self.zbl_weight = coef diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index 0b21033d31..180e5a5211 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -9,7 +9,7 @@ DenoiseNet, ) from .dipole import ( - DipoleFittingNetType, + DipoleFittingNet, ) from .ener import ( EnergyFittingNet, @@ -25,7 +25,7 @@ __all__ = [ "FittingNetAttenLcc", "DenoiseNet", - "DipoleFittingNetType", + "DipoleFittingNet", "EnergyFittingNet", "EnergyFittingNetDirect", "Fitting", diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 07a1531f33..50b4ebff3e 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -191,14 +191,14 @@ def forward( outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion if self.use_tebd: atom_dipole = self.filter_layers.networks[0](xx) - outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] + outs = outs + atom_dipole # Shape is [nframes, nloc, 3] else: 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_dipole = ll(xx) atom_dipole = atom_dipole * mask - outs = outs + atom_dipole # Shape is [nframes, natoms[0], 3] + outs = outs + atom_dipole # Shape is [nframes, nloc, 3] outs = ( torch.bmm(outs, gr).squeeze(-2).view(nframes, nloc, 3) ) # Shape is [nframes, nloc, 3] From 3797ac07610659c621f41655dc53a0a922015d7f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Feb 2024 11:40:21 +0000 Subject: [PATCH 04/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/model/linear_atomic_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/model/linear_atomic_model.py b/deepmd/dpmodel/model/linear_atomic_model.py index 89df234932..24ad89ebcd 100644 --- a/deepmd/dpmodel/model/linear_atomic_model.py +++ b/deepmd/dpmodel/model/linear_atomic_model.py @@ -318,7 +318,7 @@ def _compute_weight( ), axis=-1, ) # handle masked nnei. - with np.errstate(divide='ignore', invalid='ignore'): + with np.errstate(divide="ignore", invalid="ignore"): sigma = numerator / denominator u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin) coef = np.zeros_like(u) @@ -326,7 +326,7 @@ def _compute_weight( mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax) right_mask = sigma >= self.sw_rmax coef[left_mask] = 1 - with np.errstate(invalid='ignore'): + with np.errstate(invalid="ignore"): smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 coef[mid_mask] = smooth[mid_mask] coef[right_mask] = 0 From 86a6713e18547df773086b6b9b17f42d01cbb1bc Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 14 Feb 2024 20:56:29 +0800 Subject: [PATCH 05/37] chore: refactor InvarFitting --- deepmd/pt/model/task/ener.py | 265 ++++++++++++++++++++--------------- 1 file changed, 155 insertions(+), 110 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 1b3e2c3d65..00271e1493 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -46,8 +46,7 @@ log = logging.getLogger(__name__) -@fitting_check_output -class InvarFitting(Fitting): +class GeneralFitting(Fitting): def __init__( self, var_name: str, @@ -87,13 +86,7 @@ def __init__( self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] - if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes, self.dim_out]) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) - bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) - if not self.use_tebd: - assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" - self.register_buffer("bias_atom_e", bias_atom_e) + # init constants if self.numb_fparam > 0: self.register_buffer( @@ -119,7 +112,6 @@ def __init__( self.aparam_avg, self.aparam_inv_std = None, None in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - out_dim = 1 self.old_impl = kwargs.get("old_impl", False) if self.old_impl: @@ -144,7 +136,7 @@ def __init__( networks=[ FittingNet( in_dim, - out_dim, + self.dim_out, self.neuron, self.activation_function, self.resnet_dt, @@ -161,75 +153,6 @@ def __init__( log.info("Set seed to %d in fitting net.", kwargs["seed"]) torch.manual_seed(kwargs["seed"]) - def output_def(self) -> FittingOutputDef: - return FittingOutputDef( - [ - OutputVariableDef( - self.var_name, - [self.dim_out], - reduciable=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - - def __setitem__(self, key, value): - if key in ["bias_atom_e"]: - value = value.view([self.ntypes, self.dim_out]) - self.bias_atom_e = value - elif key in ["fparam_avg"]: - self.fparam_avg = value - elif key in ["fparam_inv_std"]: - self.fparam_inv_std = value - elif key in ["aparam_avg"]: - self.aparam_avg = value - elif key in ["aparam_inv_std"]: - self.aparam_inv_std = value - else: - raise KeyError(key) - - def __getitem__(self, key): - if key in ["bias_atom_e"]: - return self.bias_atom_e - elif key in ["fparam_avg"]: - return self.fparam_avg - elif key in ["fparam_inv_std"]: - return self.fparam_inv_std - elif key in ["aparam_avg"]: - return self.aparam_avg - elif key in ["aparam_inv_std"]: - return self.aparam_inv_std - else: - raise KeyError(key) - - @property - def data_stat_key(self): - """ - Get the keys for the data statistic of the fitting. - Return a list of statistic names needed, such as "bias_atom_e". - """ - return ["bias_atom_e"] - - def compute_output_stats(self, merged): - energy = [item["energy"] for item in merged] - mixed_type = "real_natoms_vec" in merged[0] - if mixed_type: - input_natoms = [item["real_natoms_vec"] for item in merged] - else: - input_natoms = [item["natoms"] for item in merged] - tmp = compute_output_bias(energy, input_natoms) - bias_atom_e = tmp[:, 0] - return {"bias_atom_e": bias_atom_e} - - def init_fitting_stat(self, bias_atom_e=None, **kwargs): - assert all(x is not None for x in [bias_atom_e]) - self.bias_atom_e.copy_( - torch.tensor(bias_atom_e, device=env.DEVICE).view( - [self.ntypes, self.dim_out] - ) - ) - def serialize(self) -> dict: """Serialize the fitting to dict.""" return { @@ -246,7 +169,6 @@ def serialize(self) -> dict: "distinguish_types": self.distinguish_types, "nets": self.filter_layers.serialize(), "@variables": { - "bias_atom_e": to_numpy_array(self.bias_atom_e), "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), "aparam_avg": to_numpy_array(self.aparam_avg), @@ -286,7 +208,7 @@ def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) - def forward( + def _foward_common( self, descriptor: torch.Tensor, atype: torch.Tensor, @@ -296,23 +218,12 @@ def forward( fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): - """Based on embedding net output, alculate total energy. - - Args: - - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. - - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. - - Returns - ------- - - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. - """ + xx = descriptor nf, nloc, nd = xx.shape - # NOTICE in tests/pt/test_model.py - # it happens that the user directly access the data memeber self.bias_atom_e - # and set it to a wrong shape! - self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) - # check input dim + if hasattr(self, "bias_atom_e"): + self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) + if nd != self.dim_descrpt: raise ValueError( "get an input descriptor of dim {nd}," @@ -363,33 +274,167 @@ def forward( outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None if self.use_tebd: - atom_energy = self.filter_layers_old[0](xx) + self.bias_atom_e[ + atom_property = self.filter_layers_old[0](xx) + self.bias_atom_e[ atype ].unsqueeze(-1) - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: for type_i, filter_layer in enumerate(self.filter_layers_old): mask = atype == type_i - atom_energy = filter_layer(xx) - atom_energy = atom_energy + self.bias_atom_e[type_i] - atom_energy = atom_energy * mask.unsqueeze(-1) - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] - return {"energy": outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + atom_property = filter_layer(xx) + atom_property = atom_property + self.bias_atom_e[type_i] if hasattr(self, "bias_atom_e") else atom_property + atom_property = atom_property * mask.unsqueeze(-1) + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} else: if self.use_tebd: - atom_energy = ( + atom_property = ( self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] ) - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: 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) - atom_energy = atom_energy + self.bias_atom_e[type_i] - atom_energy = atom_energy * mask - outs = outs + atom_energy # Shape is [nframes, natoms[0], 1] + atom_property = ll(xx) + atom_property = atom_property + self.bias_atom_e[type_i] if hasattr(self, "bias_atom_e") else atom_property + atom_property = atom_property * mask + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + + +@fitting_check_output +class InvarFitting(GeneralFitting): + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, + **kwargs, + ): + """Construct a fitting net for energy. + + Args: + - ntypes: Element count. + - embedding_width: Embedding width per atom. + - neuron: Number of neurons in each hidden layers of the fitting net. + - bias_atom_e: Average enery per atom for each element. + - resnet_dt: Using time-step in the ResNet construction. + """ + super().__init__() + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, self.dim_out]) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) + if not self.use_tebd: + assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" + self.register_buffer("bias_atom_e", bias_atom_e) + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + + def __setitem__(self, key, value): + if key in ["bias_atom_e"]: + value = value.view([self.ntypes, self.dim_out]) + self.bias_atom_e = value + elif key in ["fparam_avg"]: + self.fparam_avg = value + elif key in ["fparam_inv_std"]: + self.fparam_inv_std = value + elif key in ["aparam_avg"]: + self.aparam_avg = value + elif key in ["aparam_inv_std"]: + self.aparam_inv_std = value + else: + raise KeyError(key) + + def __getitem__(self, key): + if key in ["bias_atom_e"]: + return self.bias_atom_e + elif key in ["fparam_avg"]: + return self.fparam_avg + elif key in ["fparam_inv_std"]: + return self.fparam_inv_std + elif key in ["aparam_avg"]: + return self.aparam_avg + elif key in ["aparam_inv_std"]: + return self.aparam_inv_std + else: + raise KeyError(key) + + @property + def data_stat_key(self): + """ + Get the keys for the data statistic of the fitting. + Return a list of statistic names needed, such as "bias_atom_e". + """ + return ["bias_atom_e"] + + def serialize(self): + data = super().serialize() + data["@variables"]["bias_atom_e"] = to_numpy_array(self.bias_atom_e) + return data + + def compute_output_stats(self, merged): + energy = [item["energy"] for item in merged] + mixed_type = "real_natoms_vec" in merged[0] + if mixed_type: + input_natoms = [item["real_natoms_vec"] for item in merged] + else: + input_natoms = [item["natoms"] for item in merged] + tmp = compute_output_bias(energy, input_natoms) + bias_atom_e = tmp[:, 0] + return {"bias_atom_e": bias_atom_e} + + def init_fitting_stat(self, bias_atom_e=None, **kwargs): + assert all(x is not None for x in [bias_atom_e]) + self.bias_atom_e.copy_( + torch.tensor(bias_atom_e, device=env.DEVICE).view( + [self.ntypes, self.dim_out] + ) + ) + + def forward( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + """Based on embedding net output, alculate total energy. + + Args: + - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. + - natoms: Tell atom count and element count. Its shape is [2+self.ntypes]. + + Returns + ------- + - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. + """ + return super().forward(descriptor, atype, gr, g2, h2, fparam, aparam) + @Fitting.register("ener") From 95852035615d214d0f1683cab1228cfe73719a1e Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Feb 2024 12:58:01 +0000 Subject: [PATCH 06/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 22 ++++++++++++++-------- 1 file changed, 14 insertions(+), 8 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 00271e1493..4d7f25ff7d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -86,7 +86,7 @@ def __init__( self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] - + # init constants if self.numb_fparam > 0: self.register_buffer( @@ -218,7 +218,6 @@ def _foward_common( fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): - xx = descriptor nf, nloc, nd = xx.shape if hasattr(self, "bias_atom_e"): @@ -282,7 +281,11 @@ def _foward_common( for type_i, filter_layer in enumerate(self.filter_layers_old): mask = atype == type_i atom_property = filter_layer(xx) - atom_property = atom_property + self.bias_atom_e[type_i] if hasattr(self, "bias_atom_e") else atom_property + atom_property = ( + atom_property + self.bias_atom_e[type_i] + if hasattr(self, "bias_atom_e") + else atom_property + ) atom_property = atom_property * mask.unsqueeze(-1) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} @@ -297,11 +300,15 @@ def _foward_common( mask = (atype == type_i).unsqueeze(-1) mask = torch.tile(mask, (1, 1, self.dim_out)) atom_property = ll(xx) - atom_property = atom_property + self.bias_atom_e[type_i] if hasattr(self, "bias_atom_e") else atom_property + atom_property = ( + atom_property + self.bias_atom_e[type_i] + if hasattr(self, "bias_atom_e") + else atom_property + ) atom_property = atom_property * mask outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} - + @fitting_check_output class InvarFitting(GeneralFitting): @@ -330,14 +337,14 @@ def __init__( - bias_atom_e: Average enery per atom for each element. - resnet_dt: Using time-step in the ResNet construction. """ - super().__init__() + super().__init__() if bias_atom_e is None: bias_atom_e = np.zeros([self.ntypes, self.dim_out]) bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) if not self.use_tebd: assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" - self.register_buffer("bias_atom_e", bias_atom_e) + self.register_buffer("bias_atom_e", bias_atom_e) def output_def(self) -> FittingOutputDef: return FittingOutputDef( @@ -434,7 +441,6 @@ def forward( - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. """ return super().forward(descriptor, atype, gr, g2, h2, fparam, aparam) - @Fitting.register("ener") From 31d84fc6224fe9b96df6539ae2c417abebffc91f Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 14 Feb 2024 21:00:45 +0800 Subject: [PATCH 07/37] chore: refactor InvarFitting --- deepmd/pt/model/task/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 4d7f25ff7d..7e6c8abe8d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -192,7 +192,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data: dict) -> "InvarFitting": + def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) variables = data.pop("@variables") nets = data.pop("nets") @@ -440,7 +440,7 @@ def forward( ------- - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. """ - return super().forward(descriptor, atype, gr, g2, h2, fparam, aparam) + return super()._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) @Fitting.register("ener") From 5459248d505d97fa042ec803bac3b1deb15cb654 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 14 Feb 2024 21:16:26 +0800 Subject: [PATCH 08/37] chore: refactor InvarFitting --- deepmd/pt/model/task/ener.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 7e6c8abe8d..de4aaf585f 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -54,7 +54,6 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: List[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -337,7 +336,20 @@ def __init__( - bias_atom_e: Average enery per atom for each element. - resnet_dt: Using time-step in the ResNet construction. """ - super().__init__() + super().__init__( + var_name=var_name, + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=dim_out, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + distinguish_types=distinguish_types, + **kwargs + ) if bias_atom_e is None: bias_atom_e = np.zeros([self.ntypes, self.dim_out]) bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) From be5be2534b80f0aa8d7f68c39bb35ecd32610c28 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 14 Feb 2024 13:17:00 +0000 Subject: [PATCH 09/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index de4aaf585f..38636376db 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -348,7 +348,7 @@ def __init__( activation_function=activation_function, precision=precision, distinguish_types=distinguish_types, - **kwargs + **kwargs, ) if bias_atom_e is None: bias_atom_e = np.zeros([self.ntypes, self.dim_out]) From ad810df7fd6864215c51cab543348678a1a5ad6a Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Wed, 14 Feb 2024 21:49:02 +0800 Subject: [PATCH 10/37] chore: refactor InvarFitting --- deepmd/pt/model/task/ener.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 38636376db..f85820d57d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -207,7 +207,7 @@ def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) - def _foward_common( + def _forward_common( self, descriptor: torch.Tensor, atype: torch.Tensor, @@ -452,7 +452,7 @@ def forward( ------- - `torch.Tensor`: Total energy with shape [nframes, natoms[0]]. """ - return super()._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) + return self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam) @Fitting.register("ener") From 6ade44f1ff72804505f8ef56483dd029779eea69 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 11:07:38 +0800 Subject: [PATCH 11/37] chore: refactor InvarFitting --- deepmd/pt/model/task/ener.py | 285 +--------------------------- deepmd/pt/model/task/fitting.py | 318 ++++++++++++++++++++++++++++++++ 2 files changed, 320 insertions(+), 283 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index b84e40f6af..00884390bf 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -1,5 +1,4 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import copy import logging from typing import ( List, @@ -15,29 +14,25 @@ OutputVariableDef, fitting_check_output, ) -from deepmd.pt.model.network.mlp import ( - FittingNet, - NetworkCollection, -) + from deepmd.pt.model.network.network import ( ResidualDeep, ) from deepmd.pt.model.task.fitting import ( Fitting, + GeneralFitting, ) from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.env import ( DEFAULT_PRECISION, - PRECISION_DICT, ) from deepmd.pt.utils.stat import ( compute_output_bias, ) from deepmd.pt.utils.utils import ( to_numpy_array, - to_torch_tensor, ) dtype = env.GLOBAL_PT_FLOAT_PRECISION @@ -46,269 +41,6 @@ log = logging.getLogger(__name__) -class GeneralFitting(Fitting): - def __init__( - self, - var_name: str, - ntypes: int, - dim_descrpt: int, - dim_out: int, - neuron: List[int] = [128, 128, 128], - resnet_dt: bool = True, - numb_fparam: int = 0, - numb_aparam: int = 0, - activation_function: str = "tanh", - precision: str = DEFAULT_PRECISION, - distinguish_types: bool = False, - **kwargs, - ): - """Construct a fitting net for energy. - - Args: - - ntypes: Element count. - - embedding_width: Embedding width per atom. - - neuron: Number of neurons in each hidden layers of the fitting net. - - bias_atom_e: Average enery per atom for each element. - - resnet_dt: Using time-step in the ResNet construction. - """ - super().__init__() - self.var_name = var_name - self.ntypes = ntypes - self.dim_descrpt = dim_descrpt - self.dim_out = dim_out - self.neuron = neuron - self.distinguish_types = distinguish_types - self.use_tebd = not self.distinguish_types - self.resnet_dt = resnet_dt - self.numb_fparam = numb_fparam - self.numb_aparam = numb_aparam - self.activation_function = activation_function - self.precision = precision - self.prec = PRECISION_DICT[self.precision] - - # init constants - if self.numb_fparam > 0: - self.register_buffer( - "fparam_avg", - torch.zeros(self.numb_fparam, dtype=self.prec, device=device), - ) - self.register_buffer( - "fparam_inv_std", - torch.ones(self.numb_fparam, dtype=self.prec, device=device), - ) - else: - self.fparam_avg, self.fparam_inv_std = None, None - if self.numb_aparam > 0: - self.register_buffer( - "aparam_avg", - torch.zeros(self.numb_aparam, dtype=self.prec, device=device), - ) - self.register_buffer( - "aparam_inv_std", - torch.ones(self.numb_aparam, dtype=self.prec, device=device), - ) - else: - self.aparam_avg, self.aparam_inv_std = None, None - - in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - - self.old_impl = kwargs.get("old_impl", False) - if self.old_impl: - filter_layers = [] - for type_i in range(self.ntypes): - bias_type = 0.0 - one = ResidualDeep( - type_i, - self.dim_descrpt, - self.neuron, - bias_type, - resnet_dt=self.resnet_dt, - ) - filter_layers.append(one) - self.filter_layers_old = torch.nn.ModuleList(filter_layers) - self.filter_layers = None - else: - self.filter_layers = NetworkCollection( - 1 if self.distinguish_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - self.dim_out, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if self.distinguish_types else 1) - ], - ) - self.filter_layers_old = None - - # very bad design... - if "seed" in kwargs: - log.info("Set seed to %d in fitting net.", kwargs["seed"]) - torch.manual_seed(kwargs["seed"]) - - def serialize(self) -> dict: - """Serialize the fitting to dict.""" - return { - "var_name": self.var_name, - "ntypes": self.ntypes, - "dim_descrpt": self.dim_descrpt, - "dim_out": self.dim_out, - "neuron": self.neuron, - "resnet_dt": self.resnet_dt, - "numb_fparam": self.numb_fparam, - "numb_aparam": self.numb_aparam, - "activation_function": self.activation_function, - "precision": self.precision, - "distinguish_types": self.distinguish_types, - "nets": self.filter_layers.serialize(), - "@variables": { - "fparam_avg": to_numpy_array(self.fparam_avg), - "fparam_inv_std": to_numpy_array(self.fparam_inv_std), - "aparam_avg": to_numpy_array(self.aparam_avg), - "aparam_inv_std": to_numpy_array(self.aparam_inv_std), - }, - # "rcond": self.rcond , - # "tot_ener_zero": self.tot_ener_zero , - # "trainable": self.trainable , - # "atom_ener": self.atom_ener , - # "layer_name": self.layer_name , - # "use_aparam_as_mask": self.use_aparam_as_mask , - # "spin": self.spin , - ## NOTICE: not supported by far - "rcond": None, - "tot_ener_zero": False, - "trainable": True, - "atom_ener": None, - "layer_name": None, - "use_aparam_as_mask": False, - "spin": None, - } - - @classmethod - def deserialize(cls, data: dict) -> "GeneralFitting": - data = copy.deepcopy(data) - variables = data.pop("@variables") - nets = data.pop("nets") - obj = cls(**data) - for kk in variables.keys(): - obj[kk] = to_torch_tensor(variables[kk]) - obj.filter_layers = NetworkCollection.deserialize(nets) - return obj - - def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: - return torch.tile(xx.view([1, self.numb_fparam]), [nb, 1]) - - def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: - return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) - - def _forward_common( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): - xx = descriptor - nf, nloc, nd = xx.shape - if hasattr(self, "bias_atom_e"): - self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) - - if nd != self.dim_descrpt: - raise ValueError( - "get an input descriptor of dim {nd}," - "which is not consistent with {self.dim_descrpt}." - ) - # check fparam dim, concate to input descriptor - if self.numb_fparam > 0: - assert fparam is not None, "fparam should not be None" - assert self.fparam_avg is not None - assert self.fparam_inv_std is not None - if fparam.shape[-1] != self.numb_fparam: - raise ValueError( - "get an input fparam of dim {fparam.shape[-1]}, ", - "which is not consistent with {self.numb_fparam}.", - ) - fparam = fparam.view([nf, self.numb_fparam]) - nb, _ = fparam.shape - t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) - t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) - fparam = (fparam - t_fparam_avg) * t_fparam_inv_std - fparam = torch.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) - xx = torch.cat( - [xx, fparam], - dim=-1, - ) - # check aparam dim, concate to input descriptor - if self.numb_aparam > 0: - assert aparam is not None, "aparam should not be None" - assert self.aparam_avg is not None - assert self.aparam_inv_std is not None - if aparam.shape[-1] != self.numb_aparam: - raise ValueError( - "get an input aparam of dim {aparam.shape[-1]}, ", - "which is not consistent with {self.numb_aparam}.", - ) - aparam = aparam.view([nf, nloc, self.numb_aparam]) - nb, nloc, _ = aparam.shape - t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) - t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) - aparam = (aparam - t_aparam_avg) * t_aparam_inv_std - xx = torch.cat( - [xx, aparam], - dim=-1, - ) - - outs = torch.zeros(nf, nloc, self.dim_out) # jit assertion - if self.old_impl: - outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion - assert self.filter_layers_old is not None - if self.use_tebd: - atom_property = self.filter_layers_old[0](xx) + self.bias_atom_e[ - atype - ].unsqueeze(-1) - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] - else: - for type_i, filter_layer in enumerate(self.filter_layers_old): - mask = atype == type_i - atom_property = filter_layer(xx) - atom_property = ( - atom_property + self.bias_atom_e[type_i] - if hasattr(self, "bias_atom_e") - else atom_property - ) - atom_property = atom_property * mask.unsqueeze(-1) - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] - return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} - else: - if self.use_tebd: - atom_property = ( - self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] - ) - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] - else: - 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_property = ll(xx) - atom_property = ( - atom_property + self.bias_atom_e[type_i] - if hasattr(self, "bias_atom_e") - else atom_property - ) - atom_property = atom_property * mask - outs = outs + atom_property # Shape is [nframes, natoms[0], 1] - return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} - - @fitting_check_output class InvarFitting(GeneralFitting): def __init__( @@ -358,19 +90,6 @@ def __init__( assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" self.register_buffer("bias_atom_e", bias_atom_e) - def output_def(self) -> FittingOutputDef: - return FittingOutputDef( - [ - OutputVariableDef( - self.var_name, - [self.dim_out], - reduciable=True, - r_differentiable=True, - c_differentiable=True, - ), - ] - ) - def __setitem__(self, key, value): if key in ["bias_atom_e"]: value = value.view([self.ntypes, self.dim_out]) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 360f545975..068f9be219 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy import logging from typing import ( Callable, @@ -13,6 +14,21 @@ from deepmd.pt.model.task.base_fitting import ( BaseFitting, ) +from deepmd.pt.model.network.network import ( + ResidualDeep, +) +from deepmd.pt.model.network.mlp import ( + FittingNet, + NetworkCollection, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, +) from deepmd.pt.utils.dataloader import ( DpLoaderSet, ) @@ -22,9 +38,19 @@ from deepmd.pt.utils.plugin import ( Plugin, ) +from deepmd.pt.utils import ( + env, +) from deepmd.pt.utils.stat import ( make_stat_input, ) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE log = logging.getLogger(__name__) @@ -316,3 +342,295 @@ def change_energy_bias( ) ) return None + +class GeneralFitting(Fitting): + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, + **kwargs, + ): + """Construct a fitting net for energy. + + Args: + - ntypes: Element count. + - embedding_width: Embedding width per atom. + - neuron: Number of neurons in each hidden layers of the fitting net. + - bias_atom_e: Average enery per atom for each element. + - resnet_dt: Using time-step in the ResNet construction. + """ + super().__init__() + self.var_name = var_name + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.dim_out = dim_out + self.neuron = neuron + self.distinguish_types = distinguish_types + self.use_tebd = not self.distinguish_types + self.resnet_dt = resnet_dt + self.numb_fparam = numb_fparam + self.numb_aparam = numb_aparam + self.activation_function = activation_function + self.precision = precision + self.prec = PRECISION_DICT[self.precision] + + # init constants + if self.numb_fparam > 0: + self.register_buffer( + "fparam_avg", + torch.zeros(self.numb_fparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "fparam_inv_std", + torch.ones(self.numb_fparam, dtype=self.prec, device=device), + ) + else: + self.fparam_avg, self.fparam_inv_std = None, None + if self.numb_aparam > 0: + self.register_buffer( + "aparam_avg", + torch.zeros(self.numb_aparam, dtype=self.prec, device=device), + ) + self.register_buffer( + "aparam_inv_std", + torch.ones(self.numb_aparam, dtype=self.prec, device=device), + ) + else: + self.aparam_avg, self.aparam_inv_std = None, None + + in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam + + self.old_impl = kwargs.get("old_impl", False) + if self.old_impl: + filter_layers = [] + for type_i in range(self.ntypes): + bias_type = 0.0 + one = ResidualDeep( + type_i, + self.dim_descrpt, + self.neuron, + bias_type, + resnet_dt=self.resnet_dt, + ) + filter_layers.append(one) + self.filter_layers_old = torch.nn.ModuleList(filter_layers) + self.filter_layers = None + else: + self.filter_layers = NetworkCollection( + 1 if self.distinguish_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim, + self.dim_out, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if self.distinguish_types else 1) + ], + ) + self.filter_layers_old = None + + # very bad design... + if "seed" in kwargs: + log.info("Set seed to %d in fitting net.", kwargs["seed"]) + torch.manual_seed(kwargs["seed"]) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "dim_out": self.dim_out, + "neuron": self.neuron, + "resnet_dt": self.resnet_dt, + "numb_fparam": self.numb_fparam, + "numb_aparam": self.numb_aparam, + "activation_function": self.activation_function, + "precision": self.precision, + "distinguish_types": self.distinguish_types, + "nets": self.filter_layers.serialize(), + "@variables": { + "fparam_avg": to_numpy_array(self.fparam_avg), + "fparam_inv_std": to_numpy_array(self.fparam_inv_std), + "aparam_avg": to_numpy_array(self.aparam_avg), + "aparam_inv_std": to_numpy_array(self.aparam_inv_std), + }, + # "rcond": self.rcond , + # "tot_ener_zero": self.tot_ener_zero , + # "trainable": self.trainable , + # "atom_ener": self.atom_ener , + # "layer_name": self.layer_name , + # "use_aparam_as_mask": self.use_aparam_as_mask , + # "spin": self.spin , + ## NOTICE: not supported by far + "rcond": None, + "tot_ener_zero": False, + "trainable": True, + "atom_ener": None, + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + variables = data.pop("@variables") + nets = data.pop("nets") + obj = cls(**data) + for kk in variables.keys(): + obj[kk] = to_torch_tensor(variables[kk]) + obj.filter_layers = NetworkCollection.deserialize(nets) + return obj + + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this atomic model.""" + return self.numb_fparam + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this atomic model.""" + return self.numb_aparam + + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ + return [] + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=True, + c_differentiable=True, + ), + ] + ) + def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: + return torch.tile(xx.view([1, self.numb_fparam]), [nb, 1]) + + def _extend_a_avg_std(self, xx: torch.Tensor, nb: int, nloc: int) -> torch.Tensor: + return torch.tile(xx.view([1, 1, self.numb_aparam]), [nb, nloc, 1]) + + def _forward_common( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + xx = descriptor + nf, nloc, nd = xx.shape + if hasattr(self, "bias_atom_e"): + self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) + + if nd != self.dim_descrpt: + raise ValueError( + "get an input descriptor of dim {nd}," + "which is not consistent with {self.dim_descrpt}." + ) + # check fparam dim, concate to input descriptor + if self.numb_fparam > 0: + assert fparam is not None, "fparam should not be None" + assert self.fparam_avg is not None + assert self.fparam_inv_std is not None + if fparam.shape[-1] != self.numb_fparam: + raise ValueError( + "get an input fparam of dim {fparam.shape[-1]}, ", + "which is not consistent with {self.numb_fparam}.", + ) + fparam = fparam.view([nf, self.numb_fparam]) + nb, _ = fparam.shape + t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) + t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) + fparam = (fparam - t_fparam_avg) * t_fparam_inv_std + fparam = torch.tile(fparam.reshape([nf, 1, -1]), [1, nloc, 1]) + xx = torch.cat( + [xx, fparam], + dim=-1, + ) + # check aparam dim, concate to input descriptor + if self.numb_aparam > 0: + assert aparam is not None, "aparam should not be None" + assert self.aparam_avg is not None + assert self.aparam_inv_std is not None + if aparam.shape[-1] != self.numb_aparam: + raise ValueError( + "get an input aparam of dim {aparam.shape[-1]}, ", + "which is not consistent with {self.numb_aparam}.", + ) + aparam = aparam.view([nf, nloc, self.numb_aparam]) + nb, nloc, _ = aparam.shape + t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) + t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) + aparam = (aparam - t_aparam_avg) * t_aparam_inv_std + xx = torch.cat( + [xx, aparam], + dim=-1, + ) + + outs = torch.zeros(nf, nloc, self.dim_out) # jit assertion + if self.old_impl: + outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion + assert self.filter_layers_old is not None + if self.use_tebd: + atom_property = self.filter_layers_old[0](xx) + self.bias_atom_e[ + atype + ].unsqueeze(-1) + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + else: + for type_i, filter_layer in enumerate(self.filter_layers_old): + mask = atype == type_i + atom_property = filter_layer(xx) + atom_property = ( + atom_property + self.bias_atom_e[type_i] + if hasattr(self, "bias_atom_e") + else atom_property + ) + atom_property = atom_property * mask.unsqueeze(-1) + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + else: + if self.use_tebd: + atom_property = ( + self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] + ) + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + else: + 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_property = ll(xx) + atom_property = ( + atom_property + self.bias_atom_e[type_i] + if hasattr(self, "bias_atom_e") + else atom_property + ) + atom_property = atom_property * mask + outs = outs + atom_property # Shape is [nframes, natoms[0], 1] + return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + From ca957b54b7921b071aef82a6dcb2af19153b954c Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 03:08:16 +0000 Subject: [PATCH 12/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 1 - deepmd/pt/model/task/fitting.py | 65 ++++++++++++++++----------------- 2 files changed, 32 insertions(+), 34 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 00884390bf..9a93c126b9 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -14,7 +14,6 @@ OutputVariableDef, fitting_check_output, ) - from deepmd.pt.model.network.network import ( ResidualDeep, ) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 068f9be219..cbdf187172 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -11,42 +11,40 @@ import numpy as np import torch -from deepmd.pt.model.task.base_fitting import ( - BaseFitting, -) -from deepmd.pt.model.network.network import ( - ResidualDeep, +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, ) from deepmd.pt.model.network.mlp import ( FittingNet, NetworkCollection, ) -from deepmd.pt.utils.utils import ( - to_numpy_array, - to_torch_tensor, +from deepmd.pt.model.network.network import ( + ResidualDeep, ) -from deepmd.dpmodel import ( - FittingOutputDef, - OutputVariableDef, +from deepmd.pt.model.task.base_fitting import ( + BaseFitting, +) +from deepmd.pt.utils import ( + env, ) from deepmd.pt.utils.dataloader import ( DpLoaderSet, ) from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, DEVICE, + PRECISION_DICT, ) from deepmd.pt.utils.plugin import ( Plugin, ) -from deepmd.pt.utils import ( - env, -) from deepmd.pt.utils.stat import ( make_stat_input, ) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, ) dtype = env.GLOBAL_PT_FLOAT_PRECISION @@ -343,6 +341,7 @@ def change_energy_bias( ) return None + class GeneralFitting(Fitting): def __init__( self, @@ -498,23 +497,23 @@ def deserialize(cls, data: dict) -> "GeneralFitting": obj.filter_layers = NetworkCollection.deserialize(nets) return obj - def get_dim_fparam(self) -> int: - """Get the number (dimension) of frame parameters of this atomic model.""" + def get_dim_fparam(self) -> int: + """Get the number (dimension) of frame parameters of this atomic model.""" return self.numb_fparam - - def get_dim_aparam(self) -> int: - """Get the number (dimension) of atomic parameters of this atomic model.""" + + def get_dim_aparam(self) -> int: + """Get the number (dimension) of atomic parameters of this atomic model.""" return self.numb_aparam - - def get_sel_type(self) -> List[int]: - """Get the selected atom types of this model. - - Only atoms with selected atom types have atomic contribution - to the result of the model. - If returning an empty list, all atom types are selected. - """ - return [] - + + def get_sel_type(self) -> List[int]: + """Get the selected atom types of this model. + + Only atoms with selected atom types have atomic contribution + to the result of the model. + If returning an empty list, all atom types are selected. + """ + return [] + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ @@ -527,6 +526,7 @@ def output_def(self) -> FittingOutputDef: ), ] ) + def _extend_f_avg_std(self, xx: torch.Tensor, nb: int) -> torch.Tensor: return torch.tile(xx.view([1, self.numb_fparam]), [nb, 1]) @@ -633,4 +633,3 @@ def _forward_common( atom_property = atom_property * mask outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} - From 404f31ff82628a61064ed4c57b422c70fb538eaf Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 15:19:50 +0800 Subject: [PATCH 13/37] fix: internal dipole fit output shape --- deepmd/pt/model/task/dipole.py | 174 +++++-------------- deepmd/pt/model/task/fitting.py | 2 +- source/tests/pt/model/test_dipole_fitting.py | 82 +++++++++ 3 files changed, 129 insertions(+), 129 deletions(-) create mode 100644 source/tests/pt/model/test_dipole_fitting.py diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 50b4ebff3e..35eecabfe0 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -12,14 +12,13 @@ NetworkCollection, ) from deepmd.pt.model.task.fitting import ( - Fitting, + GeneralFitting, ) from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.env import ( DEFAULT_PRECISION, - PRECISION_DICT, ) dtype = env.GLOBAL_PT_FLOAT_PRECISION @@ -28,7 +27,7 @@ log = logging.getLogger(__name__) -class DipoleFittingNet(Fitting): +class DipoleFittingNet(GeneralFitting): def __init__( self, var_name: str, @@ -36,7 +35,6 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: List[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -54,69 +52,21 @@ def __init__( - bias_atom_e: Average enery per atom for each element. - resnet_dt: Using time-step in the ResNet construction. """ - super().__init__() - self.var_name = var_name - self.ntypes = ntypes - self.dim_descrpt = dim_descrpt - self.dim_out = dim_out - self.neuron = neuron - self.distinguish_types = distinguish_types - self.use_tebd = not self.distinguish_types - self.resnet_dt = resnet_dt - self.numb_fparam = numb_fparam - self.numb_aparam = numb_aparam - self.activation_function = activation_function - self.precision = precision - self.prec = PRECISION_DICT[self.precision] - - # init constants - if self.numb_fparam > 0: - self.register_buffer( - "fparam_avg", - torch.zeros(self.numb_fparam, dtype=self.prec, device=device), - ) - self.register_buffer( - "fparam_inv_std", - torch.ones(self.numb_fparam, dtype=self.prec, device=device), - ) - else: - self.fparam_avg, self.fparam_inv_std = None, None - if self.numb_aparam > 0: - self.register_buffer( - "aparam_avg", - torch.zeros(self.numb_aparam, dtype=self.prec, device=device), - ) - self.register_buffer( - "aparam_inv_std", - torch.ones(self.numb_aparam, dtype=self.prec, device=device), - ) - else: - self.aparam_avg, self.aparam_inv_std = None, None - - in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - out_dim = 3 - - self.filter_layers = NetworkCollection( - 1 if self.distinguish_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - out_dim, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if self.distinguish_types else 1) - ], + super().__init__( + var_name=var_name, + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=dim_out, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + distinguish_types=distinguish_types, + **kwargs, ) - - if "seed" in kwargs: - log.info("Set seed to %d in fitting net.", kwargs["seed"]) - torch.manual_seed(kwargs["seed"]) + self.old_impl = False # this only supports the new implementation. def forward( self, @@ -140,66 +90,34 @@ def forward( ------- - vec_out: output vector. Its shape is [nframes, nloc, 3]. """ - xx = descriptor - nframes, nloc, nd = xx.shape - # check input dim - if nd != self.dim_descrpt: - raise ValueError( - "get an input descriptor of dim {nd}," - "which is not consistent with {self.dim_descrpt}." - ) - # check fparam dim, concate to input descriptor - if self.numb_fparam > 0: - assert fparam is not None, "fparam should not be None" - assert self.fparam_avg is not None - assert self.fparam_inv_std is not None - if fparam.shape[-1] != self.numb_fparam: - raise ValueError( - "get an input fparam of dim {fparam.shape[-1]}, ", - "which is not consistent with {self.numb_fparam}.", - ) - fparam = fparam.view([nframes, self.numb_fparam]) - nb, _ = fparam.shape - t_fparam_avg = self._extend_f_avg_std(self.fparam_avg, nb) - t_fparam_inv_std = self._extend_f_avg_std(self.fparam_inv_std, nb) - fparam = (fparam - t_fparam_avg) * t_fparam_inv_std - fparam = torch.tile(fparam.reshape([nframes, 1, -1]), [1, nloc, 1]) - xx = torch.cat( - [xx, fparam], - dim=-1, - ) - # check aparam dim, concate to input descriptor - if self.numb_aparam > 0: - assert aparam is not None, "aparam should not be None" - assert self.aparam_avg is not None - assert self.aparam_inv_std is not None - if aparam.shape[-1] != self.numb_aparam: - raise ValueError( - "get an input aparam of dim {aparam.shape[-1]}, ", - "which is not consistent with {self.numb_aparam}.", + in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam + + nframes, nloc, _ = descriptor.shape + gr = gr.view(nframes, nloc, -1, 3) + out_dim = gr.shape[2] # m1 + self.filter_layers = NetworkCollection( + 1 if self.distinguish_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim, + out_dim, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, ) - aparam = aparam.view([nframes, nloc, self.numb_aparam]) - nb, nloc, _ = aparam.shape - t_aparam_avg = self._extend_a_avg_std(self.aparam_avg, nb, nloc) - t_aparam_inv_std = self._extend_a_avg_std(self.aparam_inv_std, nb, nloc) - aparam = (aparam - t_aparam_avg) * t_aparam_inv_std - xx = torch.cat( - [xx, aparam], - dim=-1, - ) - - outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion - if self.use_tebd: - atom_dipole = self.filter_layers.networks[0](xx) - outs = outs + atom_dipole # Shape is [nframes, nloc, 3] - else: - 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_dipole = ll(xx) - atom_dipole = atom_dipole * mask - outs = outs + atom_dipole # Shape is [nframes, nloc, 3] - outs = ( - torch.bmm(outs, gr).squeeze(-2).view(nframes, nloc, 3) - ) # Shape is [nframes, nloc, 3] - return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} + for ii in range(self.ntypes if self.distinguish_types else 1) + ], + ) + # (nframes, nloc, m1) + out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[self.var_name] + # (nframes * nloc, 1, m1) + out = out.view(-1, 1, out_dim) + # (nframes, nloc, 3) + out = ( + torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3) + ) + return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} \ No newline at end of file diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index cbdf187172..e607cd99b9 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -618,7 +618,7 @@ def _forward_common( if self.use_tebd: atom_property = ( self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] - ) + ) if hasattr(self, "bias_atom_e") else self.filter_layers.networks[0](xx) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: for type_i, ll in enumerate(self.filter_layers.networks): diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py new file mode 100644 index 0000000000..6a01dcd0d4 --- /dev/null +++ b/source/tests/pt/model/test_dipole_fitting.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import itertools +import unittest + +import numpy as np +import torch + +from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.pt.model.descriptor.se_a import ( + DescrptSeA, +) +from deepmd.pt.model.task.dipole import ( + DipoleFittingNet, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, +) + +from .test_env_mat import ( + TestCaseSingleFrameWithNlist, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION + + +class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): + def setUp(self): + TestCaseSingleFrameWithNlist.setUp(self) + + def test_new_old( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + dd = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + rd0, gr, _, _, _ = dd( + 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), + ) + atype = torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE) + + od = 3 + for distinguish_types in itertools.product( + [True, False], + ): + ft0 = DipoleFittingNet( + "dipole", + self.nt, + dd.dim_out, + od, + distinguish_types=distinguish_types, + ).to(env.DEVICE) + ft1 = DipoleFittingNet( + "dipole", + self.nt, + dd.dim_out, + od, + distinguish_types=distinguish_types, + ).to(env.DEVICE) + dd0 = ft0.state_dict() + 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("deep_layers", "layers") + new_kk = new_kk.replace("final_layer", "layers.3") + dd1[kk] = dd0[new_kk] + # if kk.split(".")[-1] in ["idt", "bias"]: + # dd1[kk] = dd1[kk].unsqueeze(0) + ft1.load_state_dict(dd1) + ret0 = ft0(rd0, atype, gr) + ret1 = ft1(rd0, atype, gr) + np.testing.assert_allclose( + to_numpy_array(ret0["dipole"]), + to_numpy_array(ret1["dipole"]), + ) + + \ No newline at end of file From d1fedf81d67ed0d39836a1d27f25192c987f6eda Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 07:20:24 +0000 Subject: [PATCH 14/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/dipole.py | 16 ++++++++-------- deepmd/pt/model/task/fitting.py | 6 ++++-- source/tests/pt/model/test_dipole_fitting.py | 5 +---- 3 files changed, 13 insertions(+), 14 deletions(-) diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 35eecabfe0..2b5f6f8a23 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -66,7 +66,7 @@ def __init__( distinguish_types=distinguish_types, **kwargs, ) - self.old_impl = False # this only supports the new implementation. + self.old_impl = False # this only supports the new implementation. def forward( self, @@ -91,10 +91,10 @@ def forward( - vec_out: output vector. Its shape is [nframes, nloc, 3]. """ in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - + nframes, nloc, _ = descriptor.shape gr = gr.view(nframes, nloc, -1, 3) - out_dim = gr.shape[2] # m1 + out_dim = gr.shape[2] # m1 self.filter_layers = NetworkCollection( 1 if self.distinguish_types else 0, self.ntypes, @@ -113,11 +113,11 @@ def forward( ], ) # (nframes, nloc, m1) - out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[self.var_name] + out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ + self.var_name + ] # (nframes * nloc, 1, m1) out = out.view(-1, 1, out_dim) # (nframes, nloc, 3) - out = ( - torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3) - ) - return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} \ No newline at end of file + out = torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3) + return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index e607cd99b9..54b50cbeb1 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -617,8 +617,10 @@ def _forward_common( else: if self.use_tebd: atom_property = ( - self.filter_layers.networks[0](xx) + self.bias_atom_e[atype] - ) if hasattr(self, "bias_atom_e") else self.filter_layers.networks[0](xx) + (self.filter_layers.networks[0](xx) + self.bias_atom_e[atype]) + if hasattr(self, "bias_atom_e") + else self.filter_layers.networks[0](xx) + ) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: for type_i, ll in enumerate(self.filter_layers.networks): diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py index 6a01dcd0d4..a78aa58af6 100644 --- a/source/tests/pt/model/test_dipole_fitting.py +++ b/source/tests/pt/model/test_dipole_fitting.py @@ -5,7 +5,6 @@ import numpy as np import torch -from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) @@ -70,7 +69,7 @@ def test_new_old( new_kk = new_kk.replace("final_layer", "layers.3") dd1[kk] = dd0[new_kk] # if kk.split(".")[-1] in ["idt", "bias"]: - # dd1[kk] = dd1[kk].unsqueeze(0) + # dd1[kk] = dd1[kk].unsqueeze(0) ft1.load_state_dict(dd1) ret0 = ft0(rd0, atype, gr) ret1 = ft1(rd0, atype, gr) @@ -78,5 +77,3 @@ def test_new_old( to_numpy_array(ret0["dipole"]), to_numpy_array(ret1["dipole"]), ) - - \ No newline at end of file From 514271b28e83bc099d1da8dfb584ca70ff46c0c7 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 16:07:25 +0800 Subject: [PATCH 15/37] fix: mask shape --- deepmd/pt/model/task/dipole.py | 1 + deepmd/pt/model/task/fitting.py | 3 +++ 2 files changed, 4 insertions(+) diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 2b5f6f8a23..3ff061af0a 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -118,6 +118,7 @@ def forward( ] # (nframes * nloc, 1, m1) out = out.view(-1, 1, out_dim) + gr = gr.view(nframes * nloc, -1, 3) # (nframes, nloc, 3) out = torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3) return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 54b50cbeb1..220e827f73 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -545,6 +545,9 @@ def _forward_common( ): xx = descriptor nf, nloc, nd = xx.shape + if gr is not None: + gr = gr.view(nf, nloc, -1, 3) + self.dim_out = gr.shape[2] if hasattr(self, "bias_atom_e"): self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) From 06fb5bf021add677a1d7073c59c66cf4d539e5ba Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 18:02:09 +0800 Subject: [PATCH 16/37] chore: restore dipole, split PR --- deepmd/pt/model/task/dipole.py | 117 +++++-------------- deepmd/pt/model/task/fitting.py | 4 +- source/tests/pt/model/test_dipole_fitting.py | 79 ------------- 3 files changed, 31 insertions(+), 169 deletions(-) delete mode 100644 source/tests/pt/model/test_dipole_fitting.py diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index 3ff061af0a..d46ad11717 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -1,47 +1,21 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import logging -from typing import ( - List, - Optional, -) import torch -from deepmd.pt.model.network.mlp import ( - FittingNet, - NetworkCollection, +from deepmd.pt.model.network.network import ( + ResidualDeep, ) from deepmd.pt.model.task.fitting import ( - GeneralFitting, -) -from deepmd.pt.utils import ( - env, -) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, + Fitting, ) -dtype = env.GLOBAL_PT_FLOAT_PRECISION -device = env.DEVICE - log = logging.getLogger(__name__) -class DipoleFittingNet(GeneralFitting): +class DipoleFittingNet(Fitting): def __init__( - self, - var_name: str, - ntypes: int, - dim_descrpt: int, - dim_out: int, - neuron: List[int] = [128, 128, 128], - resnet_dt: bool = True, - numb_fparam: int = 0, - numb_aparam: int = 0, - activation_function: str = "tanh", - precision: str = DEFAULT_PRECISION, - distinguish_types: bool = False, - **kwargs, + self, ntypes, embedding_width, neuron, out_dim, resnet_dt=True, **kwargs ): """Construct a fitting net for dipole. @@ -52,32 +26,23 @@ def __init__( - bias_atom_e: Average enery per atom for each element. - resnet_dt: Using time-step in the ResNet construction. """ - super().__init__( - var_name=var_name, - ntypes=ntypes, - dim_descrpt=dim_descrpt, - dim_out=dim_out, - neuron=neuron, - resnet_dt=resnet_dt, - numb_fparam=numb_fparam, - numb_aparam=numb_aparam, - activation_function=activation_function, - precision=precision, - distinguish_types=distinguish_types, - **kwargs, + super().__init__() + self.ntypes = ntypes + self.embedding_width = embedding_width + self.out_dim = out_dim + + filter_layers = [] + one = ResidualDeep( + 0, embedding_width, neuron, 0.0, out_dim=self.out_dim, resnet_dt=resnet_dt ) - self.old_impl = False # this only supports the new implementation. + filter_layers.append(one) + self.filter_layers = torch.nn.ModuleList(filter_layers) - def forward( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): + if "seed" in kwargs: + log.info("Set seed to %d in fitting net.", kwargs["seed"]) + torch.manual_seed(kwargs["seed"]) + + def forward(self, inputs, atype, atype_tebd, rot_mat): """Based on embedding net output, alculate total energy. Args: @@ -90,35 +55,13 @@ def forward( ------- - vec_out: output vector. Its shape is [nframes, nloc, 3]. """ - in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam - - nframes, nloc, _ = descriptor.shape - gr = gr.view(nframes, nloc, -1, 3) - out_dim = gr.shape[2] # m1 - self.filter_layers = NetworkCollection( - 1 if self.distinguish_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim, - out_dim, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if self.distinguish_types else 1) - ], - ) - # (nframes, nloc, m1) - out = self._forward_common(descriptor, atype, gr, g2, h2, fparam, aparam)[ - self.var_name - ] - # (nframes * nloc, 1, m1) - out = out.view(-1, 1, out_dim) - gr = gr.view(nframes * nloc, -1, 3) - # (nframes, nloc, 3) - out = torch.bmm(out, gr).squeeze(-2).view(nframes, nloc, 3) - return {self.var_name: out.to(env.GLOBAL_PT_FLOAT_PRECISION)} + nframes, nloc, _ = inputs.size() + if atype_tebd is not None: + inputs = torch.concat([inputs, atype_tebd], dim=-1) + vec_out = self.filter_layers[0](inputs) # Shape is [nframes, nloc, m1] + assert list(vec_out.size()) == [nframes, nloc, self.out_dim] + vec_out = vec_out.view(-1, 1, self.out_dim) + vec_out = ( + torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3) + ) # Shape is [nframes, nloc, 3] + return vec_out \ No newline at end of file diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 220e827f73..e8fbadbfdd 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -545,9 +545,7 @@ def _forward_common( ): xx = descriptor nf, nloc, nd = xx.shape - if gr is not None: - gr = gr.view(nf, nloc, -1, 3) - self.dim_out = gr.shape[2] + if hasattr(self, "bias_atom_e"): self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) diff --git a/source/tests/pt/model/test_dipole_fitting.py b/source/tests/pt/model/test_dipole_fitting.py deleted file mode 100644 index a78aa58af6..0000000000 --- a/source/tests/pt/model/test_dipole_fitting.py +++ /dev/null @@ -1,79 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import itertools -import unittest - -import numpy as np -import torch - -from deepmd.pt.model.descriptor.se_a import ( - DescrptSeA, -) -from deepmd.pt.model.task.dipole import ( - DipoleFittingNet, -) -from deepmd.pt.utils import ( - env, -) -from deepmd.pt.utils.utils import ( - to_numpy_array, -) - -from .test_env_mat import ( - TestCaseSingleFrameWithNlist, -) - -dtype = env.GLOBAL_PT_FLOAT_PRECISION - - -class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): - def setUp(self): - TestCaseSingleFrameWithNlist.setUp(self) - - def test_new_old( - self, - ): - rng = np.random.default_rng() - nf, nloc, nnei = self.nlist.shape - dd = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) - rd0, gr, _, _, _ = dd( - 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), - ) - atype = torch.tensor(self.atype_ext[:, :nloc], dtype=int, device=env.DEVICE) - - od = 3 - for distinguish_types in itertools.product( - [True, False], - ): - ft0 = DipoleFittingNet( - "dipole", - self.nt, - dd.dim_out, - od, - distinguish_types=distinguish_types, - ).to(env.DEVICE) - ft1 = DipoleFittingNet( - "dipole", - self.nt, - dd.dim_out, - od, - distinguish_types=distinguish_types, - ).to(env.DEVICE) - dd0 = ft0.state_dict() - 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("deep_layers", "layers") - new_kk = new_kk.replace("final_layer", "layers.3") - dd1[kk] = dd0[new_kk] - # if kk.split(".")[-1] in ["idt", "bias"]: - # dd1[kk] = dd1[kk].unsqueeze(0) - ft1.load_state_dict(dd1) - ret0 = ft0(rd0, atype, gr) - ret1 = ft1(rd0, atype, gr) - np.testing.assert_allclose( - to_numpy_array(ret0["dipole"]), - to_numpy_array(ret1["dipole"]), - ) From 41d94b5674f62033fa933387645fe2d3a78f03fb Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 18:15:58 +0800 Subject: [PATCH 17/37] chore: restore LinearAtomicModel --- deepmd/dpmodel/model/linear_atomic_model.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/deepmd/dpmodel/model/linear_atomic_model.py b/deepmd/dpmodel/model/linear_atomic_model.py index 24ad89ebcd..8a53cb9229 100644 --- a/deepmd/dpmodel/model/linear_atomic_model.py +++ b/deepmd/dpmodel/model/linear_atomic_model.py @@ -318,16 +318,14 @@ def _compute_weight( ), axis=-1, ) # handle masked nnei. - with np.errstate(divide="ignore", invalid="ignore"): - sigma = numerator / denominator + sigma = numerator / denominator u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin) coef = np.zeros_like(u) left_mask = sigma < self.sw_rmin mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax) right_mask = sigma >= self.sw_rmax coef[left_mask] = 1 - with np.errstate(invalid="ignore"): - smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 + smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 coef[mid_mask] = smooth[mid_mask] coef[right_mask] = 0 self.zbl_weight = coef From e122700d8d36121f04deeaa86829e15c73f90ae1 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 10:16:56 +0000 Subject: [PATCH 18/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/dipole.py | 2 +- deepmd/pt/model/task/fitting.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index d46ad11717..aa518d2cd3 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -64,4 +64,4 @@ def forward(self, inputs, atype, atype_tebd, rot_mat): vec_out = ( torch.bmm(vec_out, rot_mat).squeeze(-2).view(nframes, nloc, 3) ) # Shape is [nframes, nloc, 3] - return vec_out \ No newline at end of file + return vec_out diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index e8fbadbfdd..a608563a68 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -545,7 +545,7 @@ def _forward_common( ): xx = descriptor nf, nloc, nd = xx.shape - + if hasattr(self, "bias_atom_e"): self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) From 0e1182e8d5d56cc1d00732107c4a5c96274c9bd2 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 18:18:01 +0800 Subject: [PATCH 19/37] fix: ignore numpy warning --- deepmd/dpmodel/atomic_model/linear_atomic_model.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/deepmd/dpmodel/atomic_model/linear_atomic_model.py b/deepmd/dpmodel/atomic_model/linear_atomic_model.py index 48aacacdee..a9e57d6270 100644 --- a/deepmd/dpmodel/atomic_model/linear_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/linear_atomic_model.py @@ -320,14 +320,16 @@ def _compute_weight( ), axis=-1, ) # handle masked nnei. - sigma = numerator / denominator + with np.errstate(divide="ignore", invalid="ignore"): + sigma = numerator / denominator u = (sigma - self.sw_rmin) / (self.sw_rmax - self.sw_rmin) coef = np.zeros_like(u) left_mask = sigma < self.sw_rmin mid_mask = (self.sw_rmin <= sigma) & (sigma < self.sw_rmax) right_mask = sigma >= self.sw_rmax coef[left_mask] = 1 - smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 + with np.errstate(invalid="ignore"): + smooth = -6 * u**5 + 15 * u**4 - 10 * u**3 + 1 coef[mid_mask] = smooth[mid_mask] coef[right_mask] = 0 self.zbl_weight = coef From c8d97e8df362c18b4acba37b508db8d152ea0acd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 13:43:34 +0000 Subject: [PATCH 20/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 723cf49947..93d3efb81d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -80,7 +80,7 @@ def __init__( activation_function=activation_function, precision=precision, distinguish_types=distinguish_types, - rcond = rcond, + rcond=rcond, **kwargs, ) From 39a5c344907ab321c02f393a67bc57cc6ea45961 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 21:44:57 +0800 Subject: [PATCH 21/37] fix: merge conflict --- deepmd/pt/model/task/fitting.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index a608563a68..e69ebcf4f1 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -463,6 +463,7 @@ def serialize(self) -> dict: "precision": self.precision, "distinguish_types": self.distinguish_types, "nets": self.filter_layers.serialize(), + "rcond": self.rcond, "@variables": { "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), From 1c30126005ba17a4f35f37f97c08cf664047da35 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 21:48:20 +0800 Subject: [PATCH 22/37] fix: merge conflict --- deepmd/pt/model/task/ener.py | 1 - deepmd/pt/model/task/fitting.py | 2 ++ 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 93d3efb81d..5f713903dc 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -56,7 +56,6 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, - rcond: Optional[float] = None, **kwargs, ): """Construct a fitting net for energy. diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index e69ebcf4f1..b5b89311b8 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -356,6 +356,7 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, + rcond: Optional[float] = None, **kwargs, ): """Construct a fitting net for energy. @@ -381,6 +382,7 @@ def __init__( self.activation_function = activation_function self.precision = precision self.prec = PRECISION_DICT[self.precision] + self.rcond = rcond # init constants if self.numb_fparam > 0: From 1ae6ce52d386291fd7bfe48275e67400ce64609e Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 21:52:54 +0800 Subject: [PATCH 23/37] fix: merge conflict --- deepmd/pt/model/task/ener.py | 1 + deepmd/pt/model/task/fitting.py | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 5f713903dc..93d3efb81d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -56,6 +56,7 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, + rcond: Optional[float] = None, **kwargs, ): """Construct a fitting net for energy. diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index b5b89311b8..d7d611a84c 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -480,7 +480,6 @@ def serialize(self) -> dict: # "use_aparam_as_mask": self.use_aparam_as_mask , # "spin": self.spin , ## NOTICE: not supported by far - "rcond": None, "tot_ener_zero": False, "trainable": True, "atom_ener": None, From 0e98ddcb7c4fd17d5753ba9c9181c27adb01635f Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 22:20:00 +0800 Subject: [PATCH 24/37] chore: refactor --- deepmd/pt/model/task/ener.py | 22 ++++++++++++++++------ deepmd/pt/model/task/fitting.py | 29 ++++++++++++++++++++--------- 2 files changed, 36 insertions(+), 15 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 93d3efb81d..574ed06c5b 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -61,12 +61,18 @@ def __init__( ): """Construct a fitting net for energy. - Args: - - ntypes: Element count. - - embedding_width: Embedding width per atom. - - neuron: Number of neurons in each hidden layers of the fitting net. - - bias_atom_e: Average enery per atom for each element. - - resnet_dt: Using time-step in the ResNet construction. + Parameters + ---------- + + var_name: The atomic property to fit, 'energy', 'dipole', and 'polar' + ntypes: Element count. + dim_descrpt: Embedding width per atom. + dim_out: The output dimension of the fitting net. + neuron: Number of neurons in each hidden layers of the fitting net. + bias_atom_e: Average enery per atom for each element. + resnet_dt: Using time-step in the ResNet construction. + numb_fparam: Number of frame parameters. + numb_aparam: Number of atomic parameters. """ super().__init__( var_name=var_name, @@ -92,6 +98,10 @@ def __init__( assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" self.register_buffer("bias_atom_e", bias_atom_e) + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return self.dim_out + def __setitem__(self, key, value): if key in ["bias_atom_e"]: value = value.view([self.ntypes, self.dim_out]) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index d7d611a84c..6009eeeb77 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -7,6 +7,7 @@ Optional, Union, ) +from abc import abstractmethod import numpy as np import torch @@ -359,14 +360,18 @@ def __init__( rcond: Optional[float] = None, **kwargs, ): - """Construct a fitting net for energy. - - Args: - - ntypes: Element count. - - embedding_width: Embedding width per atom. - - neuron: Number of neurons in each hidden layers of the fitting net. - - bias_atom_e: Average enery per atom for each element. - - resnet_dt: Using time-step in the ResNet construction. + """Construct a general fitting net. + + Parameters + ---------- + var_name: The atomic property to fit, 'energy', 'dipole', and 'polar' + ntypes: Element count. + dim_descrpt: Embedding width per atom. + dim_out: The output dimension of the fitting net. + neuron: Number of neurons in each hidden layers of the fitting net. + resnet_dt: Using time-step in the ResNet construction. + numb_fparam: Number of frame parameters. + numb_aparam: Number of atomic parameters. """ super().__init__() self.var_name = var_name @@ -411,6 +416,7 @@ def __init__( in_dim = self.dim_descrpt + self.numb_fparam + self.numb_aparam self.old_impl = kwargs.get("old_impl", False) + net_dim_out = self._net_out_dim() if self.old_impl: filter_layers = [] for type_i in range(self.ntypes): @@ -433,7 +439,7 @@ def __init__( networks=[ FittingNet( in_dim, - self.dim_out, + net_dim_out, self.neuron, self.activation_function, self.resnet_dt, @@ -515,6 +521,11 @@ def get_sel_type(self) -> List[int]: If returning an empty list, all atom types are selected. """ return [] + + @abstractmethod + def _net_out_dim(self): + """Set the FittingNet output dim.""" + raise NotImplementedError def output_def(self) -> FittingOutputDef: return FittingOutputDef( From 22e6d82c6b9addbbc16d2a506d4325353d7fd003 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 14:20:47 +0000 Subject: [PATCH 25/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 19 +++++++++---------- deepmd/pt/model/task/fitting.py | 22 ++++++++++++---------- 2 files changed, 21 insertions(+), 20 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 574ed06c5b..1fe9c3ea50 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -63,16 +63,15 @@ def __init__( Parameters ---------- - - var_name: The atomic property to fit, 'energy', 'dipole', and 'polar' - ntypes: Element count. - dim_descrpt: Embedding width per atom. - dim_out: The output dimension of the fitting net. - neuron: Number of neurons in each hidden layers of the fitting net. - bias_atom_e: Average enery per atom for each element. - resnet_dt: Using time-step in the ResNet construction. - numb_fparam: Number of frame parameters. - numb_aparam: Number of atomic parameters. + var_name : The atomic property to fit, 'energy', 'dipole', and 'polar' + ntypes : Element count. + dim_descrpt : Embedding width per atom. + dim_out : The output dimension of the fitting net. + neuron : Number of neurons in each hidden layers of the fitting net. + bias_atom_e : Average enery per atom for each element. + resnet_dt : Using time-step in the ResNet construction. + numb_fparam : Number of frame parameters. + numb_aparam : Number of atomic parameters. """ super().__init__( var_name=var_name, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 6009eeeb77..04559c7289 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -1,13 +1,15 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging +from abc import ( + abstractmethod, +) from typing import ( Callable, List, Optional, Union, ) -from abc import abstractmethod import numpy as np import torch @@ -364,14 +366,14 @@ def __init__( Parameters ---------- - var_name: The atomic property to fit, 'energy', 'dipole', and 'polar' - ntypes: Element count. - dim_descrpt: Embedding width per atom. - dim_out: The output dimension of the fitting net. - neuron: Number of neurons in each hidden layers of the fitting net. - resnet_dt: Using time-step in the ResNet construction. - numb_fparam: Number of frame parameters. - numb_aparam: Number of atomic parameters. + var_name : The atomic property to fit, 'energy', 'dipole', and 'polar' + ntypes : Element count. + dim_descrpt : Embedding width per atom. + dim_out : The output dimension of the fitting net. + neuron : Number of neurons in each hidden layers of the fitting net. + resnet_dt : Using time-step in the ResNet construction. + numb_fparam : Number of frame parameters. + numb_aparam : Number of atomic parameters. """ super().__init__() self.var_name = var_name @@ -521,7 +523,7 @@ def get_sel_type(self) -> List[int]: If returning an empty list, all atom types are selected. """ return [] - + @abstractmethod def _net_out_dim(self): """Set the FittingNet output dim.""" From ecf242b9fb838c399999872fee57b7e35e13becc Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 23:02:18 +0800 Subject: [PATCH 26/37] chore: refactor --- deepmd/pt/model/task/ener.py | 6 +++++- deepmd/pt/model/task/fitting.py | 11 ++++++++--- 2 files changed, 13 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 1fe9c3ea50..202418ef4c 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -63,7 +63,7 @@ def __init__( Parameters ---------- - var_name : The atomic property to fit, 'energy', 'dipole', and 'polar' + var_name : The atomic property to fit, 'energy', 'dipole', and 'polar'. ntypes : Element count. dim_descrpt : Embedding width per atom. dim_out : The output dimension of the fitting net. @@ -72,6 +72,10 @@ def __init__( resnet_dt : Using time-step in the ResNet construction. numb_fparam : Number of frame parameters. numb_aparam : Number of atomic parameters. + activation_function: Activation function. + precision: Numerical precision. + distinguish_types: Neighbor list that distinguish different atomic types or not. + rcond: The condition number for the regression of atomic energy. """ super().__init__( var_name=var_name, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 04559c7289..39bbae6fe5 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -366,7 +366,7 @@ def __init__( Parameters ---------- - var_name : The atomic property to fit, 'energy', 'dipole', and 'polar' + var_name : The atomic property to fit, 'energy', 'dipole', and 'polar'. ntypes : Element count. dim_descrpt : Embedding width per atom. dim_out : The output dimension of the fitting net. @@ -374,6 +374,10 @@ def __init__( resnet_dt : Using time-step in the ResNet construction. numb_fparam : Number of frame parameters. numb_aparam : Number of atomic parameters. + activation_function: Activation function. + precision: Numerical precision. + distinguish_types: Neighbor list that distinguish different atomic types or not. + rcond: The condition number for the regression of atomic energy. """ super().__init__() self.var_name = var_name @@ -527,7 +531,7 @@ def get_sel_type(self) -> List[int]: @abstractmethod def _net_out_dim(self): """Set the FittingNet output dim.""" - raise NotImplementedError + pass def output_def(self) -> FittingOutputDef: return FittingOutputDef( @@ -639,9 +643,10 @@ def _forward_common( ) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: + net_dim_out = self._net_out_dim() 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)) + mask = torch.tile(mask, (1, 1, net_dim_out)) atom_property = ll(xx) atom_property = ( atom_property + self.bias_atom_e[type_i] From 042201ade33b5188465429a0507322bf8105a80d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 15:05:29 +0000 Subject: [PATCH 27/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 8 ++++---- deepmd/pt/model/task/fitting.py | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 202418ef4c..2bab682fa3 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -72,10 +72,10 @@ def __init__( resnet_dt : Using time-step in the ResNet construction. numb_fparam : Number of frame parameters. numb_aparam : Number of atomic parameters. - activation_function: Activation function. - precision: Numerical precision. - distinguish_types: Neighbor list that distinguish different atomic types or not. - rcond: The condition number for the regression of atomic energy. + activation_function : Activation function. + precision : Numerical precision. + distinguish_types : Neighbor list that distinguish different atomic types or not. + rcond : The condition number for the regression of atomic energy. """ super().__init__( var_name=var_name, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 39bbae6fe5..5745b3056f 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -374,10 +374,10 @@ def __init__( resnet_dt : Using time-step in the ResNet construction. numb_fparam : Number of frame parameters. numb_aparam : Number of atomic parameters. - activation_function: Activation function. - precision: Numerical precision. - distinguish_types: Neighbor list that distinguish different atomic types or not. - rcond: The condition number for the regression of atomic energy. + activation_function : Activation function. + precision : Numerical precision. + distinguish_types : Neighbor list that distinguish different atomic types or not. + rcond : The condition number for the regression of atomic energy. """ super().__init__() self.var_name = var_name From 7698483e1213ebd5fca7e37ea9326b9029594371 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Thu, 15 Feb 2024 23:40:12 +0800 Subject: [PATCH 28/37] chore: refactor --- deepmd/pt/model/task/ener.py | 53 ++++++++++++++--------------- deepmd/pt/model/task/fitting.py | 59 +++++++++++++++++++++------------ 2 files changed, 64 insertions(+), 48 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 2bab682fa3..6da466622b 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -63,19 +63,32 @@ def __init__( Parameters ---------- - var_name : The atomic property to fit, 'energy', 'dipole', and 'polar'. - ntypes : Element count. - dim_descrpt : Embedding width per atom. - dim_out : The output dimension of the fitting net. - neuron : Number of neurons in each hidden layers of the fitting net. - bias_atom_e : Average enery per atom for each element. - resnet_dt : Using time-step in the ResNet construction. - numb_fparam : Number of frame parameters. - numb_aparam : Number of atomic parameters. - activation_function : Activation function. - precision : Numerical precision. - distinguish_types : Neighbor list that distinguish different atomic types or not. - rcond : The condition number for the regression of atomic energy. + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + dim_out : int + The output dimension of the fitting net. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_e : torch.Tensor, optional + Average enery per atom for each element. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + distinguish_types : bool + Neighbor list that distinguish different atomic types or not. + rcond : float, optional + The condition number for the regression of atomic energy. """ super().__init__( var_name=var_name, @@ -83,6 +96,7 @@ def __init__( dim_descrpt=dim_descrpt, dim_out=dim_out, neuron=neuron, + bias_atom_e = bias_atom_e, resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, @@ -93,14 +107,6 @@ def __init__( **kwargs, ) - if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes, self.dim_out]) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) - bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) - if not self.use_tebd: - assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" - self.register_buffer("bias_atom_e", bias_atom_e) - def _net_out_dim(self): """Set the FittingNet output dim.""" return self.dim_out @@ -142,11 +148,6 @@ def data_stat_key(self): """ return ["bias_atom_e"] - def serialize(self): - data = super().serialize() - data["@variables"]["bias_atom_e"] = to_numpy_array(self.bias_atom_e) - return data - def compute_output_stats(self, merged): energy = [item["energy"] for item in merged] mixed_type = "real_natoms_vec" in merged[0] diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 5745b3056f..827424ff93 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -353,6 +353,7 @@ def __init__( dim_descrpt: int, dim_out: int, neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, @@ -366,18 +367,32 @@ def __init__( Parameters ---------- - var_name : The atomic property to fit, 'energy', 'dipole', and 'polar'. - ntypes : Element count. - dim_descrpt : Embedding width per atom. - dim_out : The output dimension of the fitting net. - neuron : Number of neurons in each hidden layers of the fitting net. - resnet_dt : Using time-step in the ResNet construction. - numb_fparam : Number of frame parameters. - numb_aparam : Number of atomic parameters. - activation_function : Activation function. - precision : Numerical precision. - distinguish_types : Neighbor list that distinguish different atomic types or not. - rcond : The condition number for the regression of atomic energy. + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + dim_out : int + The output dimension of the fitting net. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_e : torch.Tensor, optional + Average enery per atom for each element. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + distinguish_types : bool + Neighbor list that distinguish different atomic types or not. + rcond : float, optional + The condition number for the regression of atomic energy. """ super().__init__() self.var_name = var_name @@ -396,6 +411,14 @@ def __init__( self.rcond = rcond # init constants + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, self.dim_out]) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, self.dim_out]) + if not self.use_tebd: + assert self.ntypes == bias_atom_e.shape[0], "Element count mismatches!" + self.register_buffer("bias_atom_e", bias_atom_e) + if self.numb_fparam > 0: self.register_buffer( "fparam_avg", @@ -479,6 +502,7 @@ def serialize(self) -> dict: "nets": self.filter_layers.serialize(), "rcond": self.rcond, "@variables": { + "bias_atom_e": to_numpy_array(self.bias_atom_e), "fparam_avg": to_numpy_array(self.fparam_avg), "fparam_inv_std": to_numpy_array(self.fparam_inv_std), "aparam_avg": to_numpy_array(self.aparam_avg), @@ -565,9 +589,6 @@ def _forward_common( xx = descriptor nf, nloc, nd = xx.shape - if hasattr(self, "bias_atom_e"): - self.bias_atom_e = self.bias_atom_e.view([self.ntypes, self.dim_out]) - if nd != self.dim_descrpt: raise ValueError( "get an input descriptor of dim {nd}," @@ -613,7 +634,7 @@ def _forward_common( dim=-1, ) - outs = torch.zeros(nf, nloc, self.dim_out) # jit assertion + outs = torch.zeros((nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION).to(env.DEVICE) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None @@ -628,8 +649,6 @@ def _forward_common( atom_property = filter_layer(xx) atom_property = ( atom_property + self.bias_atom_e[type_i] - if hasattr(self, "bias_atom_e") - else atom_property ) atom_property = atom_property * mask.unsqueeze(-1) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] @@ -638,8 +657,6 @@ def _forward_common( if self.use_tebd: atom_property = ( (self.filter_layers.networks[0](xx) + self.bias_atom_e[atype]) - if hasattr(self, "bias_atom_e") - else self.filter_layers.networks[0](xx) ) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] else: @@ -650,8 +667,6 @@ def _forward_common( atom_property = ll(xx) atom_property = ( atom_property + self.bias_atom_e[type_i] - if hasattr(self, "bias_atom_e") - else atom_property ) atom_property = atom_property * mask outs = outs + atom_property # Shape is [nframes, natoms[0], 1] From c3be87cb96b26071289c3fd62b2fb8b8da5e6629 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 15:42:21 +0000 Subject: [PATCH 29/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 5 +---- deepmd/pt/model/task/fitting.py | 14 ++++++-------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 6da466622b..8f5e5e4219 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -30,9 +30,6 @@ from deepmd.pt.utils.stat import ( compute_output_bias, ) -from deepmd.pt.utils.utils import ( - to_numpy_array, -) dtype = env.GLOBAL_PT_FLOAT_PRECISION device = env.DEVICE @@ -96,7 +93,7 @@ def __init__( dim_descrpt=dim_descrpt, dim_out=dim_out, neuron=neuron, - bias_atom_e = bias_atom_e, + bias_atom_e=bias_atom_e, resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 827424ff93..0d567f86da 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -634,7 +634,9 @@ def _forward_common( dim=-1, ) - outs = torch.zeros((nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION).to(env.DEVICE) # jit assertion + outs = torch.zeros( + (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION + ).to(env.DEVICE) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None @@ -647,16 +649,14 @@ def _forward_common( for type_i, filter_layer in enumerate(self.filter_layers_old): mask = atype == type_i atom_property = filter_layer(xx) - atom_property = ( - atom_property + self.bias_atom_e[type_i] - ) + atom_property = atom_property + self.bias_atom_e[type_i] atom_property = atom_property * mask.unsqueeze(-1) outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} else: if self.use_tebd: atom_property = ( - (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_property # Shape is [nframes, natoms[0], 1] else: @@ -665,9 +665,7 @@ def _forward_common( mask = (atype == type_i).unsqueeze(-1) mask = torch.tile(mask, (1, 1, net_dim_out)) atom_property = ll(xx) - atom_property = ( - atom_property + self.bias_atom_e[type_i] - ) + atom_property = atom_property + self.bias_atom_e[type_i] atom_property = atom_property * mask outs = outs + atom_property # Shape is [nframes, natoms[0], 1] return {self.var_name: outs.to(env.GLOBAL_PT_FLOAT_PRECISION)} From 367e0adeaffab2b20a0637d3403010aa26a7afb8 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 16 Feb 2024 00:32:14 +0800 Subject: [PATCH 30/37] fix: revert device --- deepmd/pt/model/task/ener.py | 4 ++++ deepmd/pt/model/task/fitting.py | 14 ++++++++------ 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 8f5e5e4219..6f3a20ed0b 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -54,6 +54,7 @@ def __init__( precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, rcond: Optional[float] = None, + seed: Optional[int] = None, **kwargs, ): """Construct a fitting net for energy. @@ -86,6 +87,8 @@ def __init__( Neighbor list that distinguish different atomic types or not. rcond : float, optional The condition number for the regression of atomic energy. + seed : int, optional + Random seed. """ super().__init__( var_name=var_name, @@ -101,6 +104,7 @@ def __init__( precision=precision, distinguish_types=distinguish_types, rcond=rcond, + seed=seed, **kwargs, ) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 0d567f86da..3b2cb8de16 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -361,6 +361,7 @@ def __init__( precision: str = DEFAULT_PRECISION, distinguish_types: bool = False, rcond: Optional[float] = None, + seed: Optional[int] = None, **kwargs, ): """Construct a general fitting net. @@ -393,6 +394,8 @@ def __init__( Neighbor list that distinguish different atomic types or not. rcond : float, optional The condition number for the regression of atomic energy. + seed : int, optional + Random seed. """ super().__init__() self.var_name = var_name @@ -480,10 +483,9 @@ def __init__( ) self.filter_layers_old = None - # very bad design... - if "seed" in kwargs: - log.info("Set seed to %d in fitting net.", kwargs["seed"]) - torch.manual_seed(kwargs["seed"]) + if seed is not None: + log.info("Set seed to %d in fitting net.", seed) + torch.manual_seed(seed) def serialize(self) -> dict: """Serialize the fitting to dict.""" @@ -635,8 +637,8 @@ def _forward_common( ) outs = torch.zeros( - (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION - ).to(env.DEVICE) # jit assertion + (nf, nloc, self.dim_out) + ) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None From 709b0105e8b315a39e6e8846acb9a52d1a6c892f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 15 Feb 2024 16:33:52 +0000 Subject: [PATCH 31/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/fitting.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 3b2cb8de16..3eacea9943 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -636,9 +636,7 @@ def _forward_common( dim=-1, ) - outs = torch.zeros( - (nf, nloc, self.dim_out) - ) # jit assertion + outs = torch.zeros((nf, nloc, self.dim_out)) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None From b8bdbbc436cf88a2eb19198cb596c7d21b32f5ea Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 16 Feb 2024 10:09:22 +0800 Subject: [PATCH 32/37] fix: add device --- deepmd/pt/model/task/ener.py | 40 +++++++++++++++--------------- deepmd/pt/model/task/fitting.py | 44 +++++++++++++++++---------------- 2 files changed, 43 insertions(+), 41 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 6f3a20ed0b..aa9f7e76c7 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -39,25 +39,7 @@ @fitting_check_output class InvarFitting(GeneralFitting): - def __init__( - self, - var_name: str, - ntypes: int, - dim_descrpt: int, - dim_out: int, - neuron: List[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, - resnet_dt: bool = True, - numb_fparam: int = 0, - numb_aparam: int = 0, - activation_function: str = "tanh", - precision: str = DEFAULT_PRECISION, - distinguish_types: bool = False, - rcond: Optional[float] = None, - seed: Optional[int] = None, - **kwargs, - ): - """Construct a fitting net for energy. + """Construct a fitting net for energy. Parameters ---------- @@ -89,7 +71,25 @@ def __init__( The condition number for the regression of atomic energy. seed : int, optional Random seed. - """ + """ + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, + rcond: Optional[float] = None, + seed: Optional[int] = None, + **kwargs, + ): super().__init__( var_name=var_name, ntypes=ntypes, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 3eacea9943..4ed3931f6e 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -346,25 +346,7 @@ def change_energy_bias( class GeneralFitting(Fitting): - def __init__( - self, - var_name: str, - ntypes: int, - dim_descrpt: int, - dim_out: int, - neuron: List[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, - resnet_dt: bool = True, - numb_fparam: int = 0, - numb_aparam: int = 0, - activation_function: str = "tanh", - precision: str = DEFAULT_PRECISION, - distinguish_types: bool = False, - rcond: Optional[float] = None, - seed: Optional[int] = None, - **kwargs, - ): - """Construct a general fitting net. + """Construct a general fitting net. Parameters ---------- @@ -396,7 +378,25 @@ def __init__( The condition number for the regression of atomic energy. seed : int, optional Random seed. - """ + """ + def __init__( + self, + var_name: str, + ntypes: int, + dim_descrpt: int, + dim_out: int, + neuron: List[int] = [128, 128, 128], + bias_atom_e: Optional[torch.Tensor] = None, + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + distinguish_types: bool = False, + rcond: Optional[float] = None, + seed: Optional[int] = None, + **kwargs, + ): super().__init__() self.var_name = var_name self.ntypes = ntypes @@ -636,7 +636,9 @@ def _forward_common( dim=-1, ) - outs = torch.zeros((nf, nloc, self.dim_out)) # jit assertion + outs = torch.zeros( + (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION + ).to(env.DEVICE) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None From 4bb7737383efeb2f148d63de6d17ff0295715c25 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 16 Feb 2024 02:11:56 +0000 Subject: [PATCH 33/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/ener.py | 61 +++++++++++++++++---------------- deepmd/pt/model/task/fitting.py | 61 +++++++++++++++++---------------- 2 files changed, 62 insertions(+), 60 deletions(-) diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index aa9f7e76c7..f1dad4c58d 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -41,37 +41,38 @@ class InvarFitting(GeneralFitting): """Construct a fitting net for energy. - Parameters - ---------- - var_name : str - The atomic property to fit, 'energy', 'dipole', and 'polar'. - ntypes : int - Element count. - dim_descrpt : int - Embedding width per atom. - dim_out : int - The output dimension of the fitting net. - neuron : List[int] - Number of neurons in each hidden layers of the fitting net. - bias_atom_e : torch.Tensor, optional - Average enery per atom for each element. - resnet_dt : bool - Using time-step in the ResNet construction. - numb_fparam : int - Number of frame parameters. - numb_aparam : int - Number of atomic parameters. - activation_function : str - Activation function. - precision : str - Numerical precision. - distinguish_types : bool - Neighbor list that distinguish different atomic types or not. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. + Parameters + ---------- + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + dim_out : int + The output dimension of the fitting net. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_e : torch.Tensor, optional + Average enery per atom for each element. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + distinguish_types : bool + Neighbor list that distinguish different atomic types or not. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. """ + def __init__( self, var_name: str, diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 4ed3931f6e..7ecff10248 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -348,37 +348,38 @@ def change_energy_bias( class GeneralFitting(Fitting): """Construct a general fitting net. - Parameters - ---------- - var_name : str - The atomic property to fit, 'energy', 'dipole', and 'polar'. - ntypes : int - Element count. - dim_descrpt : int - Embedding width per atom. - dim_out : int - The output dimension of the fitting net. - neuron : List[int] - Number of neurons in each hidden layers of the fitting net. - bias_atom_e : torch.Tensor, optional - Average enery per atom for each element. - resnet_dt : bool - Using time-step in the ResNet construction. - numb_fparam : int - Number of frame parameters. - numb_aparam : int - Number of atomic parameters. - activation_function : str - Activation function. - precision : str - Numerical precision. - distinguish_types : bool - Neighbor list that distinguish different atomic types or not. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. + Parameters + ---------- + var_name : str + The atomic property to fit, 'energy', 'dipole', and 'polar'. + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + dim_out : int + The output dimension of the fitting net. + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_e : torch.Tensor, optional + Average enery per atom for each element. + resnet_dt : bool + Using time-step in the ResNet construction. + numb_fparam : int + Number of frame parameters. + numb_aparam : int + Number of atomic parameters. + activation_function : str + Activation function. + precision : str + Numerical precision. + distinguish_types : bool + Neighbor list that distinguish different atomic types or not. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. """ + def __init__( self, var_name: str, From 7f9292aab18475397372da44a2e20ee717231610 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 16 Feb 2024 11:23:20 +0800 Subject: [PATCH 34/37] fix: cuda --- deepmd/pt/model/task/fitting.py | 4 ++-- source/tests/pt/model/test_make_hessian_model.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 7ecff10248..0a76dfcb26 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -638,8 +638,8 @@ def _forward_common( ) outs = torch.zeros( - (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION - ).to(env.DEVICE) # jit assertion + (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion assert self.filter_layers_old is not None diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index 81aee758bf..ad2f96c003 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -72,8 +72,8 @@ def test( cell0 = 1.0 * (cell0 + cell0.T) + 5.0 * torch.eye(3) cell1 = torch.rand([3, 3], dtype=dtype) cell1 = 1.0 * (cell1 + cell1.T) + 5.0 * torch.eye(3) - cell = torch.stack([cell0, cell1]) - coord = torch.rand([nf, natoms, 3], dtype=dtype) + cell = torch.stack([cell0, cell1]).to(env.DEVICE) + coord = torch.rand([nf, natoms, 3], dtype=dtype).to(env.DEVICE) coord = torch.matmul(coord, cell) cell = cell.view([nf, 9]) coord = coord.view([nf, natoms * 3]) @@ -82,10 +82,10 @@ def test( torch.IntTensor([0, 0, 1]), torch.IntTensor([1, 0, 1]), ] - ).view([nf, natoms]) + ).view([nf, natoms]).to(env.DEVICE) nfp, nap = 2, 3 - fparam = torch.rand([nf, nfp], dtype=dtype) - aparam = torch.rand([nf, natoms * nap], dtype=dtype) + fparam = torch.rand([nf, nfp], dtype=dtype).to(env.DEVICE) + aparam = torch.rand([nf, natoms * nap], dtype=dtype).to(env.DEVICE) # forward hess and valu models ret_dict0 = self.model_hess.forward_common( coord, atype, box=cell, fparam=fparam, aparam=aparam From 4ff301924bf086b16f63dccfe7ad4e2dcb36987d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 16 Feb 2024 03:23:54 +0000 Subject: [PATCH 35/37] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/fitting.py | 4 +++- source/tests/pt/model/test_make_hessian_model.py | 16 ++++++++++------ 2 files changed, 13 insertions(+), 7 deletions(-) diff --git a/deepmd/pt/model/task/fitting.py b/deepmd/pt/model/task/fitting.py index 0a76dfcb26..b2d8c875ce 100644 --- a/deepmd/pt/model/task/fitting.py +++ b/deepmd/pt/model/task/fitting.py @@ -638,7 +638,9 @@ def _forward_common( ) outs = torch.zeros( - (nf, nloc, self.dim_out), dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + (nf, nloc, self.dim_out), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=env.DEVICE, ) # jit assertion if self.old_impl: outs = torch.zeros_like(atype).unsqueeze(-1) # jit assertion diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index ad2f96c003..ed42aa6739 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -77,12 +77,16 @@ def test( coord = torch.matmul(coord, cell) cell = cell.view([nf, 9]) coord = coord.view([nf, natoms * 3]) - atype = torch.stack( - [ - torch.IntTensor([0, 0, 1]), - torch.IntTensor([1, 0, 1]), - ] - ).view([nf, natoms]).to(env.DEVICE) + atype = ( + torch.stack( + [ + torch.IntTensor([0, 0, 1]), + torch.IntTensor([1, 0, 1]), + ] + ) + .view([nf, natoms]) + .to(env.DEVICE) + ) nfp, nap = 2, 3 fparam = torch.rand([nf, nfp], dtype=dtype).to(env.DEVICE) aparam = torch.rand([nf, natoms * nap], dtype=dtype).to(env.DEVICE) From 1fd58a4a3d0fd9b872a74007f66de9ecbaddcab6 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 16 Feb 2024 12:20:07 +0800 Subject: [PATCH 36/37] fix: cuda --- source/tests/pt/model/test_make_hessian_model.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index ed42aa6739..48c35a19f1 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -68,12 +68,12 @@ def test( natoms = self.nloc nf = self.nf nv = self.nv - cell0 = torch.rand([3, 3], dtype=dtype) + cell0 = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) cell0 = 1.0 * (cell0 + cell0.T) + 5.0 * torch.eye(3) - cell1 = torch.rand([3, 3], dtype=dtype) + cell1 = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) cell1 = 1.0 * (cell1 + cell1.T) + 5.0 * torch.eye(3) - cell = torch.stack([cell0, cell1]).to(env.DEVICE) - coord = torch.rand([nf, natoms, 3], dtype=dtype).to(env.DEVICE) + cell = torch.stack([cell0, cell1]) + coord = torch.rand([nf, natoms, 3], dtype=dtype, device=env.DEVICE) coord = torch.matmul(coord, cell) cell = cell.view([nf, 9]) coord = coord.view([nf, natoms * 3]) @@ -88,8 +88,8 @@ def test( .to(env.DEVICE) ) nfp, nap = 2, 3 - fparam = torch.rand([nf, nfp], dtype=dtype).to(env.DEVICE) - aparam = torch.rand([nf, natoms * nap], dtype=dtype).to(env.DEVICE) + fparam = torch.rand([nf, nfp], dtype=dtype, device=env.DEVICE) + aparam = torch.rand([nf, natoms * nap], dtype=dtype, device=env.DEVICE) # forward hess and valu models ret_dict0 = self.model_hess.forward_common( coord, atype, box=cell, fparam=fparam, aparam=aparam From e125be19d49fff5a2ecdf35cb69c1dd174a88178 Mon Sep 17 00:00:00 2001 From: Anyang Peng <137014849+anyangml@users.noreply.github.com> Date: Fri, 16 Feb 2024 14:22:43 +0800 Subject: [PATCH 37/37] fix: cuda --- source/tests/pt/model/test_make_hessian_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/tests/pt/model/test_make_hessian_model.py b/source/tests/pt/model/test_make_hessian_model.py index 48c35a19f1..6f321b6478 100644 --- a/source/tests/pt/model/test_make_hessian_model.py +++ b/source/tests/pt/model/test_make_hessian_model.py @@ -69,9 +69,9 @@ def test( nf = self.nf nv = self.nv cell0 = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) - cell0 = 1.0 * (cell0 + cell0.T) + 5.0 * torch.eye(3) + cell0 = 1.0 * (cell0 + cell0.T) + 5.0 * torch.eye(3, device=env.DEVICE) cell1 = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) - cell1 = 1.0 * (cell1 + cell1.T) + 5.0 * torch.eye(3) + cell1 = 1.0 * (cell1 + cell1.T) + 5.0 * torch.eye(3, device=env.DEVICE) cell = torch.stack([cell0, cell1]) coord = torch.rand([nf, natoms, 3], dtype=dtype, device=env.DEVICE) coord = torch.matmul(coord, cell)