From d1d0a0a93445ef12183475be3fc95493d7020e95 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Sat, 16 Mar 2024 17:53:15 +0000 Subject: [PATCH 01/31] 3.16 update:support property fitting(only zero bias and mean pooling). Not include denoise and data format. --- deepmd/pt/model/model/dp_model.py | 8 + deepmd/pt/model/model/property_model.py | 85 ++++++++ deepmd/pt/model/task/__init__.py | 4 + deepmd/pt/model/task/property.py | 134 ++++++++++++ deepmd/utils/argcheck.py | 35 ++++ .../tests/pt/model/test_property_fitting.py | 194 ++++++++++++++++++ 6 files changed, 460 insertions(+) create mode 100644 deepmd/pt/model/model/property_model.py create mode 100644 deepmd/pt/model/task/property.py create mode 100644 source/tests/pt/model/test_property_fitting.py diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index d7b3c4f4e2..1b29eebd9f 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -28,6 +28,9 @@ from deepmd.pt.model.task.polarizability import ( PolarFittingNet, ) +from deepmd.pt.model.task.property import ( + PropertyFittingNet, +) from .make_model import ( make_model, @@ -57,6 +60,9 @@ def __new__( from deepmd.pt.model.model.polar_model import ( PolarModel, ) + from deepmd.pt.model.model.property_model import ( + PropertyModel, + ) if atomic_model_ is not None: fitting = atomic_model_.fitting_net @@ -76,6 +82,8 @@ def __new__( cls = PolarModel elif isinstance(fitting, DOSFittingNet): cls = DOSModel + elif isinstance(fitting, PropertyFittingNet): + cls = PropertyModel # else: unknown fitting type, fall back to DPModel return super().__new__(cls) diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py new file mode 100644 index 0000000000..dc6372fd2e --- /dev/null +++ b/deepmd/pt/model/model/property_model.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + Optional, +) + +import torch + +from .dp_model import ( + DPModel, +) + + +class PropertyModel(DPModel): + model_type = "property" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = {} + model_predict["atom_property"] = model_ret["property"] + model_predict["property"] = model_ret["property_redu"]/atype.shape[-1] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ): + model_ret = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + if self.get_fitting_net() is not None: + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy"): + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["extended_virial"] = model_ret[ + "energy_derv_c" + ].squeeze(-3) + else: + assert model_ret["dforce"] is not None + model_predict["dforce"] = model_ret["dforce"] + else: + model_predict = model_ret + return model_predict diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index 9430ede766..4f43343f76 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -24,6 +24,9 @@ from .type_predict import ( TypePredictNet, ) +from .property import ( + PropertyFittingNet, +) __all__ = [ "FittingNetAttenLcc", @@ -35,4 +38,5 @@ "BaseFitting", "TypePredictNet", "PolarFittingNet", + "PropertyFittingNet", ] diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py new file mode 100644 index 0000000000..3a22b6caef --- /dev/null +++ b/deepmd/pt/model/task/property.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +from typing import ( + Callable, + Union, + List, + Optional, + Tuple, +) + +import numpy as np +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) +from deepmd.pt.model.network.network import ( + ResidualDeep, +) +from deepmd.pt.model.task.fitting import ( + Fitting, + GeneralFitting, +) +from deepmd.pt.model.task.invar_fitting import ( + InvarFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, +) +from deepmd.utils.version import ( + check_version_compatibility, +) +from deepmd.utils.path import ( + DPPath, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("property") +class PropertyFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + dim_descrpt: 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, + mixed_types: bool = True, + **kwargs, + ): + super().__init__( + "property", + ntypes, + dim_descrpt, + 1, + neuron=neuron, + bias_atom_e=bias_atom_e, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + activation_function=activation_function, + precision=precision, + mixed_types=mixed_types, + **kwargs, + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("var_name") + data.pop("dim_out") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + **super().serialize(), + "type": "property", + } + + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + # make jit happy with torch 2.0.0 + exclude_types: List[int] \ No newline at end of file diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 1c2b7935b5..91b0245050 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1175,6 +1175,41 @@ def fitting_dipole(): Argument("seed", [int, None], optional=True, doc=doc_seed), ] +@fitting_args_plugin.register("property") +def fitting_property(): + doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." + doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." + doc_seed = "Random seed for parameter initialization of the fitting net" + return [ + Argument( + "neuron", + List[int], + optional=True, + default=[120, 120, 120], + alias=["n_neuron"], + doc=doc_neuron, + ), + Argument( + "activation_function", + str, + optional=True, + default="tanh", + doc=doc_activation_function, + ), + Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument( + "sel_type", + [List[int], int, None], + optional=True, + alias=["property_type"], + doc=doc_sel_type + doc_only_tf_supported, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), + ] # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py new file mode 100644 index 0000000000..33a6d2d39c --- /dev/null +++ b/source/tests/pt/model/test_property_fitting.py @@ -0,0 +1,194 @@ +# 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.property import ( + PropertyFittingNet, + InvarFitting, +) +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_consistency( + self, + ): + rng = np.random.default_rng() + nf, nloc, nnei = self.nlist.shape + dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + rd0, _, _, _, _ = dd0( + 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) + + for od, mixed_types, nfp, nap, et, nn in itertools.product( + [1, 3], + [True, False], + [0, 3], + [0, 4], + [[], [0], [1]], + [[4, 4, 4], []], + ): + ft0 = InvarFitting( + "foo", + self.nt, + dd0.dim_out, + od, + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + exclude_types=et, + neuron=nn, + ).to(env.DEVICE) + ft1 = DPInvarFitting.deserialize(ft0.serialize()) + ft2 = InvarFitting.deserialize(ft0.serialize()) + + if nfp > 0: + ifp = torch.tensor( + rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + rng.normal(size=(self.nf, self.nloc, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + ret0 = ft0(rd0, atype, fparam=ifp, aparam=iap) + ret1 = ft1( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + ret2 = ft2(rd0, atype, fparam=ifp, aparam=iap) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + ret1["foo"], + ) + np.testing.assert_allclose( + to_numpy_array(ret0["foo"]), + to_numpy_array(ret2["foo"]), + ) + self.assertEqual(ft0.get_sel_type(), ft1.get_sel_type()) + + @unittest.skip("Don't need to test new&old in property fitting") + 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, _, _, _, _ = 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 = 1 + for foo, mixed_types in itertools.product( + [True], + [True, False], + ): + ft0 = PropertyFittingNet( + self.nt, + dd.dim_out, + mixed_types=mixed_types, + ).to(env.DEVICE) + ft1 = PropertyFittingNet( + self.nt, + dd.dim_out, + mixed_types=mixed_types, + old_impl=True, + ).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) + dd1["bias_atom_e"] = dd0["bias_atom_e"] + ft1.load_state_dict(dd1) + ret0 = ft0(rd0, atype) + ret1 = ft1(rd0, atype) + np.testing.assert_allclose( + to_numpy_array(ret0["property"]), + to_numpy_array(ret1["property"]), + ) + + @unittest.skip("Don't need to test jit at this time") + def test_jit( + self, + ): + for od, mixed_types, nfp, nap, et in itertools.product( + [1, 3], + [True, False], + [0, 3], + [0, 4], + [[], [0]], + ): + ft0 = InvarFitting( + "foo", + self.nt, + 9, + od, + numb_fparam=nfp, + numb_aparam=nap, + mixed_types=mixed_types, + exclude_types=et, + ).to(env.DEVICE) + torch.jit.script(ft0) + + def test_get_set(self): + ifn0 = InvarFitting( + "property", + self.nt, + 3, + 1, + ) + print(ifn0) + rng = np.random.default_rng() + foo = rng.normal([3, 4]) + for ii in [ + "bias_atom_e", + "fparam_avg", + "fparam_inv_std", + "aparam_avg", + "aparam_inv_std", + ]: + ifn0[ii] = torch.tensor(foo, dtype=dtype, device=env.DEVICE) + np.testing.assert_allclose( + foo, np.reshape(ifn0[ii].detach().cpu().numpy(), foo.shape) + ) From b9a7be4a8d340b4816275ed3fc361972eaa67a8c Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 12:27:35 +0000 Subject: [PATCH 02/31] 3.18 update --- deepmd/pt/loss/__init__.py | 4 + deepmd/pt/loss/property.py | 110 ++++ .../pt/model/atomic_model/dp_atomic_model.py | 1 + deepmd/pt/model/model/__init__.py | 2 +- deepmd/pt/model/model/denoise_model.py | 81 +++ deepmd/pt/model/model/dp_model.py | 8 + deepmd/pt/model/task/denoise.py | 531 +++++++++++++++++- deepmd/pt/model/task/invar_fitting.py | 3 +- deepmd/pt/model/task/property.py | 7 +- deepmd/pt/train/training.py | 5 + deepmd/utils/argcheck.py | 62 ++ 11 files changed, 802 insertions(+), 12 deletions(-) create mode 100644 deepmd/pt/loss/property.py create mode 100644 deepmd/pt/model/model/denoise_model.py diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 9c8bbc9a2a..c23df6b6d3 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -14,6 +14,9 @@ from .tensor import ( TensorLoss, ) +from .property import ( + PropertyLoss, +) __all__ = [ "DenoiseLoss", @@ -21,4 +24,5 @@ "EnergySpinLoss", "TensorLoss", "TaskLoss", + "PropertyLoss", ] diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py new file mode 100644 index 0000000000..69826f7f69 --- /dev/null +++ b/deepmd/pt/loss/property.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + List, +) + +import torch +import torch.nn.functional as F + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + GLOBAL_PT_FLOAT_PRECISION, +) +from deepmd.utils.data import ( + DataRequirementItem, +) +import logging +log = logging.getLogger(__name__) + +class PropertyLoss(TaskLoss): + def __init__( + self, + task_num, + loss_func: str = "smooth_mae", + metric: list = ["mae"], + **kwargs, + ): + r"""Construct a layer to compute loss on property. + + Parameters + ---------- + starter_learning_rate : float + The learning rate at the start of the training. + loss_func: str + The loss function, such as "smooth_mae", "mae", "rmse" + metric: list + The metric such as mae,rmse which will be printed. + **kwargs + Other keyword arguments. + """ + super().__init__() + self.loss_func = loss_func + self.metric = metric + self.task_num = task_num + self.mean = kwargs.get("mean", 0) + self.std = kwargs.get("std", 1) + self.beta = kwargs.get("beta", 1.00) + + def forward(self, model_pred, label, natoms, learning_rate, mae=False): + """Return loss on loss and force. + + Args: + - model_pred: Property prediction. + - label: Target property. + - natoms: Tell atom count. + + Returns + ------- + - loss: Loss to minimize. + """ + assert label['property'].shape[-1] == self.task_num + assert model_pred['property'].shape[-1] == self.task_num + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] + more_loss = {} + + label_mean = torch.tensor(self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + label_std = torch.tensor(self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + + # loss + if self.loss_func == "smooth_mae": + loss += F.smooth_l1_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum", beta=self.beta) + elif self.func == "mae": + loss += F.l1_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum") + elif self.func == "mse": + loss += F.mse_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum") + elif self.func == "rmse": + loss += torch.sqrt(F.mse_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="mean")) + else: + raise RuntimeError(f"Unknown loss function : {self.func}") + + # more loss + if "smooth_mae" in self.metric: + more_loss["smooth_mae"] = F.smooth_l1_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean", beta=self.beta).detach() + if "mae" in self.metric: + more_loss['mae'] = F.l1_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean").detach() + if "mse" in self.metric: + more_loss['mse'] = F.mse_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean").detach() + if "rmse" in self.metric: + more_loss['rmse'] = torch.sqrt(F.mse_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean")).detach() + + return loss, more_loss + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + label_requirement.append( + DataRequirementItem( + "property", + ndof=self.task_num, + atomic=False, + must=False, + high_prec=True, + ) + ) + return label_requirement diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index 6aa8df7aee..fd8f71fa2d 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -176,6 +176,7 @@ def forward_atomic( gr=rot_mat, g2=g2, h2=h2, + sw=sw, fparam=fparam, aparam=aparam, ) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index 7a2070e476..c42f369318 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -140,7 +140,7 @@ def get_standard_model(model_params): fitting_net["type"] = fitting_net.get("type", "ener") fitting_net["ntypes"] = descriptor.get_ntypes() fitting_net["mixed_types"] = descriptor.mixed_types() - if fitting_net["type"] in ["dipole", "polar"]: + if fitting_net["type"] in ["dipole", "polar","denoise"]: fitting_net["embedding_width"] = descriptor.get_dim_emb() fitting_net["dim_descrpt"] = descriptor.get_dim_out() grad_force = "direct" not in fitting_net["type"] diff --git a/deepmd/pt/model/model/denoise_model.py b/deepmd/pt/model/model/denoise_model.py new file mode 100644 index 0000000000..249d995f7b --- /dev/null +++ b/deepmd/pt/model/model/denoise_model.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Dict, + Optional, +) + +import torch + +from .dp_model import ( + DPModel, +) + + +class DenoiseModel(DPModel): + model_type = "denoise" + + def __init__( + self, + *args, + **kwargs, + ): + super().__init__(*args, **kwargs) + + def forward( + self, + coord, + atype, + box: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ) -> Dict[str, torch.Tensor]: + model_ret = self.forward_common( + coord, + atype, + box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + model_predict = model_ret + return model_predict + + @torch.jit.export + def forward_lower( + self, + extended_coord, + extended_atype, + nlist, + mapping: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + do_atomic_virial: bool = False, + ): + model_ret = self.forward_common_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + if self.get_fitting_net() is not None: + model_predict = {} + model_predict["atom_energy"] = model_ret["energy"] + model_predict["energy"] = model_ret["energy_redu"] + if self.do_grad_r("energy"): + model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) + if self.do_grad_c("energy"): + model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) + if do_atomic_virial: + model_predict["extended_virial"] = model_ret[ + "energy_derv_c" + ].squeeze(-3) + else: + assert model_ret["dforce"] is not None + model_predict["dforce"] = model_ret["dforce"] + else: + model_predict = model_ret + return model_predict \ No newline at end of file diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 1b29eebd9f..495a233d0b 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -31,6 +31,9 @@ from deepmd.pt.model.task.property import ( PropertyFittingNet, ) +from deepmd.pt.model.task.denoise import ( + DenoiseFittingNet +) from .make_model import ( make_model, @@ -63,6 +66,9 @@ def __new__( from deepmd.pt.model.model.property_model import ( PropertyModel, ) + from deepmd.pt.model.model.denoise_model import ( + DenoiseModel, + ) if atomic_model_ is not None: fitting = atomic_model_.fitting_net @@ -84,6 +90,8 @@ def __new__( cls = DOSModel elif isinstance(fitting, PropertyFittingNet): cls = PropertyModel + elif isinstance(fitting, DenoiseFittingNet): + cls = DenoiseModel # else: unknown fitting type, fall back to DPModel return super().__new__(cls) diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 5f1e780de3..26fb4742c5 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -1,10 +1,58 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +import os +import tempfile +from abc import ( + abstractmethod, +) from typing import ( + Callable, + List, Optional, + Union, ) +import numpy as np import torch +from deepmd.infer.deep_eval import ( + DeepEval, +) +from deepmd.pt.model.network.mlp import ( + FittingNet, + NetworkCollection, +) +from deepmd.pt.model.network.network import ( + ResidualDeep, +) +from deepmd.pt.model.task.base_fitting import ( + BaseFitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + DEVICE, + PRECISION_DICT, +) +from deepmd.pt.utils.exclude_mask import ( + AtomExcludeMask, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.utils.data_system import ( + DeepmdDataSystem, +) +from deepmd.utils.finetune import ( + change_energy_bias_lower, +) +from deepmd.pt.model.task.fitting import ( + Fitting, +) from deepmd.dpmodel import ( FittingOutputDef, OutputVariableDef, @@ -14,13 +62,486 @@ MaskLMHead, NonLinearHead, ) -from deepmd.pt.model.task.fitting import ( - Fitting, -) -from deepmd.pt.utils import ( - env, +from deepmd.utils.path import ( + DPPath, ) +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + +@Fitting.register("denoise") +class DenoiseFittingNet(Fitting): + """Construct a denoise 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. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. + exclude_types: List[int] + Atomic contributions of the excluded atom types are set zero. + trainable : Union[List[bool], bool] + If the parameters in the fitting net are trainable. + Now this only supports setting all the parameters in the fitting net at one state. + When in List[bool], the trainable will be True only if all the boolean parameters are True. + remove_vaccum_contribution: List[bool], optional + Remove vaccum contribution before the bias is added. The list assigned each + type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same + length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + embedding_width: 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, + mixed_types: bool = True, + rcond: Optional[float] = None, + seed: Optional[int] = None, + exclude_types: List[int] = [], + trainable: Union[bool, List[bool]] = True, + remove_vaccum_contribution: Optional[List[bool]] = None, + **kwargs, + ): + super().__init__() + self.var_name = ["updated_coord","logits"] + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.embedding_width = embedding_width + self.neuron = neuron + self.mixed_types = mixed_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] + self.rcond = rcond + # order matters, should be place after the assignment of ntypes + self.reinit_exclude(exclude_types) + self.trainable = trainable + # need support for each layer settings + self.trainable = ( + all(self.trainable) if isinstance(self.trainable, list) else self.trainable + ) + self.remove_vaccum_contribution = remove_vaccum_contribution + + # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] + net_dim_out = self._net_out_dim() + + # init constants + # TODO: actually bias is useless in denoise task + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) + if not self.mixed_types: + 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", + 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_coord = self.embedding_width + in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam + + # the first MLP is used to update coordinate + self.filter_layers_coord = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_coord, + 1, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + # the second MLP is used to update logits + self.filter_layers_logits = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_logits, + net_dim_out[1], + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + self.filter_layers_old = None + + if seed is not None: + torch.manual_seed(seed) + # set trainable + for param in self.parameters(): + param.requires_grad = self.trainable + + def reinit_exclude( + self, + exclude_types: List[int] = [], + ): + self.exclude_types = exclude_types + self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "@class": "Fitting", + "@version": 1, + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "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, + "mixed_types": self.mixed_types, + "nets": self.filter_layers.serialize(), + "rcond": self.rcond, + "exclude_types": self.exclude_types, + "@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), + "aparam_inv_std": to_numpy_array(self.aparam_inv_std), + }, + # "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 + "tot_ener_zero": False, + "trainable": [self.trainable] * (len(self.neuron) + 1), + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "DenoiseFitting": + 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 + + # make jit happy + exclude_types: List[int] + + 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. + """ + # make jit happy + sel_type: List[int] = [] + for ii in range(self.ntypes): + if ii not in self.exclude_types: + sel_type.append(ii) + return sel_type + + def __setitem__(self, key, value): + if key in ["bias_atom_e"]: + value = value.view([self.ntypes, self._net_out_dim()]) + 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 + elif key in ["scale"]: + self.scale = 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 + elif key in ["scale"]: + return self.scale + else: + raise KeyError(key) + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return [3, self.ntypes] + #pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "updated_coord", + [3], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + OutputVariableDef( + "logits", + [-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + 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( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + xx = descriptor + nnei = g2.shape[-2] + if self.remove_vaccum_contribution is not None: + # TODO: Idealy, the input for vaccum should be computed; + # we consider it as always zero for convenience. + # Needs a compute_input_stats for vaccum passed from the + # descriptor. + xx_zeros = torch.zeros_like(xx) + else: + xx_zeros = None + nf, nloc, nd = xx.shape + net_dim_out = self._net_out_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([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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, 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( + f"get an input aparam of dim {aparam.shape[-1]}, ", + f"which is not consistent with {self.numb_aparam}.", + ) + aparam = aparam.view([nf, -1, 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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, aparam], + dim=-1, + ) + + outs_coord = torch.zeros( + (nf, nloc, net_dim_out[0]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) # jit assertion + outs_logits = torch.zeros( + (nf, nloc, net_dim_out[1]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) + if self.mixed_types: + atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum(dim=-2) / (sw.sum(dim=-1).unsqueeze(-1)+1e-6) + atom_logits = self.filter_layers_logits[0](xx) + #Is xx_zeros useful in denoise task?????????????? + #if xx_zeros is not None: + # atom_property -= self.filter_layers.networks[0](xx_zeros) + outs_coord = ( + outs_coord + atom_updated_coord + ) # Shape is [nframes, natoms[0], net_dim_out] + outs_logits = ( + outs_logits + atom_logits + ) + # TODO: + else: + for type_i, ll in enumerate(self.filter_layers_coord.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = torch.tile(mask, (1, 1, net_dim_out)) + atom_property = ll(xx) + if xx_zeros is not None: + # must assert, otherwise jit is not happy + assert self.remove_vaccum_contribution is not None + if not ( + len(self.remove_vaccum_contribution) > type_i + and not self.remove_vaccum_contribution[type_i] + ): + atom_property -= ll(xx_zeros) + atom_property = atom_property + self.bias_atom_e[type_i] + atom_property = atom_property * mask + outs = ( + outs + atom_property + ) # Shape is [nframes, natoms[0], net_dim_out] + # nf x nloc + mask = self.emask(atype) + # nf x nloc x nod + outs_coord = outs_coord * mask[:, :, None] + outs_logits = outs_logits * mask[:, :, None] + return {self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), + self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION)} + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass @fitting_check_output class DenoiseNet(Fitting): diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 1699b440ac..b0ff1323eb 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -195,8 +195,9 @@ def forward( h2: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, ): - """Based on embedding net output, alculate total energy. + """Based on embedding net output, calculate total energy. Args: - inputs: Embedding matrix. Its shape is [nframes, natoms[0], self.dim_descrpt]. diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 3a22b6caef..4ecc74e798 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -15,10 +15,6 @@ from deepmd.dpmodel import ( FittingOutputDef, OutputVariableDef, - fitting_check_output, -) -from deepmd.pt.model.network.network import ( - ResidualDeep, ) from deepmd.pt.model.task.fitting import ( Fitting, @@ -52,6 +48,7 @@ def __init__( self, ntypes: int, dim_descrpt: int, + task_num: int = 1, neuron: List[int] = [128, 128, 128], bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, @@ -66,7 +63,7 @@ def __init__( "property", ntypes, dim_descrpt, - 1, + task_num, neuron=neuron, bias_atom_e=bias_atom_e, resnet_dt=resnet_dt, diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index fc293f70ec..4480800ec8 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -28,6 +28,7 @@ EnergySpinLoss, EnergyStdLoss, TensorLoss, + PropertyLoss, ) from deepmd.pt.model.model import ( DPZBLModel, @@ -297,6 +298,10 @@ def get_loss(loss_params, start_lr, _ntypes, _model): label_name = "polarizability" loss_params["label_name"] = label_name return TensorLoss(**loss_params) + elif loss_type == "property": + task_num = _model.model_output_def()["property"].output_size + loss_params["task_num"] = task_num + return PropertyLoss(**loss_params) else: raise NotImplementedError diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 91b0245050..b3124cb859 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1183,6 +1183,7 @@ def fitting_property(): doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." doc_seed = "Random seed for parameter initialization of the fitting net" + doc_task_num = "The number of outputs of fitting net" return [ Argument( "neuron", @@ -1209,6 +1210,43 @@ def fitting_property(): doc=doc_sel_type + doc_only_tf_supported, ), Argument("seed", [int, None], optional=True, doc=doc_seed), + Argument("task_num", int, optional=True, default=1, doc=doc_task_num), + ] + +@fitting_args_plugin.register("denoise") +def fitting_denoise(): + doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." + doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' + doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' + doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." + doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." + doc_seed = "Random seed for parameter initialization of the fitting net" + return [ + Argument( + "neuron", + List[int], + optional=True, + default=[120, 120, 120], + alias=["n_neuron"], + doc=doc_neuron, + ), + Argument( + "activation_function", + str, + optional=True, + default="tanh", + doc=doc_activation_function, + ), + Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), + Argument("precision", str, optional=True, default="default", doc=doc_precision), + Argument( + "sel_type", + [List[int], int, None], + optional=True, + alias=["denoise_type"], + doc=doc_sel_type + doc_only_tf_supported, + ), + Argument("seed", [int, None], optional=True, doc=doc_seed), ] # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. @@ -1953,6 +1991,30 @@ def loss_tensor(): ), ] +@loss_args_plugin.register("property") +def loss_property(): + doc_loss_func = "The loss function, such as 'mae','smooth_mae'." + doc_metric = "The metric such as 'mae','smooth_mae' which will be printed." + doc_mean = "The averge value of target." + doc_std = "The standard deviation of the target." + doc_beta = "The 'beta' parameter in 'smooth_mae' loss." + return [ + Argument( + "loss_func", str, optional=True, default="smooth_mae", doc=doc_loss_func, + ), + Argument( + "metric", list, optional=True, default=["mae"], doc=doc_metric, + ), + Argument( + "mean", [float, int, list], optional=True, default=0, doc=doc_mean, + ), + Argument( + "std", [float, int, list], optional=True, default=1, doc=doc_std, + ), + Argument( + "beta", [float, int], optional=True, default=1.00, doc=doc_beta, + ), + ] def loss_variant_type_args(): doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener` or left unset. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." From 3f95a82d013e14cc9ef1f96eb5026811cfe24949 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 15:47:39 +0000 Subject: [PATCH 03/31] Add DeepProperty and UT --- deepmd/dpmodel/fitting/__init__.py | 4 + deepmd/dpmodel/fitting/property_fitting.py | 137 +++++++ deepmd/infer/deep_eval.py | 2 + deepmd/infer/deep_property.py | 136 +++++++ deepmd/pt/infer/deep_eval.py | 5 + deepmd/pt/model/task/property.py | 16 +- .../tests/pt/model/test_property_fitting.py | 358 +++++++++++++----- 7 files changed, 548 insertions(+), 110 deletions(-) create mode 100644 deepmd/dpmodel/fitting/property_fitting.py create mode 100644 deepmd/infer/deep_property.py diff --git a/deepmd/dpmodel/fitting/__init__.py b/deepmd/dpmodel/fitting/__init__.py index 866a710a3b..2a35a68d35 100644 --- a/deepmd/dpmodel/fitting/__init__.py +++ b/deepmd/dpmodel/fitting/__init__.py @@ -17,6 +17,9 @@ from .polarizability_fitting import ( PolarFitting, ) +from .property_fitting import ( + PropertyFittingNet, +) __all__ = [ "InvarFitting", @@ -25,4 +28,5 @@ "EnergyFittingNet", "PolarFitting", "DOSFittingNet", + "PropertyFittingNet", ] diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py new file mode 100644 index 0000000000..ce03ca0be8 --- /dev/null +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +from typing import ( + Callable, + Union, + TYPE_CHECKING, + Any, + List, + Optional, +) + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + OutputVariableDef, +) +from deepmd.utils.path import ( + DPPath, +) + +from deepmd.dpmodel.common import ( + DEFAULT_PRECISION, +) +from deepmd.dpmodel.fitting.invar_fitting import ( + InvarFitting, +) + +if TYPE_CHECKING: + from deepmd.dpmodel.fitting.general_fitting import ( + GeneralFitting, + ) +from deepmd.utils.version import ( + check_version_compatibility, +) + + +@InvarFitting.register("property") +class PropertyFittingNet(InvarFitting): + def __init__( + self, + ntypes: int, + dim_descrpt: int, + task_num: int = 1, + neuron: List[int] = [128, 128, 128], + resnet_dt: bool = True, + numb_fparam: int = 0, + numb_aparam: int = 0, + rcond: Optional[float] = None, + tot_ener_zero: bool = False, + trainable: Optional[List[bool]] = None, + atom_ener: Optional[List[float]] = None, + activation_function: str = "tanh", + precision: str = DEFAULT_PRECISION, + layer_name: Optional[List[Optional[str]]] = None, + use_aparam_as_mask: bool = False, + spin: Any = None, + mixed_types: bool = False, + exclude_types: List[int] = [], + # not used + seed: Optional[int] = None, + ): + self.task_num = task_num + super().__init__( + var_name="property", + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=task_num, + neuron=neuron, + resnet_dt=resnet_dt, + numb_fparam=numb_fparam, + numb_aparam=numb_aparam, + rcond=rcond, + tot_ener_zero=tot_ener_zero, + trainable=trainable, + atom_ener=atom_ener, + activation_function=activation_function, + precision=precision, + layer_name=layer_name, + use_aparam_as_mask=use_aparam_as_mask, + spin=spin, + mixed_types=mixed_types, + exclude_types=exclude_types, + ) + + @classmethod + def deserialize(cls, data: dict) -> "GeneralFitting": + data = copy.deepcopy(data) + check_version_compatibility(data.pop("@version", 1), 1, 1) + data.pop("var_name") + data.pop("dim_out") + return super().deserialize(data) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + **super().serialize(), + "type": "property", + "task_num": self.task_num + } + + def output_def(self) -> FittingOutputDef: + return FittingOutputDef( + [ + OutputVariableDef( + self.var_name, + [self.dim_out], + reduciable=True, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + # make jit happy with torch 2.0.0 + exclude_types: List[int] \ No newline at end of file diff --git a/deepmd/infer/deep_eval.py b/deepmd/infer/deep_eval.py index aae2082e13..5d1cc6ed43 100644 --- a/deepmd/infer/deep_eval.py +++ b/deepmd/infer/deep_eval.py @@ -73,6 +73,8 @@ class DeepEvalBackend(ABC): "dipole_derv_c_redu": "virial", "dos": "atom_dos", "dos_redu": "dos", + "property": "atom_property", + "property_redu": "property", "mask_mag": "mask_mag", "mask": "mask", } diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py new file mode 100644 index 0000000000..d8eea8608f --- /dev/null +++ b/deepmd/infer/deep_property.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + Dict, + List, + Optional, + Tuple, + Union, +) + +import numpy as np + +from deepmd.dpmodel.output_def import ( + FittingOutputDef, + ModelOutputDef, + OutputVariableDef, +) + +from .deep_eval import ( + DeepEval, +) + + +class DeepProperty(DeepEval): + """Properties of structures. + + Parameters + ---------- + model_file : Path + The name of the frozen model file. + *args : list + Positional arguments. + auto_batch_size : bool or int or AutoBatchSize, default: True + If True, automatic batch size will be used. If int, it will be used + as the initial batch size. + neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional + The ASE neighbor list class to produce the neighbor list. If None, the + neighbor list will be built natively in the model. + **kwargs : dict + Keyword arguments. + """ + + @property + def output_def(self) -> ModelOutputDef: + """Get the output definition of this model.""" + return ModelOutputDef( + FittingOutputDef( + [ + OutputVariableDef( + "property", + shape=[-1], + reduciable=True, + atomic=True, + ), + ] + ) + ) + + @property + def numb_task(self) -> int: + """Get the number of task.""" + return self.get_numb_task() + + def eval( + self, + coords: np.ndarray, + cells: Optional[np.ndarray], + atom_types: Union[List[int], np.ndarray], + atomic: bool = False, + fparam: Optional[np.ndarray] = None, + aparam: Optional[np.ndarray] = None, + mixed_type: bool = False, + **kwargs: Dict[str, Any], + ) -> Tuple[np.ndarray, ...]: + """Evaluate energy, force, and virial. If atomic is True, + also return atomic energy and atomic virial. + + Parameters + ---------- + coords : np.ndarray + The coordinates of the atoms, in shape (nframes, natoms, 3). + cells : np.ndarray + The cell vectors of the system, in shape (nframes, 9). If the system + is not periodic, set it to None. + atom_types : List[int] or np.ndarray + The types of the atoms. If mixed_type is False, the shape is (natoms,); + otherwise, the shape is (nframes, natoms). + atomic : bool, optional + Whether to return atomic energy and atomic virial, by default False. + fparam : np.ndarray, optional + The frame parameters, by default None. + aparam : np.ndarray, optional + The atomic parameters, by default None. + mixed_type : bool, optional + Whether the atom_types is mixed type, by default False. + **kwargs : Dict[str, Any] + Keyword arguments. + + Returns + ------- + property + The properties of the system, in shape (nframes, num_tasks). + """ + ( + coords, + cells, + atom_types, + fparam, + aparam, + nframes, + natoms, + ) = self._standard_input(coords, cells, atom_types, fparam, aparam, mixed_type) + results = self.deep_eval.eval( + coords, + cells, + atom_types, + atomic, + fparam=fparam, + aparam=aparam, + **kwargs, + ) + atomic_property = results["property"].reshape(nframes, natoms, -1) + property = np.sum(atomic_property, axis=1) + + if atomic: + return ( + property, + atomic_property, + ) + else: + return (property,) + + def get_numb_task(self) -> int: + return self.deep_eval.get_numb_task() + +__all__ = ["DeepProperty"] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index b8031993c0..aaf37b96d8 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -38,6 +38,9 @@ from deepmd.infer.deep_wfc import ( DeepWFC, ) +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.pt.model.model import ( get_model, ) @@ -164,6 +167,8 @@ def model_type(self) -> "DeepEvalWrapper": return DeepGlobalPolar elif "wfc" in model_output_type: return DeepWFC + elif "property" in model_output_type: + return DeepProperty else: raise RuntimeError("Unknown model type") diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 4ecc74e798..8524974058 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -6,7 +6,6 @@ Union, List, Optional, - Tuple, ) import numpy as np @@ -59,11 +58,12 @@ def __init__( mixed_types: bool = True, **kwargs, ): + self.task_num = task_num super().__init__( - "property", - ntypes, - dim_descrpt, - task_num, + var_name="property", + ntypes=ntypes, + dim_descrpt=dim_descrpt, + dim_out=task_num, neuron=neuron, bias_atom_e=bias_atom_e, resnet_dt=resnet_dt, @@ -79,8 +79,8 @@ def __init__( def deserialize(cls, data: dict) -> "GeneralFitting": data = copy.deepcopy(data) check_version_compatibility(data.pop("@version", 1), 1, 1) - data.pop("var_name") data.pop("dim_out") + data.pop("var_name") return super().deserialize(data) def serialize(self) -> dict: @@ -88,9 +88,9 @@ def serialize(self) -> dict: return { **super().serialize(), "type": "property", + "task_num": self.task_num } - - + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index 33a6d2d39c..64d1d2c4f5 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -1,21 +1,33 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import itertools +import os import unittest import numpy as np import torch +from scipy.stats import ( + special_ortho_group, +) -from deepmd.dpmodel.fitting import InvarFitting as DPInvarFitting +from deepmd.dpmodel.fitting import PropertyFittingNet as DPProperFittingNet +from deepmd.infer.deep_property import ( + DeepProperty, +) from deepmd.pt.model.descriptor.se_a import ( DescrptSeA, ) +from deepmd.pt.model.model.property_model import ( + PropertyModel, +) from deepmd.pt.model.task.property import ( PropertyFittingNet, - InvarFitting, ) from deepmd.pt.utils import ( env, ) +from deepmd.pt.utils.nlist import ( + extend_input_and_build_neighbor_list, +) from deepmd.pt.utils.utils import ( to_numpy_array, ) @@ -27,168 +39,310 @@ dtype = env.GLOBAL_PT_FLOAT_PRECISION -class TestInvarFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): +class TestPropertyFitting(unittest.TestCase, TestCaseSingleFrameWithNlist): def setUp(self): TestCaseSingleFrameWithNlist.setUp(self) + self.rng = np.random.default_rng() + self.nf, self.nloc, _ = self.nlist.shape + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() def test_consistency( self, ): - rng = np.random.default_rng() - nf, nloc, nnei = self.nlist.shape - dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) - rd0, _, _, _, _ = dd0( + rd0, gr, _, _, _ = self.dd0( 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) + atype = torch.tensor( + self.atype_ext[:, : self.nloc], dtype=int, device=env.DEVICE + ) - for od, mixed_types, nfp, nap, et, nn in itertools.product( - [1, 3], + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( [True, False], [0, 3], [0, 4], - [[], [0], [1]], - [[4, 4, 4], []], + [True, False], + [None, self.scale], ): - ft0 = InvarFitting( - "foo", + ft0 = PropertyFittingNet( self.nt, - dd0.dim_out, - od, + self.dd0.dim_out, + task_num=5, numb_fparam=nfp, numb_aparam=nap, mixed_types=mixed_types, - exclude_types=et, - neuron=nn, + fit_diag=fit_diag, + scale=scale, ).to(env.DEVICE) - ft1 = DPInvarFitting.deserialize(ft0.serialize()) - ft2 = InvarFitting.deserialize(ft0.serialize()) + + ft1 = DPProperFittingNet.deserialize(ft0.serialize()) + ft2 = PropertyFittingNet.deserialize(ft0.serialize()) + ft3 = DPProperFittingNet.deserialize(ft1.serialize()) if nfp > 0: ifp = torch.tensor( - rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE ) else: ifp = None if nap > 0: iap = torch.tensor( - rng.normal(size=(self.nf, self.nloc, nap)), + self.rng.normal(size=(self.nf, self.nloc, nap)), dtype=dtype, device=env.DEVICE, ) else: iap = None - ret0 = ft0(rd0, atype, fparam=ifp, aparam=iap) + ret0 = ft0(rd0, atype, gr, fparam=ifp, aparam=iap) ret1 = ft1( rd0.detach().cpu().numpy(), atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), fparam=to_numpy_array(ifp), aparam=to_numpy_array(iap), ) - ret2 = ft2(rd0, atype, fparam=ifp, aparam=iap) + ret2 = ft2(rd0, atype, gr, fparam=ifp, aparam=iap) + ret3 = ft3( + rd0.detach().cpu().numpy(), + atype.detach().cpu().numpy(), + gr.detach().cpu().numpy(), + fparam=to_numpy_array(ifp), + aparam=to_numpy_array(iap), + ) + np.testing.assert_allclose( + to_numpy_array(ret0["property"]), + ret1["property"], + ) np.testing.assert_allclose( - to_numpy_array(ret0["foo"]), - ret1["foo"], + to_numpy_array(ret0["property"]), + to_numpy_array(ret2["property"]), ) np.testing.assert_allclose( - to_numpy_array(ret0["foo"]), - to_numpy_array(ret2["foo"]), + to_numpy_array(ret0["property"]), + ret3["property"], ) - self.assertEqual(ft0.get_sel_type(), ft1.get_sel_type()) - @unittest.skip("Don't need to test new&old in property fitting") - def test_new_old( + def test_jit( 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, _, _, _, _ = 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 = 1 - for foo, mixed_types in itertools.product( - [True], + for mixed_types, nfp, nap, fit_diag in itertools.product( + [True, False], + [0, 3], + [0, 4], [True, False], ): ft0 = PropertyFittingNet( self.nt, - dd.dim_out, - mixed_types=mixed_types, - ).to(env.DEVICE) - ft1 = PropertyFittingNet( - self.nt, - dd.dim_out, + self.dd0.dim_out, + task_num=5, + numb_fparam=nfp, + numb_aparam=nap, mixed_types=mixed_types, - old_impl=True, + fit_diag=fit_diag, ).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) - dd1["bias_atom_e"] = dd0["bias_atom_e"] - ft1.load_state_dict(dd1) - ret0 = ft0(rd0, atype) - ret1 = ft1(rd0, atype) - np.testing.assert_allclose( - to_numpy_array(ret0["property"]), - to_numpy_array(ret1["property"]), - ) + torch.jit.script(ft0) - @unittest.skip("Don't need to test jit at this time") - def test_jit( - self, - ): - for od, mixed_types, nfp, nap, et in itertools.product( - [1, 3], +class TestInvariance(unittest.TestCase): + def setUp(self) -> None: + self.natoms = 5 + self.rcut = 4 + self.rcut_smth = 0.5 + self.sel = [46, 92, 4] + self.nf = 1 + self.nt = 3 + self.rng = np.random.default_rng() + self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device=env.DEVICE) + self.shift = torch.tensor([4, 4, 4], dtype=dtype, device=env.DEVICE) + self.atype = torch.tensor([0, 0, 0, 1, 1], dtype=torch.int32, device=env.DEVICE) + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.cell = torch.rand([3, 3], dtype=dtype, device=env.DEVICE) + self.cell = (self.cell + self.cell.T) + 5.0 * torch.eye(3, device=env.DEVICE) + self.scale = self.rng.uniform(0, 1, self.nt).tolist() + + def test_rot(self): + atype = self.atype.reshape(1, 5) + rmat = torch.tensor(special_ortho_group.rvs(3), dtype=dtype, device=env.DEVICE) + coord_rot = torch.matmul(self.coord, rmat) + + for mixed_types, nfp, nap, fit_diag, scale in itertools.product( [True, False], [0, 3], [0, 4], - [[], [0]], + [True, False], + [None, self.scale], ): - ft0 = InvarFitting( - "foo", + ft0 = PropertyFittingNet( self.nt, - 9, - od, + self.dd0.dim_out, # dim_descrpt + task_num=9, numb_fparam=nfp, numb_aparam=nap, - mixed_types=mixed_types, - exclude_types=et, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, ).to(env.DEVICE) - torch.jit.script(ft0) + if nfp > 0: + ifp = torch.tensor( + self.rng.normal(size=(self.nf, nfp)), dtype=dtype, device=env.DEVICE + ) + else: + ifp = None + if nap > 0: + iap = torch.tensor( + self.rng.normal(size=(self.nf, self.natoms, nap)), + dtype=dtype, + device=env.DEVICE, + ) + else: + iap = None + + res = [] + for xyz in [self.coord, coord_rot]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz + self.shift, atype, self.rcut, self.sel, mixed_types + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=ifp, aparam=iap) + res.append(ret0["property"]) + np.testing.assert_allclose( + to_numpy_array(res[1]), + to_numpy_array(res[0]), + ) + + def test_permu(self): + coord = torch.matmul(self.coord, self.cell) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_num = 8, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for idx_perm in [[0, 1, 2, 3, 4], [1, 0, 4, 3, 2]]: + atype = self.atype[idx_perm].reshape(1, 5) + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + coord[idx_perm], atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=None, aparam=None) + res.append(ret0["property"]) - def test_get_set(self): - ifn0 = InvarFitting( - "property", - self.nt, - 3, - 1, - ) - print(ifn0) - rng = np.random.default_rng() - foo = rng.normal([3, 4]) - for ii in [ - "bias_atom_e", - "fparam_avg", - "fparam_inv_std", - "aparam_avg", - "aparam_inv_std", - ]: - ifn0[ii] = torch.tensor(foo, dtype=dtype, device=env.DEVICE) np.testing.assert_allclose( - foo, np.reshape(ifn0[ii].detach().cpu().numpy(), foo.shape) + to_numpy_array(res[0][:, idx_perm]), + to_numpy_array(res[1]), ) + + def test_trans(self): + atype = self.atype.reshape(1, 5) + coord_s = torch.matmul( + torch.remainder( + torch.matmul(self.coord + self.shift, torch.linalg.inv(self.cell)), 1.0 + ), + self.cell, + ) + for fit_diag, scale in itertools.product([True, False], [None, self.scale]): + ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_num = 11, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + fit_diag=fit_diag, + scale=scale, + ).to(env.DEVICE) + res = [] + for xyz in [self.coord, coord_s]: + ( + extended_coord, + extended_atype, + _, + nlist, + ) = extend_input_and_build_neighbor_list( + xyz, atype, self.rcut, self.sel, False + ) + + rd0, gr0, _, _, _ = self.dd0( + extended_coord, + extended_atype, + nlist, + ) + + ret0 = ft0(rd0, extended_atype, gr0, fparam=0, aparam=0) + res.append(ret0["property"]) + + np.testing.assert_allclose(to_numpy_array(res[0]), to_numpy_array(res[1])) + + +class TestPropertyModel(unittest.TestCase): + def setUp(self): + self.natoms = 5 + self.rcut = 4.0 + self.nt = 3 + self.rcut_smth = 0.5 + self.sel = [46, 92, 4] + self.nf = 1 + self.coord = 2 * torch.rand([self.natoms, 3], dtype=dtype, device="cpu") + cell = torch.rand([3, 3], dtype=dtype, device="cpu") + self.cell = (cell + cell.T) + 5.0 * torch.eye(3, device="cpu") + self.atype = torch.IntTensor([0, 0, 0, 1, 1], device="cpu") + self.dd0 = DescrptSeA(self.rcut, self.rcut_smth, self.sel).to(env.DEVICE) + self.ft0 = PropertyFittingNet( + self.nt, + self.dd0.dim_out, + task_num = 3, + numb_fparam=0, + numb_aparam=0, + mixed_types=True, + ).to(env.DEVICE) + self.type_mapping = ["O", "H", "B"] + self.model = PropertyModel(self.dd0, self.ft0, self.type_mapping) + self.file_path = "model_output.pth" + + def test_deepproperty_infer(self): + atype = self.atype.view(self.nf, self.natoms) + coord = self.coord.reshape(1, 5, 3) + cell = self.cell.reshape(1, 9) + jit_md = torch.jit.script(self.model) + torch.jit.save(jit_md, self.file_path) + load_md = DeepProperty(self.file_path) + load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=True) + load_md.eval(coords=coord, atom_types=atype, cells=cell, atomic=False) + + def tearDown(self) -> None: + if os.path.exists(self.file_path): + os.remove(self.file_path) + + +if __name__ == "__main__": + unittest.main() \ No newline at end of file From 786b52881a81f2c2ade35ebd0ac6e13f86275633 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 Mar 2024 16:01:36 +0000 Subject: [PATCH 04/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/fitting/property_fitting.py | 26 +++--- deepmd/infer/deep_property.py | 1 + deepmd/pt/infer/deep_eval.py | 6 +- deepmd/pt/loss/__init__.py | 6 +- deepmd/pt/loss/property.py | 79 ++++++++++++++----- deepmd/pt/model/model/__init__.py | 2 +- deepmd/pt/model/model/denoise_model.py | 4 +- deepmd/pt/model/model/dp_model.py | 12 +-- deepmd/pt/model/model/property_model.py | 2 +- deepmd/pt/model/task/__init__.py | 6 +- deepmd/pt/model/task/denoise.py | 61 ++++++-------- deepmd/pt/model/task/property.py | 19 ++--- deepmd/pt/train/training.py | 2 +- deepmd/utils/argcheck.py | 35 ++++++-- .../tests/pt/model/test_property_fitting.py | 9 ++- 15 files changed, 156 insertions(+), 114 deletions(-) diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index ce03ca0be8..056a1af861 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -1,14 +1,20 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy from typing import ( - Callable, - Union, TYPE_CHECKING, Any, + Callable, List, Optional, + Union, ) +from deepmd.dpmodel.common import ( + DEFAULT_PRECISION, +) +from deepmd.dpmodel.fitting.invar_fitting import ( + InvarFitting, +) from deepmd.dpmodel.output_def import ( FittingOutputDef, OutputVariableDef, @@ -17,17 +23,11 @@ DPPath, ) -from deepmd.dpmodel.common import ( - DEFAULT_PRECISION, -) -from deepmd.dpmodel.fitting.invar_fitting import ( - InvarFitting, -) - if TYPE_CHECKING: from deepmd.dpmodel.fitting.general_fitting import ( GeneralFitting, ) + from deepmd.utils.version import ( check_version_compatibility, ) @@ -91,11 +91,7 @@ def deserialize(cls, data: dict) -> "GeneralFitting": def serialize(self) -> dict: """Serialize the fitting to dict.""" - return { - **super().serialize(), - "type": "property", - "task_num": self.task_num - } + return {**super().serialize(), "type": "property", "task_num": self.task_num} def output_def(self) -> FittingOutputDef: return FittingOutputDef( @@ -134,4 +130,4 @@ def compute_output_stats( pass # make jit happy with torch 2.0.0 - exclude_types: List[int] \ No newline at end of file + exclude_types: List[int] diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index d8eea8608f..f53d7acd25 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -133,4 +133,5 @@ def eval( def get_numb_task(self) -> int: return self.deep_eval.get_numb_task() + __all__ = ["DeepProperty"] diff --git a/deepmd/pt/infer/deep_eval.py b/deepmd/pt/infer/deep_eval.py index aaf37b96d8..f83a3c4782 100644 --- a/deepmd/pt/infer/deep_eval.py +++ b/deepmd/pt/infer/deep_eval.py @@ -35,12 +35,12 @@ from deepmd.infer.deep_pot import ( DeepPot, ) -from deepmd.infer.deep_wfc import ( - DeepWFC, -) from deepmd.infer.deep_property import ( DeepProperty, ) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) from deepmd.pt.model.model import ( get_model, ) diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index c23df6b6d3..0656db8dbe 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -11,12 +11,12 @@ from .loss import ( TaskLoss, ) -from .tensor import ( - TensorLoss, -) from .property import ( PropertyLoss, ) +from .tensor import ( + TensorLoss, +) __all__ = [ "DenoiseLoss", diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 69826f7f69..5aaaf0fcf4 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -1,4 +1,5 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +import logging from typing import ( List, ) @@ -12,15 +13,13 @@ from deepmd.pt.utils import ( env, ) -from deepmd.pt.utils.env import ( - GLOBAL_PT_FLOAT_PRECISION, -) from deepmd.utils.data import ( DataRequirementItem, ) -import logging + log = logging.getLogger(__name__) + class PropertyLoss(TaskLoss): def __init__( self, @@ -33,11 +32,11 @@ def __init__( Parameters ---------- - starter_learning_rate : float + task_num : float The learning rate at the start of the training. - loss_func: str + loss_func : str The loss function, such as "smooth_mae", "mae", "rmse" - metric: list + metric : list The metric such as mae,rmse which will be printed. **kwargs Other keyword arguments. @@ -62,35 +61,77 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): ------- - loss: Loss to minimize. """ - assert label['property'].shape[-1] == self.task_num - assert model_pred['property'].shape[-1] == self.task_num + assert label["property"].shape[-1] == self.task_num + assert model_pred["property"].shape[-1] == self.task_num loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} - label_mean = torch.tensor(self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) - label_std = torch.tensor(self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE) + label_mean = torch.tensor( + self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + label_std = torch.tensor( + self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) # loss if self.loss_func == "smooth_mae": - loss += F.smooth_l1_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum", beta=self.beta) + loss += F.smooth_l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + beta=self.beta, + ) elif self.func == "mae": - loss += F.l1_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum") + loss += F.l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) elif self.func == "mse": - loss += F.mse_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="sum") + loss += F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) elif self.func == "rmse": - loss += torch.sqrt(F.mse_loss((label['property'] - label_mean) / label_std, model_pred['property'], reduction="mean")) + loss += torch.sqrt( + F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="mean", + ) + ) else: raise RuntimeError(f"Unknown loss function : {self.func}") # more loss if "smooth_mae" in self.metric: - more_loss["smooth_mae"] = F.smooth_l1_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean", beta=self.beta).detach() + more_loss["smooth_mae"] = F.smooth_l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + beta=self.beta, + ).detach() if "mae" in self.metric: - more_loss['mae'] = F.l1_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean").detach() + more_loss["mae"] = F.l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() if "mse" in self.metric: - more_loss['mse'] = F.mse_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean").detach() + more_loss["mse"] = F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() if "rmse" in self.metric: - more_loss['rmse'] = torch.sqrt(F.mse_loss(label['property'], (model_pred['property'] * label_std) + label_mean, reduction="mean")).detach() + more_loss["rmse"] = torch.sqrt( + F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ) + ).detach() return loss, more_loss diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index c42f369318..efbef794c0 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -140,7 +140,7 @@ def get_standard_model(model_params): fitting_net["type"] = fitting_net.get("type", "ener") fitting_net["ntypes"] = descriptor.get_ntypes() fitting_net["mixed_types"] = descriptor.mixed_types() - if fitting_net["type"] in ["dipole", "polar","denoise"]: + if fitting_net["type"] in ["dipole", "polar", "denoise"]: fitting_net["embedding_width"] = descriptor.get_dim_emb() fitting_net["dim_descrpt"] = descriptor.get_dim_out() grad_force = "direct" not in fitting_net["type"] diff --git a/deepmd/pt/model/model/denoise_model.py b/deepmd/pt/model/model/denoise_model.py index 249d995f7b..f5db6f5835 100644 --- a/deepmd/pt/model/model/denoise_model.py +++ b/deepmd/pt/model/model/denoise_model.py @@ -37,7 +37,7 @@ def forward( fparam=fparam, aparam=aparam, do_atomic_virial=do_atomic_virial, - ) + ) model_predict = model_ret return model_predict @@ -78,4 +78,4 @@ def forward_lower( model_predict["dforce"] = model_ret["dforce"] else: model_predict = model_ret - return model_predict \ No newline at end of file + return model_predict diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 495a233d0b..1cf8282981 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -15,6 +15,9 @@ from deepmd.pt.model.model.model import ( BaseModel, ) +from deepmd.pt.model.task.denoise import ( + DenoiseFittingNet, +) from deepmd.pt.model.task.dipole import ( DipoleFittingNet, ) @@ -31,9 +34,6 @@ from deepmd.pt.model.task.property import ( PropertyFittingNet, ) -from deepmd.pt.model.task.denoise import ( - DenoiseFittingNet -) from .make_model import ( make_model, @@ -51,6 +51,9 @@ def __new__( atomic_model_: Optional[DPAtomicModel] = None, **kwargs, ): + from deepmd.pt.model.model.denoise_model import ( + DenoiseModel, + ) from deepmd.pt.model.model.dipole_model import ( DipoleModel, ) @@ -66,9 +69,6 @@ def __new__( from deepmd.pt.model.model.property_model import ( PropertyModel, ) - from deepmd.pt.model.model.denoise_model import ( - DenoiseModel, - ) if atomic_model_ is not None: fitting = atomic_model_.fitting_net diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index dc6372fd2e..66f58ca665 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -40,7 +40,7 @@ def forward( ) model_predict = {} model_predict["atom_property"] = model_ret["property"] - model_predict["property"] = model_ret["property_redu"]/atype.shape[-1] + model_predict["property"] = model_ret["property_redu"] / atype.shape[-1] if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict diff --git a/deepmd/pt/model/task/__init__.py b/deepmd/pt/model/task/__init__.py index 4f43343f76..9906d435cb 100644 --- a/deepmd/pt/model/task/__init__.py +++ b/deepmd/pt/model/task/__init__.py @@ -21,12 +21,12 @@ from .polarizability import ( PolarFittingNet, ) -from .type_predict import ( - TypePredictNet, -) from .property import ( PropertyFittingNet, ) +from .type_predict import ( + TypePredictNet, +) __all__ = [ "FittingNetAttenLcc", diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 26fb4742c5..6ec85ef71d 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -1,11 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import copy import logging -import os -import tempfile -from abc import ( - abstractmethod, -) from typing import ( Callable, List, @@ -16,25 +11,27 @@ import numpy as np import torch -from deepmd.infer.deep_eval import ( - DeepEval, +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, ) from deepmd.pt.model.network.mlp import ( FittingNet, NetworkCollection, ) from deepmd.pt.model.network.network import ( - ResidualDeep, + MaskLMHead, + NonLinearHead, ) -from deepmd.pt.model.task.base_fitting import ( - BaseFitting, +from deepmd.pt.model.task.fitting import ( + Fitting, ) from deepmd.pt.utils import ( env, ) from deepmd.pt.utils.env import ( DEFAULT_PRECISION, - DEVICE, PRECISION_DICT, ) from deepmd.pt.utils.exclude_mask import ( @@ -44,24 +41,6 @@ to_numpy_array, to_torch_tensor, ) -from deepmd.utils.data_system import ( - DeepmdDataSystem, -) -from deepmd.utils.finetune import ( - change_energy_bias_lower, -) -from deepmd.pt.model.task.fitting import ( - Fitting, -) -from deepmd.dpmodel import ( - FittingOutputDef, - OutputVariableDef, - fitting_check_output, -) -from deepmd.pt.model.network.network import ( - MaskLMHead, - NonLinearHead, -) from deepmd.utils.path import ( DPPath, ) @@ -71,6 +50,7 @@ log = logging.getLogger(__name__) + @Fitting.register("denoise") class DenoiseFittingNet(Fitting): """Construct a denoise fitting net. @@ -139,7 +119,7 @@ def __init__( **kwargs, ): super().__init__() - self.var_name = ["updated_coord","logits"] + self.var_name = ["updated_coord", "logits"] self.ntypes = ntypes self.dim_descrpt = dim_descrpt self.embedding_width = embedding_width @@ -364,7 +344,7 @@ def __getitem__(self, key): def _net_out_dim(self): """Set the FittingNet output dim.""" return [3, self.ntypes] - #pass + # pass def output_def(self): return FittingOutputDef( @@ -482,17 +462,17 @@ def forward( device=descriptor.device, ) if self.mixed_types: - atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum(dim=-2) / (sw.sum(dim=-1).unsqueeze(-1)+1e-6) + atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( + dim=-2 + ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) atom_logits = self.filter_layers_logits[0](xx) - #Is xx_zeros useful in denoise task?????????????? - #if xx_zeros is not None: + # Is xx_zeros useful in denoise task?????????????? + # if xx_zeros is not None: # atom_property -= self.filter_layers.networks[0](xx_zeros) outs_coord = ( outs_coord + atom_updated_coord ) # Shape is [nframes, natoms[0], net_dim_out] - outs_logits = ( - outs_logits + atom_logits - ) + outs_logits = outs_logits + atom_logits # TODO: else: for type_i, ll in enumerate(self.filter_layers_coord.networks): @@ -517,8 +497,10 @@ def forward( # nf x nloc x nod outs_coord = outs_coord * mask[:, :, None] outs_logits = outs_logits * mask[:, :, None] - return {self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), - self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION)} + return { + self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), + self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), + } def compute_output_stats( self, @@ -543,6 +525,7 @@ def compute_output_stats( """ pass + @fitting_check_output class DenoiseNet(Fitting): def __init__( diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 8524974058..9f3a7a2373 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -3,12 +3,11 @@ import logging from typing import ( Callable, - Union, List, Optional, + Union, ) -import numpy as np import torch from deepmd.dpmodel import ( @@ -28,12 +27,12 @@ from deepmd.pt.utils.env import ( DEFAULT_PRECISION, ) -from deepmd.utils.version import ( - check_version_compatibility, -) from deepmd.utils.path import ( DPPath, ) +from deepmd.utils.version import ( + check_version_compatibility, +) dtype = env.GLOBAL_PT_FLOAT_PRECISION device = env.DEVICE @@ -85,12 +84,8 @@ def deserialize(cls, data: dict) -> "GeneralFitting": def serialize(self) -> dict: """Serialize the fitting to dict.""" - return { - **super().serialize(), - "type": "property", - "task_num": self.task_num - } - + return {**super().serialize(), "type": "property", "task_num": self.task_num} + def output_def(self) -> FittingOutputDef: return FittingOutputDef( [ @@ -128,4 +123,4 @@ def compute_output_stats( pass # make jit happy with torch 2.0.0 - exclude_types: List[int] \ No newline at end of file + exclude_types: List[int] diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 4480800ec8..fda9196df1 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -27,8 +27,8 @@ DenoiseLoss, EnergySpinLoss, EnergyStdLoss, - TensorLoss, PropertyLoss, + TensorLoss, ) from deepmd.pt.model.model import ( DPZBLModel, diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index b3124cb859..84f33f3a07 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1175,6 +1175,7 @@ def fitting_dipole(): Argument("seed", [int, None], optional=True, doc=doc_seed), ] + @fitting_args_plugin.register("property") def fitting_property(): doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." @@ -1213,6 +1214,7 @@ def fitting_property(): Argument("task_num", int, optional=True, default=1, doc=doc_task_num), ] + @fitting_args_plugin.register("denoise") def fitting_denoise(): doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." @@ -1249,6 +1251,7 @@ def fitting_denoise(): Argument("seed", [int, None], optional=True, doc=doc_seed), ] + # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ @@ -1991,6 +1994,7 @@ def loss_tensor(): ), ] + @loss_args_plugin.register("property") def loss_property(): doc_loss_func = "The loss function, such as 'mae','smooth_mae'." @@ -2000,22 +2004,43 @@ def loss_property(): doc_beta = "The 'beta' parameter in 'smooth_mae' loss." return [ Argument( - "loss_func", str, optional=True, default="smooth_mae", doc=doc_loss_func, + "loss_func", + str, + optional=True, + default="smooth_mae", + doc=doc_loss_func, ), Argument( - "metric", list, optional=True, default=["mae"], doc=doc_metric, + "metric", + list, + optional=True, + default=["mae"], + doc=doc_metric, ), Argument( - "mean", [float, int, list], optional=True, default=0, doc=doc_mean, + "mean", + [float, int, list], + optional=True, + default=0, + doc=doc_mean, ), Argument( - "std", [float, int, list], optional=True, default=1, doc=doc_std, + "std", + [float, int, list], + optional=True, + default=1, + doc=doc_std, ), Argument( - "beta", [float, int], optional=True, default=1.00, doc=doc_beta, + "beta", + [float, int], + optional=True, + default=1.00, + doc=doc_beta, ), ] + def loss_variant_type_args(): doc_loss = "The type of the loss. When the fitting type is `ener`, the loss type should be set to `ener` or left unset. When the fitting type is `dipole` or `polar`, the loss type should be set to `tensor`." diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index 64d1d2c4f5..00800a6251 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -145,6 +145,7 @@ def test_jit( ).to(env.DEVICE) torch.jit.script(ft0) + class TestInvariance(unittest.TestCase): def setUp(self) -> None: self.natoms = 5 @@ -229,7 +230,7 @@ def test_permu(self): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num = 8, + task_num=8, numb_fparam=0, numb_aparam=0, mixed_types=True, @@ -274,7 +275,7 @@ def test_trans(self): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num = 11, + task_num=11, numb_fparam=0, numb_aparam=0, mixed_types=True, @@ -320,7 +321,7 @@ def setUp(self): self.ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num = 3, + task_num=3, numb_fparam=0, numb_aparam=0, mixed_types=True, @@ -345,4 +346,4 @@ def tearDown(self) -> None: if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() From cea7476a15cb44e3d3014065fa08849e90002019 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 16:11:25 +0000 Subject: [PATCH 05/31] fix pre-commit --- deepmd/pt/model/task/denoise.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 26fb4742c5..58b8b71689 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -293,7 +293,7 @@ def serialize(self) -> dict: } @classmethod - def deserialize(cls, data: dict) -> "DenoiseFitting": + def deserialize(cls, data: dict) -> "DenoiseFittingNet": data = copy.deepcopy(data) variables = data.pop("@variables") nets = data.pop("nets") @@ -494,6 +494,7 @@ def forward( outs_logits + atom_logits ) # TODO: + ''' else: for type_i, ll in enumerate(self.filter_layers_coord.networks): mask = (atype == type_i).unsqueeze(-1) @@ -512,6 +513,7 @@ def forward( outs = ( outs + atom_property ) # Shape is [nframes, natoms[0], net_dim_out] + ''' # nf x nloc mask = self.emask(atype) # nf x nloc x nod From 49655298c2ab646cd374e40ed04b4bec749fe700 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 Mar 2024 16:14:35 +0000 Subject: [PATCH 06/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/denoise.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index da1f78b4db..48d6893f8b 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -474,7 +474,7 @@ def forward( ) # Shape is [nframes, natoms[0], net_dim_out] outs_logits = outs_logits + atom_logits # TODO: - ''' + """ else: for type_i, ll in enumerate(self.filter_layers_coord.networks): mask = (atype == type_i).unsqueeze(-1) @@ -493,7 +493,7 @@ def forward( outs = ( outs + atom_property ) # Shape is [nframes, natoms[0], net_dim_out] - ''' + """ # nf x nloc mask = self.emask(atype) # nf x nloc x nod From a05caced5d5c9a04011c84c9e84d19b3bed08651 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 17:19:23 +0000 Subject: [PATCH 07/31] Add example --- checkpoint | 1 + deepmd/pt/model/task/dipole.py | 1 + deepmd/pt/model/task/ener.py | 1 + deepmd/pt/model/task/invar_fitting.py | 2 +- deepmd/pt/model/task/polarizability.py | 1 + examples/property/data/3/nopbc | 0 examples/property/data/3/set.000000/coord.npy | Bin 0 -> 272 bytes .../property/data/3/set.000000/property.npy | Bin 0 -> 176 bytes .../data/3/set.000000/real_atom_types.npy | Bin 0 -> 176 bytes examples/property/data/3/type.raw | 3 + examples/property/data/3/type_map.raw | 118 ++++++++++ examples/property/data/5/nopbc | 0 examples/property/data/5/set.000000/coord.npy | Bin 0 -> 608 bytes .../property/data/5/set.000000/property.npy | Bin 0 -> 224 bytes .../data/5/set.000000/real_atom_types.npy | Bin 0 -> 288 bytes examples/property/data/5/type.raw | 5 + examples/property/data/5/type_map.raw | 118 ++++++++++ examples/property/data/6/nopbc | 0 examples/property/data/6/set.000000/coord.npy | Bin 0 -> 1568 bytes .../property/data/6/set.000000/property.npy | Bin 0 -> 368 bytes .../data/6/set.000000/real_atom_types.npy | Bin 0 -> 608 bytes examples/property/data/6/type.raw | 6 + examples/property/data/6/type_map.raw | 118 ++++++++++ examples/property/train/input_torch.json | 212 ++++++++++++++++++ source/tests/tf/input.json | 83 +++++++ source/tests/tf/input_se_a_ebd_v2.json | 66 ++++++ 26 files changed, 734 insertions(+), 1 deletion(-) create mode 100644 checkpoint create mode 100644 examples/property/data/3/nopbc create mode 100644 examples/property/data/3/set.000000/coord.npy create mode 100644 examples/property/data/3/set.000000/property.npy create mode 100644 examples/property/data/3/set.000000/real_atom_types.npy create mode 100644 examples/property/data/3/type.raw create mode 100644 examples/property/data/3/type_map.raw create mode 100644 examples/property/data/5/nopbc create mode 100644 examples/property/data/5/set.000000/coord.npy create mode 100644 examples/property/data/5/set.000000/property.npy create mode 100644 examples/property/data/5/set.000000/real_atom_types.npy create mode 100644 examples/property/data/5/type.raw create mode 100644 examples/property/data/5/type_map.raw create mode 100644 examples/property/data/6/nopbc create mode 100644 examples/property/data/6/set.000000/coord.npy create mode 100644 examples/property/data/6/set.000000/property.npy create mode 100644 examples/property/data/6/set.000000/real_atom_types.npy create mode 100644 examples/property/data/6/type.raw create mode 100644 examples/property/data/6/type_map.raw create mode 100644 examples/property/train/input_torch.json create mode 100644 source/tests/tf/input.json create mode 100644 source/tests/tf/input_se_a_ebd_v2.json diff --git a/checkpoint b/checkpoint new file mode 100644 index 0000000000..d6c7373516 --- /dev/null +++ b/checkpoint @@ -0,0 +1 @@ +model.ckpt-1.pt \ No newline at end of file diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index ca445c8588..ca15ae0c20 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -176,6 +176,7 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index 12c0917dd2..c8e2a77431 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -189,6 +189,7 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index b6a4e0abf1..0f9028f182 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -189,9 +189,9 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, ): """Based on embedding net output, calculate total energy. diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index d7428c4d53..587bce2709 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -278,6 +278,7 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): diff --git a/examples/property/data/3/nopbc b/examples/property/data/3/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/property/data/3/set.000000/coord.npy b/examples/property/data/3/set.000000/coord.npy new file mode 100644 index 0000000000000000000000000000000000000000..324d8f63ba7f9beaec0c59e6af810aebb2fdc60d GIT binary patch literal 272 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$7ItrGWItsN4WCJb+28Mv~NdhG$jYDJ>} literal 0 HcmV?d00001 diff --git a/examples/property/data/3/set.000000/property.npy b/examples/property/data/3/set.000000/property.npy new file mode 100644 index 0000000000000000000000000000000000000000..fcd3092954e321c710d0d2b29304735321f61d60 GIT binary patch literal 176 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$7Its>`ItsN4WCJb+1_lS4y_fa_=>tL2HrfN}2N#sCgJ^}*0@p#}4xjrb MgZKwxy>Hk90P`0sT>t<8 literal 0 HcmV?d00001 diff --git a/examples/property/data/3/set.000000/real_atom_types.npy b/examples/property/data/3/set.000000/real_atom_types.npy new file mode 100644 index 0000000000000000000000000000000000000000..fb312ec6b8ba7ada7b8796463ff5cfdfcbc4e6c9 GIT binary patch literal 176 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlWC!@qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= hXCxM+0{I$7Its>`ItsN4WCJdC2JC literal 0 HcmV?d00001 diff --git a/examples/property/data/3/type.raw b/examples/property/data/3/type.raw new file mode 100644 index 0000000000..bb0b1cf658 --- /dev/null +++ b/examples/property/data/3/type.raw @@ -0,0 +1,3 @@ +0 +0 +0 diff --git a/examples/property/data/3/type_map.raw b/examples/property/data/3/type_map.raw new file mode 100644 index 0000000000..e37cba22a9 --- /dev/null +++ b/examples/property/data/3/type_map.raw @@ -0,0 +1,118 @@ +H +He +Li +Be +B +C +N +O +F +Ne +Na +Mg +Al +Si +P +S +Cl +Ar +K +Ca +Sc +Ti +V +Cr +Mn +Fe +Co +Ni +Cu +Zn +Ga +Ge +As +Se +Br +Kr +Rb +Sr +Y +Zr +Nb +Mo +Tc +Ru +Rh +Pd +Ag +Cd +In +Sn +Sb +Te +I +Xe +Cs +Ba +La +Ce +Pr +Nd +Pm +Sm +Eu +Gd +Tb +Dy +Ho +Er +Tm +Yb +Lu +Hf +Ta +W +Re +Os +Ir +Pt +Au +Hg +Tl +Pb +Bi +Po +At +Rn +Fr +Ra +Ac +Th +Pa +U +Np +Pu +Am +Cm +Bk +Cf +Es +Fm +Md +No +Lr +Rf +Db +Sg +Bh +Hs +Mt +Ds +Rg +Cn +Nh +Fl +Mc +Lv +Ts +Og diff --git a/examples/property/data/5/nopbc b/examples/property/data/5/nopbc new file mode 100644 index 0000000000..e69de29bb2 diff --git a/examples/property/data/5/set.000000/coord.npy b/examples/property/data/5/set.000000/coord.npy new file mode 100644 index 0000000000000000000000000000000000000000..4aa899142cfa21c278b4611021f0805c0713c9e2 GIT binary patch literal 608 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I#yItqrSnmP)#3giMV1_lO&Ql_r`K>9&`{6~8b?VHp9q8awECV}V!C6g<_ z^1ew4Ab!Fd{SP3zVQ2Pt5bdC~(ilV^u=abnA4GrE0E;(B^S%M`9ex~o0HPBXGlInz zSTVi@sRx=L1fmn3Uf2Yp9UfMtfM}pO6ZeB?$-r+QKEsT=&0z6q6F!6JhB=G=f@rwC zpZ`AD4-$VJ2oW#(`2oaFFbV~WGh`@%#TlOWfcILek1_)&w^*N5c?1kVDW@T zmd_ybkN{Zxz>X^aj)`oOdaV0(b!P`n>xj{TCiAbNrEqCX&Y J!0?@94**dOpQQi* literal 0 HcmV?d00001 diff --git a/examples/property/data/5/set.000000/property.npy b/examples/property/data/5/set.000000/property.npy new file mode 100644 index 0000000000000000000000000000000000000000..4b9d4f94532c868ad57c10f624bf1008c156a380 GIT binary patch literal 224 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I#yIts>`ItsN4WCJb+28IQ~7jEna(hM%Y_Syrfgdp_?Ao_q_-ffWhg4st- ufcOq}FCK$vg}_^v!Q!T8SA*z;RS6eB;tO<~uYmXmzH00Ni{C3eWe)&`t2|Z! literal 0 HcmV?d00001 diff --git a/examples/property/data/5/set.000000/real_atom_types.npy b/examples/property/data/5/set.000000/real_atom_types.npy new file mode 100644 index 0000000000000000000000000000000000000000..497a9722c47eddfbd6d933c3dff48939195e652d GIT binary patch literal 288 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlWC!@qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I#yItr$mItsN4WCJc%2I2q*R2fV?jE_#kYB5 z@0r?f^79rfO%V}MKwX8Av#AwmdeNAl)Duu@7($wXie_ffx5~#v5582 z6S}F7X?+CUjb5x_{nYBGtiPD}j`c`j9Z=VQ_6&7rPwrcAySlKJdefg>(5+zb0qYm{ z+t8iMS5^Y;L}Md#J2f68udMhPT<9r@K3)aXIbS{pTEnLF0?j^Wvz{zUQHI;cP7WgK z`ihw0a`@0BP~85N`*f_ud*GLwT2fqt>ROiPudF!4K5laxu>WHloqhTh^q?=^eM}#y zkG>Hd?Dp#c3SAALt0hHe(cdmOevJ4&pVW8+-0eKVe5^=d3)XY~IOXLt=PVC+4oQx>HSBKA6;(c`}pXX(bzVHS`$M|?&Zks=ey4acIc^Us$_(WF+_i5Kgy>j`(sX^TO$g*d#*-4p|3TYIl9e9>32&T?xCMD zxdX-EV&*TK+G5~lZu9Y;)BbguGN|e@+QD^T*8$Y=e*n}oH@t){%v=G@{1yQ*_ds*c zONfcjIgjd}KqBVp6+oGl#@wX;FE6*2jivCrp;caPdE347ir!~#T0CHG&cDm3>v*!$ F{~uq!r_%rc literal 0 HcmV?d00001 diff --git a/examples/property/data/6/set.000000/property.npy b/examples/property/data/6/set.000000/property.npy new file mode 100644 index 0000000000000000000000000000000000000000..66162f99db87c32f1c573d91bcdc2f00279a0999 GIT binary patch literal 368 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlV+i=qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$-209AHnmP)#3giMV1_p+PVz&$Xf%Jj>Z|2#9`1YokK=gqXORj;$7sUKK z43@X{I|&kBV7v&T{^851An^xh^%d z8*-i>28kzd{=5R>Go+l^3!)!9;64o!f3UUX28izv@Mi{C{KJ=PVDUA-Aoj8fZ2|Ee r^!{H3(Fy7cu7cDl=-=G}rY+JhfW#l{c?>pxfuF--5P!ij@9XveHaBu^ literal 0 HcmV?d00001 diff --git a/examples/property/data/6/set.000000/real_atom_types.npy b/examples/property/data/6/set.000000/real_atom_types.npy new file mode 100644 index 0000000000000000000000000000000000000000..0b0f7e0d55991e977e80b8efb12092c3cca424d4 GIT binary patch literal 608 zcmbR27wQ`j$;eQ~P_3SlTAW;@Zl$1ZlWC!@qoAIaUsO_*m=~X4l#&V(cT3DEP6dh= zXCxM+0{I$-20996nmP)#3giMVHUS5;KqG9fa(Hu|%VCrFf7!4DLsfY1#*^llXm^yU*Fg}ci>4TXMqtWGI>R>bg D?Qk4@ literal 0 HcmV?d00001 diff --git a/examples/property/data/6/type.raw b/examples/property/data/6/type.raw new file mode 100644 index 0000000000..f7eb44d66e --- /dev/null +++ b/examples/property/data/6/type.raw @@ -0,0 +1,6 @@ +0 +0 +0 +0 +0 +0 diff --git a/examples/property/data/6/type_map.raw b/examples/property/data/6/type_map.raw new file mode 100644 index 0000000000..e37cba22a9 --- /dev/null +++ b/examples/property/data/6/type_map.raw @@ -0,0 +1,118 @@ +H +He +Li +Be +B +C +N +O +F +Ne +Na +Mg +Al +Si +P +S +Cl +Ar +K +Ca +Sc +Ti +V +Cr +Mn +Fe +Co +Ni +Cu +Zn +Ga +Ge +As +Se +Br +Kr +Rb +Sr +Y +Zr +Nb +Mo +Tc +Ru +Rh +Pd +Ag +Cd +In +Sn +Sb +Te +I +Xe +Cs +Ba +La +Ce +Pr +Nd +Pm +Sm +Eu +Gd +Tb +Dy +Ho +Er +Tm +Yb +Lu +Hf +Ta +W +Re +Os +Ir +Pt +Au +Hg +Tl +Pb +Bi +Po +At +Rn +Fr +Ra +Ac +Th +Pa +U +Np +Pu +Am +Cm +Bk +Cf +Es +Fm +Md +No +Lr +Rf +Db +Sg +Bh +Hs +Mt +Ds +Rg +Cn +Nh +Fl +Mc +Lv +Ts +Og diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json new file mode 100644 index 0000000000..ed1afc3a7c --- /dev/null +++ b/examples/property/train/input_torch.json @@ -0,0 +1,212 @@ +{ + "_comment": "that's all", + "model": { + "type_map": [ + "H", + "He", + "Li", + "Be", + "B", + "C", + "N", + "O", + "F", + "Ne", + "Na", + "Mg", + "Al", + "Si", + "P", + "S", + "Cl", + "Ar", + "K", + "Ca", + "Sc", + "Ti", + "V", + "Cr", + "Mn", + "Fe", + "Co", + "Ni", + "Cu", + "Zn", + "Ga", + "Ge", + "As", + "Se", + "Br", + "Kr", + "Rb", + "Sr", + "Y", + "Zr", + "Nb", + "Mo", + "Tc", + "Ru", + "Rh", + "Pd", + "Ag", + "Cd", + "In", + "Sn", + "Sb", + "Te", + "I", + "Xe", + "Cs", + "Ba", + "La", + "Ce", + "Pr", + "Nd", + "Pm", + "Sm", + "Eu", + "Gd", + "Tb", + "Dy", + "Ho", + "Er", + "Tm", + "Yb", + "Lu", + "Hf", + "Ta", + "W", + "Re", + "Os", + "Ir", + "Pt", + "Au", + "Hg", + "Tl", + "Pb", + "Bi", + "Po", + "At", + "Rn", + "Fr", + "Ra", + "Ac", + "Th", + "Pa", + "U", + "Np", + "Pu", + "Am", + "Cm", + "Bk", + "Cf", + "Es", + "Fm", + "Md", + "No", + "Lr", + "Rf", + "Db", + "Sg", + "Bh", + "Hs", + "Mt", + "Ds", + "Rg", + "Cn", + "Nh", + "Fl", + "Mc", + "Lv", + "Ts", + "Og" + ], + "descriptor": { + "type": "dpa2", + "tebd_dim": 8, + "repinit_rcut": 9.0, + "repinit_rcut_smth": 8.0, + "repinit_nsel": 120, + "repformer_rcut": 4.0, + "repformer_rcut_smth": 3.5, + "repformer_nsel": 40, + "repinit_neuron": [ + 25, + 50, + 100 + ], + "repinit_axis_neuron": 12, + "repinit_activation": "tanh", + "repformer_nlayers": 12, + "repformer_g1_dim": 128, + "repformer_g2_dim": 32, + "repformer_attn2_hidden": 32, + "repformer_attn2_nhead": 4, + "repformer_attn1_hidden": 128, + "repformer_attn1_nhead": 4, + "repformer_axis_dim": 4, + "repformer_update_h2": false, + "repformer_update_g1_has_conv": true, + "repformer_update_g1_has_grrg": true, + "repformer_update_g1_has_drrd": true, + "repformer_update_g1_has_attn": true, + "repformer_update_g2_has_g1g1": true, + "repformer_update_g2_has_attn": true, + "repformer_attn2_has_gate": true, + "repformer_add_type_ebd_to_seq": false + }, + "fitting_net": { + "type": "property", + "task_num": 3, + "neuron": [ + 240, + 240, + 240 + ], + "resnet_dt": true, + "seed": 1, + "_comment": " that's all" + }, + "_comment": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.0002, + "stop_lr": 3.51e-08, + "_comment": "that's all" + }, + "loss": { + "type": "property", + "loss_func": "smooth_mae", + "metric": ["mae"], + "beta": 1.00, + "_comment": " that's all" + }, + "training": { + "stat_file": "./dpa2", + "training_data": { + "systems": [ + "../data/3", + "../data/6" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "validation_data": { + "systems": [ + "../data/5" + ], + "batch_size": 1, + "_comment": "that's all" + }, + "numb_steps": 1000000, + "warmup_steps": 0, + "gradient_max_norm": 5.0, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 100, + "save_freq": 2000, + "_comment": "that's all" + } +} diff --git a/source/tests/tf/input.json b/source/tests/tf/input.json new file mode 100644 index 0000000000..d976ef43ff --- /dev/null +++ b/source/tests/tf/input.json @@ -0,0 +1,83 @@ +{ + "_comment1": " model parameters", + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_e2_a", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.5, + "rcut": 6.0, + "_comment": "N2=2N1, N2=N1, and otherwise can be tested", + "neuron": [ + 4, + 8, + 17, + 17 + ], + "resnet_dt": false, + "axis_neuron": 16, + "seed": 1, + "_comment2": " that's all", + "exclude_types": [ + [ + 0, + 1 + ] + ] + }, + "fitting_net": { + "neuron": [ + 20, + 20, + 20 + ], + "resnet_dt": true, + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 1, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} \ No newline at end of file diff --git a/source/tests/tf/input_se_a_ebd_v2.json b/source/tests/tf/input_se_a_ebd_v2.json new file mode 100644 index 0000000000..0ff86b8cc0 --- /dev/null +++ b/source/tests/tf/input_se_a_ebd_v2.json @@ -0,0 +1,66 @@ +{ + "_comment1": " model parameters", + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_a_ebd_v2", + "type_one_side": false + }, + "fitting_net": { + "neuron": [ + 20, + 20, + 20 + ], + "resnet_dt": true, + "seed": 1, + "_comment3": " that's all" + }, + "_comment4": " that's all", + "type_embedding": { + "neuron": [ + 1 + ] + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": "auto", + "_comment7": "that's all" + }, + "validation_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" + }, + "numb_steps": 1, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 1, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} \ No newline at end of file From 3050172fbb1daab919653e7bdb0fec0d47601c4f Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 18 Mar 2024 17:20:43 +0000 Subject: [PATCH 08/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- checkpoint | 2 +- examples/property/train/input_torch.json | 4 +- source/tests/tf/input.json | 156 +++++++++++------------ source/tests/tf/input_se_a_ebd_v2.json | 122 +++++++++--------- 4 files changed, 143 insertions(+), 141 deletions(-) diff --git a/checkpoint b/checkpoint index d6c7373516..aa6f74047b 100644 --- a/checkpoint +++ b/checkpoint @@ -1 +1 @@ -model.ckpt-1.pt \ No newline at end of file +model.ckpt-1.pt diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json index ed1afc3a7c..915edb46a1 100644 --- a/examples/property/train/input_torch.json +++ b/examples/property/train/input_torch.json @@ -179,7 +179,9 @@ "loss": { "type": "property", "loss_func": "smooth_mae", - "metric": ["mae"], + "metric": [ + "mae" + ], "beta": 1.00, "_comment": " that's all" }, diff --git a/source/tests/tf/input.json b/source/tests/tf/input.json index d976ef43ff..f39227f7ce 100644 --- a/source/tests/tf/input.json +++ b/source/tests/tf/input.json @@ -1,83 +1,83 @@ { - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_e2_a", - "sel": [ - 46, - 92 - ], - "rcut_smth": 0.5, - "rcut": 6.0, - "_comment": "N2=2N1, N2=N1, and otherwise can be tested", - "neuron": [ - 4, - 8, - 17, - 17 - ], - "resnet_dt": false, - "axis_neuron": 16, - "seed": 1, - "_comment2": " that's all", - "exclude_types": [ - [ - 0, - 1 - ] - ] - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all" + "_comment1": " model parameters", + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_e2_a", + "sel": [ + 46, + 92 + ], + "rcut_smth": 0.5, + "rcut": 6.0, + "_comment": "N2=2N1, N2=N1, and otherwise can be tested", + "neuron": [ + 4, + 8, + 17, + 17 + ], + "resnet_dt": false, + "axis_neuron": 16, + "seed": 1, + "_comment2": " that's all", + "exclude_types": [ + [ + 0, + 1 + ] + ] }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" + "fitting_net": { + "neuron": [ + 20, + 20, + 20 + ], + "resnet_dt": true, + "seed": 1, + "_comment3": " that's all" }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" + "_comment4": " that's all" + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": "auto", + "_comment7": "that's all" }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" + "validation_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" }, - "_comment10": "that's all" -} \ No newline at end of file + "numb_steps": 1, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 1, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} diff --git a/source/tests/tf/input_se_a_ebd_v2.json b/source/tests/tf/input_se_a_ebd_v2.json index 0ff86b8cc0..530d996383 100644 --- a/source/tests/tf/input_se_a_ebd_v2.json +++ b/source/tests/tf/input_se_a_ebd_v2.json @@ -1,66 +1,66 @@ { - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_a_ebd_v2", - "type_one_side": false - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all", - "type_embedding": { - "neuron": [ - 1 - ] - } + "_comment1": " model parameters", + "model": { + "type_map": [ + "O", + "H" + ], + "descriptor": { + "type": "se_a_ebd_v2", + "type_one_side": false }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" + "fitting_net": { + "neuron": [ + 20, + 20, + 20 + ], + "resnet_dt": true, + "seed": 1, + "_comment3": " that's all" }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" + "_comment4": " that's all", + "type_embedding": { + "neuron": [ + 1 + ] + } + }, + "learning_rate": { + "type": "exp", + "decay_steps": 5000, + "start_lr": 0.001, + "stop_lr": 3.51e-08, + "_comment5": "that's all" + }, + "loss": { + "type": "ener", + "start_pref_e": 0.02, + "limit_pref_e": 1, + "start_pref_f": 1000, + "limit_pref_f": 1, + "start_pref_v": 0, + "limit_pref_v": 0, + "_comment6": " that's all" + }, + "training": { + "training_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": "auto", + "_comment7": "that's all" }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" + "validation_data": { + "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", + "batch_size": 1, + "numb_btch": 3, + "_comment8": "that's all" }, - "_comment10": "that's all" -} \ No newline at end of file + "numb_steps": 1, + "seed": 10, + "disp_file": "lcurve.out", + "disp_freq": 1, + "save_freq": 1, + "_comment9": "that's all" + }, + "_comment10": "that's all" +} From 20cee41a9402355fa05fc92023bc49e7397f0f51 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 17:21:20 +0000 Subject: [PATCH 09/31] delete input --- source/tests/tf/input.json | 83 -------------------------- source/tests/tf/input_se_a_ebd_v2.json | 66 -------------------- 2 files changed, 149 deletions(-) delete mode 100644 source/tests/tf/input.json delete mode 100644 source/tests/tf/input_se_a_ebd_v2.json diff --git a/source/tests/tf/input.json b/source/tests/tf/input.json deleted file mode 100644 index d976ef43ff..0000000000 --- a/source/tests/tf/input.json +++ /dev/null @@ -1,83 +0,0 @@ -{ - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_e2_a", - "sel": [ - 46, - 92 - ], - "rcut_smth": 0.5, - "rcut": 6.0, - "_comment": "N2=2N1, N2=N1, and otherwise can be tested", - "neuron": [ - 4, - 8, - 17, - 17 - ], - "resnet_dt": false, - "axis_neuron": 16, - "seed": 1, - "_comment2": " that's all", - "exclude_types": [ - [ - 0, - 1 - ] - ] - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all" - }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" - }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" - }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" - }, - "_comment10": "that's all" -} \ No newline at end of file diff --git a/source/tests/tf/input_se_a_ebd_v2.json b/source/tests/tf/input_se_a_ebd_v2.json deleted file mode 100644 index 0ff86b8cc0..0000000000 --- a/source/tests/tf/input_se_a_ebd_v2.json +++ /dev/null @@ -1,66 +0,0 @@ -{ - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_a_ebd_v2", - "type_one_side": false - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all", - "type_embedding": { - "neuron": [ - 1 - ] - } - }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" - }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" - }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" - }, - "_comment10": "that's all" -} \ No newline at end of file From 501b46da4bb77f55d0da86e5a8c1ef778ed8d042 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Mon, 18 Mar 2024 17:24:15 +0000 Subject: [PATCH 10/31] fix push bug --- checkpoint | 1 - source/tests/tf/input.json | 83 -------------------------- source/tests/tf/input_se_a_ebd_v2.json | 66 -------------------- 3 files changed, 150 deletions(-) delete mode 100644 checkpoint delete mode 100644 source/tests/tf/input.json delete mode 100644 source/tests/tf/input_se_a_ebd_v2.json diff --git a/checkpoint b/checkpoint deleted file mode 100644 index aa6f74047b..0000000000 --- a/checkpoint +++ /dev/null @@ -1 +0,0 @@ -model.ckpt-1.pt diff --git a/source/tests/tf/input.json b/source/tests/tf/input.json deleted file mode 100644 index f39227f7ce..0000000000 --- a/source/tests/tf/input.json +++ /dev/null @@ -1,83 +0,0 @@ -{ - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_e2_a", - "sel": [ - 46, - 92 - ], - "rcut_smth": 0.5, - "rcut": 6.0, - "_comment": "N2=2N1, N2=N1, and otherwise can be tested", - "neuron": [ - 4, - 8, - 17, - 17 - ], - "resnet_dt": false, - "axis_neuron": 16, - "seed": 1, - "_comment2": " that's all", - "exclude_types": [ - [ - 0, - 1 - ] - ] - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all" - }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" - }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" - }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" - }, - "_comment10": "that's all" -} diff --git a/source/tests/tf/input_se_a_ebd_v2.json b/source/tests/tf/input_se_a_ebd_v2.json deleted file mode 100644 index 530d996383..0000000000 --- a/source/tests/tf/input_se_a_ebd_v2.json +++ /dev/null @@ -1,66 +0,0 @@ -{ - "_comment1": " model parameters", - "model": { - "type_map": [ - "O", - "H" - ], - "descriptor": { - "type": "se_a_ebd_v2", - "type_one_side": false - }, - "fitting_net": { - "neuron": [ - 20, - 20, - 20 - ], - "resnet_dt": true, - "seed": 1, - "_comment3": " that's all" - }, - "_comment4": " that's all", - "type_embedding": { - "neuron": [ - 1 - ] - } - }, - "learning_rate": { - "type": "exp", - "decay_steps": 5000, - "start_lr": 0.001, - "stop_lr": 3.51e-08, - "_comment5": "that's all" - }, - "loss": { - "type": "ener", - "start_pref_e": 0.02, - "limit_pref_e": 1, - "start_pref_f": 1000, - "limit_pref_f": 1, - "start_pref_v": 0, - "limit_pref_v": 0, - "_comment6": " that's all" - }, - "training": { - "training_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": "auto", - "_comment7": "that's all" - }, - "validation_data": { - "systems": "/home/data/zcq/deepmd-kit/source/tests/tf/model_compression/data", - "batch_size": 1, - "numb_btch": 3, - "_comment8": "that's all" - }, - "numb_steps": 1, - "seed": 10, - "disp_file": "lcurve.out", - "disp_freq": 1, - "save_freq": 1, - "_comment9": "that's all" - }, - "_comment10": "that's all" -} From c8dad8d4f8f854f0990787b8d76966a49bb5118a Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:23:54 +0000 Subject: [PATCH 11/31] 3.19 update --- deepmd/infer/deep_property.py | 5 +- deepmd/pt/loss/property.py | 19 +- .../pt/model/atomic_model/dp_atomic_model.py | 1 - deepmd/pt/model/model/__init__.py | 2 +- deepmd/pt/model/model/denoise_model.py | 81 --- deepmd/pt/model/model/dp_model.py | 8 - deepmd/pt/model/model/property_model.py | 22 +- deepmd/pt/model/task/denoise.py | 643 ------------------ deepmd/pt/model/task/dipole.py | 1 - deepmd/pt/model/task/ener.py | 1 - deepmd/pt/model/task/invar_fitting.py | 1 - deepmd/pt/model/task/polarizability.py | 1 - deepmd/utils/argcheck.py | 10 +- 13 files changed, 22 insertions(+), 773 deletions(-) delete mode 100644 deepmd/pt/model/model/denoise_model.py delete mode 100644 deepmd/pt/model/task/denoise.py diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index f53d7acd25..e3b1866eb5 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -72,8 +72,7 @@ def eval( mixed_type: bool = False, **kwargs: Dict[str, Any], ) -> Tuple[np.ndarray, ...]: - """Evaluate energy, force, and virial. If atomic is True, - also return atomic energy and atomic virial. + """Evaluate properties. If atomic is True, also return atomic property. Parameters ---------- @@ -86,7 +85,7 @@ def eval( The types of the atoms. If mixed_type is False, the shape is (natoms,); otherwise, the shape is (nframes, natoms). atomic : bool, optional - Whether to return atomic energy and atomic virial, by default False. + Whether to return atomic property, by default False. fparam : np.ndarray, optional The frame parameters, by default None. aparam : np.ndarray, optional diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 5aaaf0fcf4..aed50eb28d 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -50,16 +50,23 @@ def __init__( self.beta = kwargs.get("beta", 1.00) def forward(self, model_pred, label, natoms, learning_rate, mae=False): - """Return loss on loss and force. + """Return loss on properties . - Args: - - model_pred: Property prediction. - - label: Target property. - - natoms: Tell atom count. + Parameters: + ---------- + model_pred : dict[str, torch.Tensor] + Model predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. Returns ------- - - loss: Loss to minimize. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + Other losses for display. """ assert label["property"].shape[-1] == self.task_num assert model_pred["property"].shape[-1] == self.task_num diff --git a/deepmd/pt/model/atomic_model/dp_atomic_model.py b/deepmd/pt/model/atomic_model/dp_atomic_model.py index fd8f71fa2d..6aa8df7aee 100644 --- a/deepmd/pt/model/atomic_model/dp_atomic_model.py +++ b/deepmd/pt/model/atomic_model/dp_atomic_model.py @@ -176,7 +176,6 @@ def forward_atomic( gr=rot_mat, g2=g2, h2=h2, - sw=sw, fparam=fparam, aparam=aparam, ) diff --git a/deepmd/pt/model/model/__init__.py b/deepmd/pt/model/model/__init__.py index efbef794c0..7a2070e476 100644 --- a/deepmd/pt/model/model/__init__.py +++ b/deepmd/pt/model/model/__init__.py @@ -140,7 +140,7 @@ def get_standard_model(model_params): fitting_net["type"] = fitting_net.get("type", "ener") fitting_net["ntypes"] = descriptor.get_ntypes() fitting_net["mixed_types"] = descriptor.mixed_types() - if fitting_net["type"] in ["dipole", "polar", "denoise"]: + if fitting_net["type"] in ["dipole", "polar"]: fitting_net["embedding_width"] = descriptor.get_dim_emb() fitting_net["dim_descrpt"] = descriptor.get_dim_out() grad_force = "direct" not in fitting_net["type"] diff --git a/deepmd/pt/model/model/denoise_model.py b/deepmd/pt/model/model/denoise_model.py deleted file mode 100644 index f5db6f5835..0000000000 --- a/deepmd/pt/model/model/denoise_model.py +++ /dev/null @@ -1,81 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -from typing import ( - Dict, - Optional, -) - -import torch - -from .dp_model import ( - DPModel, -) - - -class DenoiseModel(DPModel): - model_type = "denoise" - - def __init__( - self, - *args, - **kwargs, - ): - super().__init__(*args, **kwargs) - - def forward( - self, - coord, - atype, - box: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - do_atomic_virial: bool = False, - ) -> Dict[str, torch.Tensor]: - model_ret = self.forward_common( - coord, - atype, - box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - model_predict = model_ret - return model_predict - - @torch.jit.export - def forward_lower( - self, - extended_coord, - extended_atype, - nlist, - mapping: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - do_atomic_virial: bool = False, - ): - model_ret = self.forward_common_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - if self.get_fitting_net() is not None: - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy"): - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy"): - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["extended_virial"] = model_ret[ - "energy_derv_c" - ].squeeze(-3) - else: - assert model_ret["dforce"] is not None - model_predict["dforce"] = model_ret["dforce"] - else: - model_predict = model_ret - return model_predict diff --git a/deepmd/pt/model/model/dp_model.py b/deepmd/pt/model/model/dp_model.py index 1cf8282981..1b29eebd9f 100644 --- a/deepmd/pt/model/model/dp_model.py +++ b/deepmd/pt/model/model/dp_model.py @@ -15,9 +15,6 @@ from deepmd.pt.model.model.model import ( BaseModel, ) -from deepmd.pt.model.task.denoise import ( - DenoiseFittingNet, -) from deepmd.pt.model.task.dipole import ( DipoleFittingNet, ) @@ -51,9 +48,6 @@ def __new__( atomic_model_: Optional[DPAtomicModel] = None, **kwargs, ): - from deepmd.pt.model.model.denoise_model import ( - DenoiseModel, - ) from deepmd.pt.model.model.dipole_model import ( DipoleModel, ) @@ -90,8 +84,6 @@ def __new__( cls = DOSModel elif isinstance(fitting, PropertyFittingNet): cls = PropertyModel - elif isinstance(fitting, DenoiseFittingNet): - cls = DenoiseModel # else: unknown fitting type, fall back to DPModel return super().__new__(cls) diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 66f58ca665..74b258013b 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -65,21 +65,9 @@ def forward_lower( aparam=aparam, do_atomic_virial=do_atomic_virial, ) - if self.get_fitting_net() is not None: - model_predict = {} - model_predict["atom_energy"] = model_ret["energy"] - model_predict["energy"] = model_ret["energy_redu"] - if self.do_grad_r("energy"): - model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2) - if self.do_grad_c("energy"): - model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2) - if do_atomic_virial: - model_predict["extended_virial"] = model_ret[ - "energy_derv_c" - ].squeeze(-3) - else: - assert model_ret["dforce"] is not None - model_predict["dforce"] = model_ret["dforce"] - else: - model_predict = model_ret + model_predict = {} + model_predict["atom_property"] = model_ret["property"] + model_predict["property"] = model_ret["property_redu"] / atype.shape[-1] + if "mask" in model_ret: + model_predict["mask"] = model_ret["mask"] return model_predict diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py deleted file mode 100644 index 48d6893f8b..0000000000 --- a/deepmd/pt/model/task/denoise.py +++ /dev/null @@ -1,643 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import copy -import logging -from typing import ( - Callable, - List, - Optional, - Union, -) - -import numpy as np -import torch - -from deepmd.dpmodel import ( - FittingOutputDef, - OutputVariableDef, - fitting_check_output, -) -from deepmd.pt.model.network.mlp import ( - FittingNet, - NetworkCollection, -) -from deepmd.pt.model.network.network import ( - MaskLMHead, - NonLinearHead, -) -from deepmd.pt.model.task.fitting import ( - Fitting, -) -from deepmd.pt.utils import ( - env, -) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, -) -from deepmd.pt.utils.exclude_mask import ( - AtomExcludeMask, -) -from deepmd.pt.utils.utils import ( - to_numpy_array, - to_torch_tensor, -) -from deepmd.utils.path import ( - DPPath, -) - -dtype = env.GLOBAL_PT_FLOAT_PRECISION -device = env.DEVICE - -log = logging.getLogger(__name__) - - -@Fitting.register("denoise") -class DenoiseFittingNet(Fitting): - """Construct a denoise 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. - mixed_types : bool - If true, use a uniform fitting net for all atom types, otherwise use - different fitting nets for different atom types. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. - exclude_types: List[int] - Atomic contributions of the excluded atom types are set zero. - trainable : Union[List[bool], bool] - If the parameters in the fitting net are trainable. - Now this only supports setting all the parameters in the fitting net at one state. - When in List[bool], the trainable will be True only if all the boolean parameters are True. - remove_vaccum_contribution: List[bool], optional - Remove vaccum contribution before the bias is added. The list assigned each - type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same - length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. - """ - - def __init__( - self, - ntypes: int, - dim_descrpt: int, - embedding_width: 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, - mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[int] = None, - exclude_types: List[int] = [], - trainable: Union[bool, List[bool]] = True, - remove_vaccum_contribution: Optional[List[bool]] = None, - **kwargs, - ): - super().__init__() - self.var_name = ["updated_coord", "logits"] - self.ntypes = ntypes - self.dim_descrpt = dim_descrpt - self.embedding_width = embedding_width - self.neuron = neuron - self.mixed_types = mixed_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] - self.rcond = rcond - # order matters, should be place after the assignment of ntypes - self.reinit_exclude(exclude_types) - self.trainable = trainable - # need support for each layer settings - self.trainable = ( - all(self.trainable) if isinstance(self.trainable, list) else self.trainable - ) - self.remove_vaccum_contribution = remove_vaccum_contribution - - # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] - net_dim_out = self._net_out_dim() - - # init constants - # TODO: actually bias is useless in denoise task - if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) - bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) - if not self.mixed_types: - 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", - 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_coord = self.embedding_width - in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam - - # the first MLP is used to update coordinate - self.filter_layers_coord = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_coord, - 1, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - # the second MLP is used to update logits - self.filter_layers_logits = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_logits, - net_dim_out[1], - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - self.filter_layers_old = None - - if seed is not None: - torch.manual_seed(seed) - # set trainable - for param in self.parameters(): - param.requires_grad = self.trainable - - def reinit_exclude( - self, - exclude_types: List[int] = [], - ): - self.exclude_types = exclude_types - self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) - - def serialize(self) -> dict: - """Serialize the fitting to dict.""" - return { - "@class": "Fitting", - "@version": 1, - "var_name": self.var_name, - "ntypes": self.ntypes, - "dim_descrpt": self.dim_descrpt, - "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, - "mixed_types": self.mixed_types, - "nets": self.filter_layers.serialize(), - "rcond": self.rcond, - "exclude_types": self.exclude_types, - "@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), - "aparam_inv_std": to_numpy_array(self.aparam_inv_std), - }, - # "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 - "tot_ener_zero": False, - "trainable": [self.trainable] * (len(self.neuron) + 1), - "layer_name": None, - "use_aparam_as_mask": False, - "spin": None, - } - - @classmethod - def deserialize(cls, data: dict) -> "DenoiseFittingNet": - 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 - - # make jit happy - exclude_types: List[int] - - 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. - """ - # make jit happy - sel_type: List[int] = [] - for ii in range(self.ntypes): - if ii not in self.exclude_types: - sel_type.append(ii) - return sel_type - - def __setitem__(self, key, value): - if key in ["bias_atom_e"]: - value = value.view([self.ntypes, self._net_out_dim()]) - 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 - elif key in ["scale"]: - self.scale = 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 - elif key in ["scale"]: - return self.scale - else: - raise KeyError(key) - - def _net_out_dim(self): - """Set the FittingNet output dim.""" - return [3, self.ntypes] - # pass - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "updated_coord", - [3], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - OutputVariableDef( - "logits", - [-1], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - ] - ) - - 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( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): - xx = descriptor - nnei = g2.shape[-2] - if self.remove_vaccum_contribution is not None: - # TODO: Idealy, the input for vaccum should be computed; - # we consider it as always zero for convenience. - # Needs a compute_input_stats for vaccum passed from the - # descriptor. - xx_zeros = torch.zeros_like(xx) - else: - xx_zeros = None - nf, nloc, nd = xx.shape - net_dim_out = self._net_out_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([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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, 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( - f"get an input aparam of dim {aparam.shape[-1]}, ", - f"which is not consistent with {self.numb_aparam}.", - ) - aparam = aparam.view([nf, -1, 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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, aparam], - dim=-1, - ) - - outs_coord = torch.zeros( - (nf, nloc, net_dim_out[0]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) # jit assertion - outs_logits = torch.zeros( - (nf, nloc, net_dim_out[1]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) - if self.mixed_types: - atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( - dim=-2 - ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) - atom_logits = self.filter_layers_logits[0](xx) - # Is xx_zeros useful in denoise task?????????????? - # if xx_zeros is not None: - # atom_property -= self.filter_layers.networks[0](xx_zeros) - outs_coord = ( - outs_coord + atom_updated_coord - ) # Shape is [nframes, natoms[0], net_dim_out] - outs_logits = outs_logits + atom_logits - # TODO: - """ - else: - for type_i, ll in enumerate(self.filter_layers_coord.networks): - mask = (atype == type_i).unsqueeze(-1) - mask = torch.tile(mask, (1, 1, net_dim_out)) - atom_property = ll(xx) - if xx_zeros is not None: - # must assert, otherwise jit is not happy - assert self.remove_vaccum_contribution is not None - if not ( - len(self.remove_vaccum_contribution) > type_i - and not self.remove_vaccum_contribution[type_i] - ): - atom_property -= ll(xx_zeros) - atom_property = atom_property + self.bias_atom_e[type_i] - atom_property = atom_property * mask - outs = ( - outs + atom_property - ) # Shape is [nframes, natoms[0], net_dim_out] - """ - # nf x nloc - mask = self.emask(atype) - # nf x nloc x nod - outs_coord = outs_coord * mask[:, :, None] - outs_logits = outs_logits * mask[:, :, None] - return { - self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), - self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), - } - - def compute_output_stats( - self, - merged: Union[Callable[[], List[dict]], List[dict]], - stat_file_path: Optional[DPPath] = None, - ): - """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. - - Parameters - ---------- - merged : Union[Callable[[], List[dict]], List[dict]] - - List[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], List[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. - - """ - pass - - -@fitting_check_output -class DenoiseNet(Fitting): - def __init__( - self, - feature_dim, - ntypes, - attn_head=8, - prefactor=[0.5, 0.5], - activation_function="gelu", - **kwargs, - ): - """Construct a denoise net. - - 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.feature_dim = feature_dim - self.ntypes = ntypes - self.attn_head = attn_head - self.prefactor = torch.tensor( - prefactor, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE - ) - - self.lm_head = MaskLMHead( - embed_dim=self.feature_dim, - output_dim=ntypes, - activation_fn=activation_function, - weight=None, - ) - - if not isinstance(self.attn_head, list): - self.pair2coord_proj = NonLinearHead( - self.attn_head, 1, activation_fn=activation_function - ) - else: - self.pair2coord_proj = [] - self.ndescriptor = len(self.attn_head) - for ii in range(self.ndescriptor): - _pair2coord_proj = NonLinearHead( - self.attn_head[ii], 1, activation_fn=activation_function - ) - self.pair2coord_proj.append(_pair2coord_proj) - self.pair2coord_proj = torch.nn.ModuleList(self.pair2coord_proj) - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "updated_coord", - [3], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - OutputVariableDef( - "logits", - [-1], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - ] - ) - - def forward( - self, - pair_weights, - diff, - nlist_mask, - features, - sw, - masked_tokens: Optional[torch.Tensor] = None, - ): - """Calculate the updated coord. - Args: - - coord: Input noisy coord with shape [nframes, nloc, 3]. - - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. - - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. - - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei]. - - Returns - ------- - - denoised_coord: Denoised updated coord with shape [nframes, nloc, 3]. - """ - # [nframes, nloc, nnei, 1] - logits = self.lm_head(features, masked_tokens=masked_tokens) - if not isinstance(self.attn_head, list): - attn_probs = self.pair2coord_proj(pair_weights) - out_coord = (attn_probs * diff).sum(dim=-2) / ( - sw.sum(dim=-1).unsqueeze(-1) + 1e-6 - ) - else: - assert len(self.prefactor) == self.ndescriptor - all_coord_update = [] - assert len(pair_weights) == len(diff) == len(nlist_mask) == self.ndescriptor - for ii in range(self.ndescriptor): - _attn_probs = self.pair2coord_proj[ii](pair_weights[ii]) - _coord_update = (_attn_probs * diff[ii]).sum(dim=-2) / ( - nlist_mask[ii].sum(dim=-1).unsqueeze(-1) + 1e-6 - ) - all_coord_update.append(_coord_update) - out_coord = self.prefactor[0] * all_coord_update[0] - for ii in range(self.ndescriptor - 1): - out_coord += self.prefactor[ii + 1] * all_coord_update[ii + 1] - return { - "updated_coord": out_coord, - "logits": logits, - } diff --git a/deepmd/pt/model/task/dipole.py b/deepmd/pt/model/task/dipole.py index ca15ae0c20..ca445c8588 100644 --- a/deepmd/pt/model/task/dipole.py +++ b/deepmd/pt/model/task/dipole.py @@ -176,7 +176,6 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): diff --git a/deepmd/pt/model/task/ener.py b/deepmd/pt/model/task/ener.py index c8e2a77431..12c0917dd2 100644 --- a/deepmd/pt/model/task/ener.py +++ b/deepmd/pt/model/task/ener.py @@ -189,7 +189,6 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, None]: diff --git a/deepmd/pt/model/task/invar_fitting.py b/deepmd/pt/model/task/invar_fitting.py index 0f9028f182..787ed8740e 100644 --- a/deepmd/pt/model/task/invar_fitting.py +++ b/deepmd/pt/model/task/invar_fitting.py @@ -189,7 +189,6 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): diff --git a/deepmd/pt/model/task/polarizability.py b/deepmd/pt/model/task/polarizability.py index 587bce2709..d7428c4d53 100644 --- a/deepmd/pt/model/task/polarizability.py +++ b/deepmd/pt/model/task/polarizability.py @@ -278,7 +278,6 @@ def forward( gr: Optional[torch.Tensor] = None, g2: Optional[torch.Tensor] = None, h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, fparam: Optional[torch.Tensor] = None, aparam: Optional[torch.Tensor] = None, ): diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 84f33f3a07..34855a56c7 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1182,9 +1182,8 @@ def fitting_property(): doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." doc_seed = "Random seed for parameter initialization of the fitting net" - doc_task_num = "The number of outputs of fitting net" + doc_task_num = "The dimension of outputs of fitting net" return [ Argument( "neuron", @@ -1203,13 +1202,6 @@ def fitting_property(): ), Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument( - "sel_type", - [List[int], int, None], - optional=True, - alias=["property_type"], - doc=doc_sel_type + doc_only_tf_supported, - ), Argument("seed", [int, None], optional=True, doc=doc_seed), Argument("task_num", int, optional=True, default=1, doc=doc_task_num), ] From d5eaf30d77bc319bece2e26d49be12fd4de1f8ab Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 19 Mar 2024 13:26:13 +0000 Subject: [PATCH 12/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/loss/property.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index aed50eb28d..e6e82e3f90 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -52,7 +52,7 @@ def __init__( def forward(self, model_pred, label, natoms, learning_rate, mae=False): """Return loss on properties . - Parameters: + Parameters ---------- model_pred : dict[str, torch.Tensor] Model predictions. From 70674f62515ce679dbf5c6bfeb496387aa4e6bcf Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:27:40 +0000 Subject: [PATCH 13/31] recover --- denoise.py | 643 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 643 insertions(+) create mode 100644 denoise.py diff --git a/denoise.py b/denoise.py new file mode 100644 index 0000000000..48d6893f8b --- /dev/null +++ b/denoise.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +from typing import ( + Callable, + List, + Optional, + Union, +) + +import numpy as np +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) +from deepmd.pt.model.network.mlp import ( + FittingNet, + NetworkCollection, +) +from deepmd.pt.model.network.network import ( + MaskLMHead, + NonLinearHead, +) +from deepmd.pt.model.task.fitting import ( + Fitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) +from deepmd.pt.utils.exclude_mask import ( + AtomExcludeMask, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.utils.path import ( + DPPath, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("denoise") +class DenoiseFittingNet(Fitting): + """Construct a denoise 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. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. + exclude_types: List[int] + Atomic contributions of the excluded atom types are set zero. + trainable : Union[List[bool], bool] + If the parameters in the fitting net are trainable. + Now this only supports setting all the parameters in the fitting net at one state. + When in List[bool], the trainable will be True only if all the boolean parameters are True. + remove_vaccum_contribution: List[bool], optional + Remove vaccum contribution before the bias is added. The list assigned each + type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same + length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + embedding_width: 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, + mixed_types: bool = True, + rcond: Optional[float] = None, + seed: Optional[int] = None, + exclude_types: List[int] = [], + trainable: Union[bool, List[bool]] = True, + remove_vaccum_contribution: Optional[List[bool]] = None, + **kwargs, + ): + super().__init__() + self.var_name = ["updated_coord", "logits"] + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.embedding_width = embedding_width + self.neuron = neuron + self.mixed_types = mixed_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] + self.rcond = rcond + # order matters, should be place after the assignment of ntypes + self.reinit_exclude(exclude_types) + self.trainable = trainable + # need support for each layer settings + self.trainable = ( + all(self.trainable) if isinstance(self.trainable, list) else self.trainable + ) + self.remove_vaccum_contribution = remove_vaccum_contribution + + # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] + net_dim_out = self._net_out_dim() + + # init constants + # TODO: actually bias is useless in denoise task + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) + if not self.mixed_types: + 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", + 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_coord = self.embedding_width + in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam + + # the first MLP is used to update coordinate + self.filter_layers_coord = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_coord, + 1, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + # the second MLP is used to update logits + self.filter_layers_logits = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_logits, + net_dim_out[1], + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + self.filter_layers_old = None + + if seed is not None: + torch.manual_seed(seed) + # set trainable + for param in self.parameters(): + param.requires_grad = self.trainable + + def reinit_exclude( + self, + exclude_types: List[int] = [], + ): + self.exclude_types = exclude_types + self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "@class": "Fitting", + "@version": 1, + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "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, + "mixed_types": self.mixed_types, + "nets": self.filter_layers.serialize(), + "rcond": self.rcond, + "exclude_types": self.exclude_types, + "@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), + "aparam_inv_std": to_numpy_array(self.aparam_inv_std), + }, + # "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 + "tot_ener_zero": False, + "trainable": [self.trainable] * (len(self.neuron) + 1), + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "DenoiseFittingNet": + 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 + + # make jit happy + exclude_types: List[int] + + 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. + """ + # make jit happy + sel_type: List[int] = [] + for ii in range(self.ntypes): + if ii not in self.exclude_types: + sel_type.append(ii) + return sel_type + + def __setitem__(self, key, value): + if key in ["bias_atom_e"]: + value = value.view([self.ntypes, self._net_out_dim()]) + 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 + elif key in ["scale"]: + self.scale = 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 + elif key in ["scale"]: + return self.scale + else: + raise KeyError(key) + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return [3, self.ntypes] + # pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "updated_coord", + [3], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + OutputVariableDef( + "logits", + [-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + 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( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + xx = descriptor + nnei = g2.shape[-2] + if self.remove_vaccum_contribution is not None: + # TODO: Idealy, the input for vaccum should be computed; + # we consider it as always zero for convenience. + # Needs a compute_input_stats for vaccum passed from the + # descriptor. + xx_zeros = torch.zeros_like(xx) + else: + xx_zeros = None + nf, nloc, nd = xx.shape + net_dim_out = self._net_out_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([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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, 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( + f"get an input aparam of dim {aparam.shape[-1]}, ", + f"which is not consistent with {self.numb_aparam}.", + ) + aparam = aparam.view([nf, -1, 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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, aparam], + dim=-1, + ) + + outs_coord = torch.zeros( + (nf, nloc, net_dim_out[0]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) # jit assertion + outs_logits = torch.zeros( + (nf, nloc, net_dim_out[1]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) + if self.mixed_types: + atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( + dim=-2 + ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) + atom_logits = self.filter_layers_logits[0](xx) + # Is xx_zeros useful in denoise task?????????????? + # if xx_zeros is not None: + # atom_property -= self.filter_layers.networks[0](xx_zeros) + outs_coord = ( + outs_coord + atom_updated_coord + ) # Shape is [nframes, natoms[0], net_dim_out] + outs_logits = outs_logits + atom_logits + # TODO: + """ + else: + for type_i, ll in enumerate(self.filter_layers_coord.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = torch.tile(mask, (1, 1, net_dim_out)) + atom_property = ll(xx) + if xx_zeros is not None: + # must assert, otherwise jit is not happy + assert self.remove_vaccum_contribution is not None + if not ( + len(self.remove_vaccum_contribution) > type_i + and not self.remove_vaccum_contribution[type_i] + ): + atom_property -= ll(xx_zeros) + atom_property = atom_property + self.bias_atom_e[type_i] + atom_property = atom_property * mask + outs = ( + outs + atom_property + ) # Shape is [nframes, natoms[0], net_dim_out] + """ + # nf x nloc + mask = self.emask(atype) + # nf x nloc x nod + outs_coord = outs_coord * mask[:, :, None] + outs_logits = outs_logits * mask[:, :, None] + return { + self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), + self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), + } + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + +@fitting_check_output +class DenoiseNet(Fitting): + def __init__( + self, + feature_dim, + ntypes, + attn_head=8, + prefactor=[0.5, 0.5], + activation_function="gelu", + **kwargs, + ): + """Construct a denoise net. + + 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.feature_dim = feature_dim + self.ntypes = ntypes + self.attn_head = attn_head + self.prefactor = torch.tensor( + prefactor, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + + self.lm_head = MaskLMHead( + embed_dim=self.feature_dim, + output_dim=ntypes, + activation_fn=activation_function, + weight=None, + ) + + if not isinstance(self.attn_head, list): + self.pair2coord_proj = NonLinearHead( + self.attn_head, 1, activation_fn=activation_function + ) + else: + self.pair2coord_proj = [] + self.ndescriptor = len(self.attn_head) + for ii in range(self.ndescriptor): + _pair2coord_proj = NonLinearHead( + self.attn_head[ii], 1, activation_fn=activation_function + ) + self.pair2coord_proj.append(_pair2coord_proj) + self.pair2coord_proj = torch.nn.ModuleList(self.pair2coord_proj) + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "updated_coord", + [3], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + OutputVariableDef( + "logits", + [-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def forward( + self, + pair_weights, + diff, + nlist_mask, + features, + sw, + masked_tokens: Optional[torch.Tensor] = None, + ): + """Calculate the updated coord. + Args: + - coord: Input noisy coord with shape [nframes, nloc, 3]. + - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. + - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. + - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei]. + + Returns + ------- + - denoised_coord: Denoised updated coord with shape [nframes, nloc, 3]. + """ + # [nframes, nloc, nnei, 1] + logits = self.lm_head(features, masked_tokens=masked_tokens) + if not isinstance(self.attn_head, list): + attn_probs = self.pair2coord_proj(pair_weights) + out_coord = (attn_probs * diff).sum(dim=-2) / ( + sw.sum(dim=-1).unsqueeze(-1) + 1e-6 + ) + else: + assert len(self.prefactor) == self.ndescriptor + all_coord_update = [] + assert len(pair_weights) == len(diff) == len(nlist_mask) == self.ndescriptor + for ii in range(self.ndescriptor): + _attn_probs = self.pair2coord_proj[ii](pair_weights[ii]) + _coord_update = (_attn_probs * diff[ii]).sum(dim=-2) / ( + nlist_mask[ii].sum(dim=-1).unsqueeze(-1) + 1e-6 + ) + all_coord_update.append(_coord_update) + out_coord = self.prefactor[0] * all_coord_update[0] + for ii in range(self.ndescriptor - 1): + out_coord += self.prefactor[ii + 1] * all_coord_update[ii + 1] + return { + "updated_coord": out_coord, + "logits": logits, + } From 55c5f9635dd389c5a983fd99e0f72a5d94db0cb3 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:29:32 +0000 Subject: [PATCH 14/31] recover --- deepmd/pt/model/task/denoise.py | 643 ++++++++++++++++++++++++++++++++ 1 file changed, 643 insertions(+) create mode 100644 deepmd/pt/model/task/denoise.py diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py new file mode 100644 index 0000000000..48d6893f8b --- /dev/null +++ b/deepmd/pt/model/task/denoise.py @@ -0,0 +1,643 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import copy +import logging +from typing import ( + Callable, + List, + Optional, + Union, +) + +import numpy as np +import torch + +from deepmd.dpmodel import ( + FittingOutputDef, + OutputVariableDef, + fitting_check_output, +) +from deepmd.pt.model.network.mlp import ( + FittingNet, + NetworkCollection, +) +from deepmd.pt.model.network.network import ( + MaskLMHead, + NonLinearHead, +) +from deepmd.pt.model.task.fitting import ( + Fitting, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.pt.utils.env import ( + DEFAULT_PRECISION, + PRECISION_DICT, +) +from deepmd.pt.utils.exclude_mask import ( + AtomExcludeMask, +) +from deepmd.pt.utils.utils import ( + to_numpy_array, + to_torch_tensor, +) +from deepmd.utils.path import ( + DPPath, +) + +dtype = env.GLOBAL_PT_FLOAT_PRECISION +device = env.DEVICE + +log = logging.getLogger(__name__) + + +@Fitting.register("denoise") +class DenoiseFittingNet(Fitting): + """Construct a denoise 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. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + rcond : float, optional + The condition number for the regression of atomic energy. + seed : int, optional + Random seed. + exclude_types: List[int] + Atomic contributions of the excluded atom types are set zero. + trainable : Union[List[bool], bool] + If the parameters in the fitting net are trainable. + Now this only supports setting all the parameters in the fitting net at one state. + When in List[bool], the trainable will be True only if all the boolean parameters are True. + remove_vaccum_contribution: List[bool], optional + Remove vaccum contribution before the bias is added. The list assigned each + type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same + length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. + """ + + def __init__( + self, + ntypes: int, + dim_descrpt: int, + embedding_width: 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, + mixed_types: bool = True, + rcond: Optional[float] = None, + seed: Optional[int] = None, + exclude_types: List[int] = [], + trainable: Union[bool, List[bool]] = True, + remove_vaccum_contribution: Optional[List[bool]] = None, + **kwargs, + ): + super().__init__() + self.var_name = ["updated_coord", "logits"] + self.ntypes = ntypes + self.dim_descrpt = dim_descrpt + self.embedding_width = embedding_width + self.neuron = neuron + self.mixed_types = mixed_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] + self.rcond = rcond + # order matters, should be place after the assignment of ntypes + self.reinit_exclude(exclude_types) + self.trainable = trainable + # need support for each layer settings + self.trainable = ( + all(self.trainable) if isinstance(self.trainable, list) else self.trainable + ) + self.remove_vaccum_contribution = remove_vaccum_contribution + + # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] + net_dim_out = self._net_out_dim() + + # init constants + # TODO: actually bias is useless in denoise task + if bias_atom_e is None: + bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) + bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) + bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) + if not self.mixed_types: + 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", + 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_coord = self.embedding_width + in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam + + # the first MLP is used to update coordinate + self.filter_layers_coord = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_coord, + 1, + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + # the second MLP is used to update logits + self.filter_layers_logits = NetworkCollection( + 1 if not self.mixed_types else 0, + self.ntypes, + network_type="fitting_network", + networks=[ + FittingNet( + in_dim_logits, + net_dim_out[1], + self.neuron, + self.activation_function, + self.resnet_dt, + self.precision, + bias_out=True, + ) + for ii in range(self.ntypes if not self.mixed_types else 1) + ], + ) + + self.filter_layers_old = None + + if seed is not None: + torch.manual_seed(seed) + # set trainable + for param in self.parameters(): + param.requires_grad = self.trainable + + def reinit_exclude( + self, + exclude_types: List[int] = [], + ): + self.exclude_types = exclude_types + self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) + + def serialize(self) -> dict: + """Serialize the fitting to dict.""" + return { + "@class": "Fitting", + "@version": 1, + "var_name": self.var_name, + "ntypes": self.ntypes, + "dim_descrpt": self.dim_descrpt, + "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, + "mixed_types": self.mixed_types, + "nets": self.filter_layers.serialize(), + "rcond": self.rcond, + "exclude_types": self.exclude_types, + "@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), + "aparam_inv_std": to_numpy_array(self.aparam_inv_std), + }, + # "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 + "tot_ener_zero": False, + "trainable": [self.trainable] * (len(self.neuron) + 1), + "layer_name": None, + "use_aparam_as_mask": False, + "spin": None, + } + + @classmethod + def deserialize(cls, data: dict) -> "DenoiseFittingNet": + 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 + + # make jit happy + exclude_types: List[int] + + 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. + """ + # make jit happy + sel_type: List[int] = [] + for ii in range(self.ntypes): + if ii not in self.exclude_types: + sel_type.append(ii) + return sel_type + + def __setitem__(self, key, value): + if key in ["bias_atom_e"]: + value = value.view([self.ntypes, self._net_out_dim()]) + 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 + elif key in ["scale"]: + self.scale = 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 + elif key in ["scale"]: + return self.scale + else: + raise KeyError(key) + + def _net_out_dim(self): + """Set the FittingNet output dim.""" + return [3, self.ntypes] + # pass + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "updated_coord", + [3], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + OutputVariableDef( + "logits", + [-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + 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( + self, + descriptor: torch.Tensor, + atype: torch.Tensor, + gr: Optional[torch.Tensor] = None, + g2: Optional[torch.Tensor] = None, + h2: Optional[torch.Tensor] = None, + sw: Optional[torch.Tensor] = None, + fparam: Optional[torch.Tensor] = None, + aparam: Optional[torch.Tensor] = None, + ): + xx = descriptor + nnei = g2.shape[-2] + if self.remove_vaccum_contribution is not None: + # TODO: Idealy, the input for vaccum should be computed; + # we consider it as always zero for convenience. + # Needs a compute_input_stats for vaccum passed from the + # descriptor. + xx_zeros = torch.zeros_like(xx) + else: + xx_zeros = None + nf, nloc, nd = xx.shape + net_dim_out = self._net_out_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([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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, 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( + f"get an input aparam of dim {aparam.shape[-1]}, ", + f"which is not consistent with {self.numb_aparam}.", + ) + aparam = aparam.view([nf, -1, 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, + ) + if xx_zeros is not None: + xx_zeros = torch.cat( + [xx_zeros, aparam], + dim=-1, + ) + + outs_coord = torch.zeros( + (nf, nloc, net_dim_out[0]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) # jit assertion + outs_logits = torch.zeros( + (nf, nloc, net_dim_out[1]), + dtype=env.GLOBAL_PT_FLOAT_PRECISION, + device=descriptor.device, + ) + if self.mixed_types: + atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( + dim=-2 + ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) + atom_logits = self.filter_layers_logits[0](xx) + # Is xx_zeros useful in denoise task?????????????? + # if xx_zeros is not None: + # atom_property -= self.filter_layers.networks[0](xx_zeros) + outs_coord = ( + outs_coord + atom_updated_coord + ) # Shape is [nframes, natoms[0], net_dim_out] + outs_logits = outs_logits + atom_logits + # TODO: + """ + else: + for type_i, ll in enumerate(self.filter_layers_coord.networks): + mask = (atype == type_i).unsqueeze(-1) + mask = torch.tile(mask, (1, 1, net_dim_out)) + atom_property = ll(xx) + if xx_zeros is not None: + # must assert, otherwise jit is not happy + assert self.remove_vaccum_contribution is not None + if not ( + len(self.remove_vaccum_contribution) > type_i + and not self.remove_vaccum_contribution[type_i] + ): + atom_property -= ll(xx_zeros) + atom_property = atom_property + self.bias_atom_e[type_i] + atom_property = atom_property * mask + outs = ( + outs + atom_property + ) # Shape is [nframes, natoms[0], net_dim_out] + """ + # nf x nloc + mask = self.emask(atype) + # nf x nloc x nod + outs_coord = outs_coord * mask[:, :, None] + outs_logits = outs_logits * mask[:, :, None] + return { + self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), + self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), + } + + def compute_output_stats( + self, + merged: Union[Callable[[], List[dict]], List[dict]], + stat_file_path: Optional[DPPath] = None, + ): + """ + Compute the output statistics (e.g. energy bias) for the fitting net from packed data. + + Parameters + ---------- + merged : Union[Callable[[], List[dict]], List[dict]] + - List[dict]: A list of data samples from various data systems. + Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` + originating from the `i`-th data system. + - Callable[[], List[dict]]: A lazy function that returns data samples in the above format + only when needed. Since the sampling process can be slow and memory-intensive, + the lazy function helps by only sampling once. + stat_file_path : Optional[DPPath] + The path to the stat file. + + """ + pass + + +@fitting_check_output +class DenoiseNet(Fitting): + def __init__( + self, + feature_dim, + ntypes, + attn_head=8, + prefactor=[0.5, 0.5], + activation_function="gelu", + **kwargs, + ): + """Construct a denoise net. + + 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.feature_dim = feature_dim + self.ntypes = ntypes + self.attn_head = attn_head + self.prefactor = torch.tensor( + prefactor, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + + self.lm_head = MaskLMHead( + embed_dim=self.feature_dim, + output_dim=ntypes, + activation_fn=activation_function, + weight=None, + ) + + if not isinstance(self.attn_head, list): + self.pair2coord_proj = NonLinearHead( + self.attn_head, 1, activation_fn=activation_function + ) + else: + self.pair2coord_proj = [] + self.ndescriptor = len(self.attn_head) + for ii in range(self.ndescriptor): + _pair2coord_proj = NonLinearHead( + self.attn_head[ii], 1, activation_fn=activation_function + ) + self.pair2coord_proj.append(_pair2coord_proj) + self.pair2coord_proj = torch.nn.ModuleList(self.pair2coord_proj) + + def output_def(self): + return FittingOutputDef( + [ + OutputVariableDef( + "updated_coord", + [3], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + OutputVariableDef( + "logits", + [-1], + reduciable=False, + r_differentiable=False, + c_differentiable=False, + ), + ] + ) + + def forward( + self, + pair_weights, + diff, + nlist_mask, + features, + sw, + masked_tokens: Optional[torch.Tensor] = None, + ): + """Calculate the updated coord. + Args: + - coord: Input noisy coord with shape [nframes, nloc, 3]. + - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. + - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. + - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei]. + + Returns + ------- + - denoised_coord: Denoised updated coord with shape [nframes, nloc, 3]. + """ + # [nframes, nloc, nnei, 1] + logits = self.lm_head(features, masked_tokens=masked_tokens) + if not isinstance(self.attn_head, list): + attn_probs = self.pair2coord_proj(pair_weights) + out_coord = (attn_probs * diff).sum(dim=-2) / ( + sw.sum(dim=-1).unsqueeze(-1) + 1e-6 + ) + else: + assert len(self.prefactor) == self.ndescriptor + all_coord_update = [] + assert len(pair_weights) == len(diff) == len(nlist_mask) == self.ndescriptor + for ii in range(self.ndescriptor): + _attn_probs = self.pair2coord_proj[ii](pair_weights[ii]) + _coord_update = (_attn_probs * diff[ii]).sum(dim=-2) / ( + nlist_mask[ii].sum(dim=-1).unsqueeze(-1) + 1e-6 + ) + all_coord_update.append(_coord_update) + out_coord = self.prefactor[0] * all_coord_update[0] + for ii in range(self.ndescriptor - 1): + out_coord += self.prefactor[ii + 1] * all_coord_update[ii + 1] + return { + "updated_coord": out_coord, + "logits": logits, + } From b645176f7ed92f7a933dfabe5eef0639cda8864b Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:35:48 +0000 Subject: [PATCH 15/31] delete denoise --- deepmd/pt/model/task/denoise.py | 505 -------------------------------- 1 file changed, 505 deletions(-) diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 48d6893f8b..51058b832f 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -1,11 +1,6 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import copy -import logging from typing import ( - Callable, - List, Optional, - Union, ) import numpy as np @@ -16,10 +11,6 @@ OutputVariableDef, fitting_check_output, ) -from deepmd.pt.model.network.mlp import ( - FittingNet, - NetworkCollection, -) from deepmd.pt.model.network.network import ( MaskLMHead, NonLinearHead, @@ -30,502 +21,6 @@ from deepmd.pt.utils import ( env, ) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, -) -from deepmd.pt.utils.exclude_mask import ( - AtomExcludeMask, -) -from deepmd.pt.utils.utils import ( - to_numpy_array, - to_torch_tensor, -) -from deepmd.utils.path import ( - DPPath, -) - -dtype = env.GLOBAL_PT_FLOAT_PRECISION -device = env.DEVICE - -log = logging.getLogger(__name__) - - -@Fitting.register("denoise") -class DenoiseFittingNet(Fitting): - """Construct a denoise 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. - mixed_types : bool - If true, use a uniform fitting net for all atom types, otherwise use - different fitting nets for different atom types. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. - exclude_types: List[int] - Atomic contributions of the excluded atom types are set zero. - trainable : Union[List[bool], bool] - If the parameters in the fitting net are trainable. - Now this only supports setting all the parameters in the fitting net at one state. - When in List[bool], the trainable will be True only if all the boolean parameters are True. - remove_vaccum_contribution: List[bool], optional - Remove vaccum contribution before the bias is added. The list assigned each - type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same - length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. - """ - - def __init__( - self, - ntypes: int, - dim_descrpt: int, - embedding_width: 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, - mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[int] = None, - exclude_types: List[int] = [], - trainable: Union[bool, List[bool]] = True, - remove_vaccum_contribution: Optional[List[bool]] = None, - **kwargs, - ): - super().__init__() - self.var_name = ["updated_coord", "logits"] - self.ntypes = ntypes - self.dim_descrpt = dim_descrpt - self.embedding_width = embedding_width - self.neuron = neuron - self.mixed_types = mixed_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] - self.rcond = rcond - # order matters, should be place after the assignment of ntypes - self.reinit_exclude(exclude_types) - self.trainable = trainable - # need support for each layer settings - self.trainable = ( - all(self.trainable) if isinstance(self.trainable, list) else self.trainable - ) - self.remove_vaccum_contribution = remove_vaccum_contribution - - # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] - net_dim_out = self._net_out_dim() - - # init constants - # TODO: actually bias is useless in denoise task - if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) - bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) - if not self.mixed_types: - 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", - 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_coord = self.embedding_width - in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam - - # the first MLP is used to update coordinate - self.filter_layers_coord = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_coord, - 1, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - # the second MLP is used to update logits - self.filter_layers_logits = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_logits, - net_dim_out[1], - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - self.filter_layers_old = None - - if seed is not None: - torch.manual_seed(seed) - # set trainable - for param in self.parameters(): - param.requires_grad = self.trainable - - def reinit_exclude( - self, - exclude_types: List[int] = [], - ): - self.exclude_types = exclude_types - self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) - - def serialize(self) -> dict: - """Serialize the fitting to dict.""" - return { - "@class": "Fitting", - "@version": 1, - "var_name": self.var_name, - "ntypes": self.ntypes, - "dim_descrpt": self.dim_descrpt, - "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, - "mixed_types": self.mixed_types, - "nets": self.filter_layers.serialize(), - "rcond": self.rcond, - "exclude_types": self.exclude_types, - "@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), - "aparam_inv_std": to_numpy_array(self.aparam_inv_std), - }, - # "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 - "tot_ener_zero": False, - "trainable": [self.trainable] * (len(self.neuron) + 1), - "layer_name": None, - "use_aparam_as_mask": False, - "spin": None, - } - - @classmethod - def deserialize(cls, data: dict) -> "DenoiseFittingNet": - 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 - - # make jit happy - exclude_types: List[int] - - 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. - """ - # make jit happy - sel_type: List[int] = [] - for ii in range(self.ntypes): - if ii not in self.exclude_types: - sel_type.append(ii) - return sel_type - - def __setitem__(self, key, value): - if key in ["bias_atom_e"]: - value = value.view([self.ntypes, self._net_out_dim()]) - 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 - elif key in ["scale"]: - self.scale = 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 - elif key in ["scale"]: - return self.scale - else: - raise KeyError(key) - - def _net_out_dim(self): - """Set the FittingNet output dim.""" - return [3, self.ntypes] - # pass - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "updated_coord", - [3], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - OutputVariableDef( - "logits", - [-1], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - ] - ) - - 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( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): - xx = descriptor - nnei = g2.shape[-2] - if self.remove_vaccum_contribution is not None: - # TODO: Idealy, the input for vaccum should be computed; - # we consider it as always zero for convenience. - # Needs a compute_input_stats for vaccum passed from the - # descriptor. - xx_zeros = torch.zeros_like(xx) - else: - xx_zeros = None - nf, nloc, nd = xx.shape - net_dim_out = self._net_out_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([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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, 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( - f"get an input aparam of dim {aparam.shape[-1]}, ", - f"which is not consistent with {self.numb_aparam}.", - ) - aparam = aparam.view([nf, -1, 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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, aparam], - dim=-1, - ) - - outs_coord = torch.zeros( - (nf, nloc, net_dim_out[0]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) # jit assertion - outs_logits = torch.zeros( - (nf, nloc, net_dim_out[1]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) - if self.mixed_types: - atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( - dim=-2 - ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) - atom_logits = self.filter_layers_logits[0](xx) - # Is xx_zeros useful in denoise task?????????????? - # if xx_zeros is not None: - # atom_property -= self.filter_layers.networks[0](xx_zeros) - outs_coord = ( - outs_coord + atom_updated_coord - ) # Shape is [nframes, natoms[0], net_dim_out] - outs_logits = outs_logits + atom_logits - # TODO: - """ - else: - for type_i, ll in enumerate(self.filter_layers_coord.networks): - mask = (atype == type_i).unsqueeze(-1) - mask = torch.tile(mask, (1, 1, net_dim_out)) - atom_property = ll(xx) - if xx_zeros is not None: - # must assert, otherwise jit is not happy - assert self.remove_vaccum_contribution is not None - if not ( - len(self.remove_vaccum_contribution) > type_i - and not self.remove_vaccum_contribution[type_i] - ): - atom_property -= ll(xx_zeros) - atom_property = atom_property + self.bias_atom_e[type_i] - atom_property = atom_property * mask - outs = ( - outs + atom_property - ) # Shape is [nframes, natoms[0], net_dim_out] - """ - # nf x nloc - mask = self.emask(atype) - # nf x nloc x nod - outs_coord = outs_coord * mask[:, :, None] - outs_logits = outs_logits * mask[:, :, None] - return { - self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), - self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), - } - - def compute_output_stats( - self, - merged: Union[Callable[[], List[dict]], List[dict]], - stat_file_path: Optional[DPPath] = None, - ): - """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. - - Parameters - ---------- - merged : Union[Callable[[], List[dict]], List[dict]] - - List[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], List[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. - - """ - pass @fitting_check_output From 03cbabaaf19a2d704f1e92ec6432a411c41b13dd Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:36:20 +0000 Subject: [PATCH 16/31] delete denoise --- deepmd/pt/model/task/denoise.py | 1 - 1 file changed, 1 deletion(-) diff --git a/deepmd/pt/model/task/denoise.py b/deepmd/pt/model/task/denoise.py index 51058b832f..5f1e780de3 100644 --- a/deepmd/pt/model/task/denoise.py +++ b/deepmd/pt/model/task/denoise.py @@ -3,7 +3,6 @@ Optional, ) -import numpy as np import torch from deepmd.dpmodel import ( From 196eb0d0bed705921cf736b6da23b0a5e38b3fef Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:37:48 +0000 Subject: [PATCH 17/31] delete denoise argcheck --- deepmd/utils/argcheck.py | 38 -------------------------------------- 1 file changed, 38 deletions(-) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 34855a56c7..40b1910f86 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1206,44 +1206,6 @@ def fitting_property(): Argument("task_num", int, optional=True, default=1, doc=doc_task_num), ] - -@fitting_args_plugin.register("denoise") -def fitting_denoise(): - doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." - doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' - doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' - doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." - doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected." - doc_seed = "Random seed for parameter initialization of the fitting net" - return [ - Argument( - "neuron", - List[int], - optional=True, - default=[120, 120, 120], - alias=["n_neuron"], - doc=doc_neuron, - ), - Argument( - "activation_function", - str, - optional=True, - default="tanh", - doc=doc_activation_function, - ), - Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), - Argument("precision", str, optional=True, default="default", doc=doc_precision), - Argument( - "sel_type", - [List[int], int, None], - optional=True, - alias=["denoise_type"], - doc=doc_sel_type + doc_only_tf_supported, - ), - Argument("seed", [int, None], optional=True, doc=doc_seed), - ] - - # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ From e137ad5565898e42a5e3ab073de39f8c18cd9177 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 19 Mar 2024 13:38:24 +0000 Subject: [PATCH 18/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/utils/argcheck.py | 1 + 1 file changed, 1 insertion(+) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 40b1910f86..1a9421161b 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1206,6 +1206,7 @@ def fitting_property(): Argument("task_num", int, optional=True, default=1, doc=doc_task_num), ] + # YWolfeee: Delete global polar mode, merge it into polar mode and use loss setting to support. def fitting_variant_type_args(): doc_descrpt_type = "The type of the fitting. See explanation below. \n\n\ From 8afc64fbfd9293ba9bb14c170850fd81e5e2cb09 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:39:19 +0000 Subject: [PATCH 19/31] delete denoise file --- denoise.py | 643 ----------------------------------------------------- 1 file changed, 643 deletions(-) delete mode 100644 denoise.py diff --git a/denoise.py b/denoise.py deleted file mode 100644 index 48d6893f8b..0000000000 --- a/denoise.py +++ /dev/null @@ -1,643 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import copy -import logging -from typing import ( - Callable, - List, - Optional, - Union, -) - -import numpy as np -import torch - -from deepmd.dpmodel import ( - FittingOutputDef, - OutputVariableDef, - fitting_check_output, -) -from deepmd.pt.model.network.mlp import ( - FittingNet, - NetworkCollection, -) -from deepmd.pt.model.network.network import ( - MaskLMHead, - NonLinearHead, -) -from deepmd.pt.model.task.fitting import ( - Fitting, -) -from deepmd.pt.utils import ( - env, -) -from deepmd.pt.utils.env import ( - DEFAULT_PRECISION, - PRECISION_DICT, -) -from deepmd.pt.utils.exclude_mask import ( - AtomExcludeMask, -) -from deepmd.pt.utils.utils import ( - to_numpy_array, - to_torch_tensor, -) -from deepmd.utils.path import ( - DPPath, -) - -dtype = env.GLOBAL_PT_FLOAT_PRECISION -device = env.DEVICE - -log = logging.getLogger(__name__) - - -@Fitting.register("denoise") -class DenoiseFittingNet(Fitting): - """Construct a denoise 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. - mixed_types : bool - If true, use a uniform fitting net for all atom types, otherwise use - different fitting nets for different atom types. - rcond : float, optional - The condition number for the regression of atomic energy. - seed : int, optional - Random seed. - exclude_types: List[int] - Atomic contributions of the excluded atom types are set zero. - trainable : Union[List[bool], bool] - If the parameters in the fitting net are trainable. - Now this only supports setting all the parameters in the fitting net at one state. - When in List[bool], the trainable will be True only if all the boolean parameters are True. - remove_vaccum_contribution: List[bool], optional - Remove vaccum contribution before the bias is added. The list assigned each - type. For `mixed_types` provide `[True]`, otherwise it should be a list of the same - length as `ntypes` signaling if or not removing the vaccum contribution for the atom types in the list. - """ - - def __init__( - self, - ntypes: int, - dim_descrpt: int, - embedding_width: 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, - mixed_types: bool = True, - rcond: Optional[float] = None, - seed: Optional[int] = None, - exclude_types: List[int] = [], - trainable: Union[bool, List[bool]] = True, - remove_vaccum_contribution: Optional[List[bool]] = None, - **kwargs, - ): - super().__init__() - self.var_name = ["updated_coord", "logits"] - self.ntypes = ntypes - self.dim_descrpt = dim_descrpt - self.embedding_width = embedding_width - self.neuron = neuron - self.mixed_types = mixed_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] - self.rcond = rcond - # order matters, should be place after the assignment of ntypes - self.reinit_exclude(exclude_types) - self.trainable = trainable - # need support for each layer settings - self.trainable = ( - all(self.trainable) if isinstance(self.trainable, list) else self.trainable - ) - self.remove_vaccum_contribution = remove_vaccum_contribution - - # in denoise task, net_dim_out is a list which has 2 elements: [3, ntypes] - net_dim_out = self._net_out_dim() - - # init constants - # TODO: actually bias is useless in denoise task - if bias_atom_e is None: - bias_atom_e = np.zeros([self.ntypes, net_dim_out[0]], dtype=np.float64) - bias_atom_e = torch.tensor(bias_atom_e, dtype=self.prec, device=device) - bias_atom_e = bias_atom_e.view([self.ntypes, net_dim_out[0]]) - if not self.mixed_types: - 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", - 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_coord = self.embedding_width - in_dim_logits = self.dim_descrpt + self.numb_fparam + self.numb_aparam - - # the first MLP is used to update coordinate - self.filter_layers_coord = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_coord, - 1, - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - # the second MLP is used to update logits - self.filter_layers_logits = NetworkCollection( - 1 if not self.mixed_types else 0, - self.ntypes, - network_type="fitting_network", - networks=[ - FittingNet( - in_dim_logits, - net_dim_out[1], - self.neuron, - self.activation_function, - self.resnet_dt, - self.precision, - bias_out=True, - ) - for ii in range(self.ntypes if not self.mixed_types else 1) - ], - ) - - self.filter_layers_old = None - - if seed is not None: - torch.manual_seed(seed) - # set trainable - for param in self.parameters(): - param.requires_grad = self.trainable - - def reinit_exclude( - self, - exclude_types: List[int] = [], - ): - self.exclude_types = exclude_types - self.emask = AtomExcludeMask(self.ntypes, self.exclude_types) - - def serialize(self) -> dict: - """Serialize the fitting to dict.""" - return { - "@class": "Fitting", - "@version": 1, - "var_name": self.var_name, - "ntypes": self.ntypes, - "dim_descrpt": self.dim_descrpt, - "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, - "mixed_types": self.mixed_types, - "nets": self.filter_layers.serialize(), - "rcond": self.rcond, - "exclude_types": self.exclude_types, - "@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), - "aparam_inv_std": to_numpy_array(self.aparam_inv_std), - }, - # "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 - "tot_ener_zero": False, - "trainable": [self.trainable] * (len(self.neuron) + 1), - "layer_name": None, - "use_aparam_as_mask": False, - "spin": None, - } - - @classmethod - def deserialize(cls, data: dict) -> "DenoiseFittingNet": - 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 - - # make jit happy - exclude_types: List[int] - - 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. - """ - # make jit happy - sel_type: List[int] = [] - for ii in range(self.ntypes): - if ii not in self.exclude_types: - sel_type.append(ii) - return sel_type - - def __setitem__(self, key, value): - if key in ["bias_atom_e"]: - value = value.view([self.ntypes, self._net_out_dim()]) - 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 - elif key in ["scale"]: - self.scale = 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 - elif key in ["scale"]: - return self.scale - else: - raise KeyError(key) - - def _net_out_dim(self): - """Set the FittingNet output dim.""" - return [3, self.ntypes] - # pass - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "updated_coord", - [3], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - OutputVariableDef( - "logits", - [-1], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - ] - ) - - 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( - self, - descriptor: torch.Tensor, - atype: torch.Tensor, - gr: Optional[torch.Tensor] = None, - g2: Optional[torch.Tensor] = None, - h2: Optional[torch.Tensor] = None, - sw: Optional[torch.Tensor] = None, - fparam: Optional[torch.Tensor] = None, - aparam: Optional[torch.Tensor] = None, - ): - xx = descriptor - nnei = g2.shape[-2] - if self.remove_vaccum_contribution is not None: - # TODO: Idealy, the input for vaccum should be computed; - # we consider it as always zero for convenience. - # Needs a compute_input_stats for vaccum passed from the - # descriptor. - xx_zeros = torch.zeros_like(xx) - else: - xx_zeros = None - nf, nloc, nd = xx.shape - net_dim_out = self._net_out_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([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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, 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( - f"get an input aparam of dim {aparam.shape[-1]}, ", - f"which is not consistent with {self.numb_aparam}.", - ) - aparam = aparam.view([nf, -1, 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, - ) - if xx_zeros is not None: - xx_zeros = torch.cat( - [xx_zeros, aparam], - dim=-1, - ) - - outs_coord = torch.zeros( - (nf, nloc, net_dim_out[0]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) # jit assertion - outs_logits = torch.zeros( - (nf, nloc, net_dim_out[1]), - dtype=env.GLOBAL_PT_FLOAT_PRECISION, - device=descriptor.device, - ) - if self.mixed_types: - atom_updated_coord = ((self.filter_layers_coord.networks[0](g2)) * h2).sum( - dim=-2 - ) / (sw.sum(dim=-1).unsqueeze(-1) + 1e-6) - atom_logits = self.filter_layers_logits[0](xx) - # Is xx_zeros useful in denoise task?????????????? - # if xx_zeros is not None: - # atom_property -= self.filter_layers.networks[0](xx_zeros) - outs_coord = ( - outs_coord + atom_updated_coord - ) # Shape is [nframes, natoms[0], net_dim_out] - outs_logits = outs_logits + atom_logits - # TODO: - """ - else: - for type_i, ll in enumerate(self.filter_layers_coord.networks): - mask = (atype == type_i).unsqueeze(-1) - mask = torch.tile(mask, (1, 1, net_dim_out)) - atom_property = ll(xx) - if xx_zeros is not None: - # must assert, otherwise jit is not happy - assert self.remove_vaccum_contribution is not None - if not ( - len(self.remove_vaccum_contribution) > type_i - and not self.remove_vaccum_contribution[type_i] - ): - atom_property -= ll(xx_zeros) - atom_property = atom_property + self.bias_atom_e[type_i] - atom_property = atom_property * mask - outs = ( - outs + atom_property - ) # Shape is [nframes, natoms[0], net_dim_out] - """ - # nf x nloc - mask = self.emask(atype) - # nf x nloc x nod - outs_coord = outs_coord * mask[:, :, None] - outs_logits = outs_logits * mask[:, :, None] - return { - self.var_name[0]: outs_coord.to(env.GLOBAL_PT_FLOAT_PRECISION), - self.var_name[1]: outs_logits.to(env.GLOBAL_PT_FLOAT_PRECISION), - } - - def compute_output_stats( - self, - merged: Union[Callable[[], List[dict]], List[dict]], - stat_file_path: Optional[DPPath] = None, - ): - """ - Compute the output statistics (e.g. energy bias) for the fitting net from packed data. - - Parameters - ---------- - merged : Union[Callable[[], List[dict]], List[dict]] - - List[dict]: A list of data samples from various data systems. - Each element, `merged[i]`, is a data dictionary containing `keys`: `torch.Tensor` - originating from the `i`-th data system. - - Callable[[], List[dict]]: A lazy function that returns data samples in the above format - only when needed. Since the sampling process can be slow and memory-intensive, - the lazy function helps by only sampling once. - stat_file_path : Optional[DPPath] - The path to the stat file. - - """ - pass - - -@fitting_check_output -class DenoiseNet(Fitting): - def __init__( - self, - feature_dim, - ntypes, - attn_head=8, - prefactor=[0.5, 0.5], - activation_function="gelu", - **kwargs, - ): - """Construct a denoise net. - - 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.feature_dim = feature_dim - self.ntypes = ntypes - self.attn_head = attn_head - self.prefactor = torch.tensor( - prefactor, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE - ) - - self.lm_head = MaskLMHead( - embed_dim=self.feature_dim, - output_dim=ntypes, - activation_fn=activation_function, - weight=None, - ) - - if not isinstance(self.attn_head, list): - self.pair2coord_proj = NonLinearHead( - self.attn_head, 1, activation_fn=activation_function - ) - else: - self.pair2coord_proj = [] - self.ndescriptor = len(self.attn_head) - for ii in range(self.ndescriptor): - _pair2coord_proj = NonLinearHead( - self.attn_head[ii], 1, activation_fn=activation_function - ) - self.pair2coord_proj.append(_pair2coord_proj) - self.pair2coord_proj = torch.nn.ModuleList(self.pair2coord_proj) - - def output_def(self): - return FittingOutputDef( - [ - OutputVariableDef( - "updated_coord", - [3], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - OutputVariableDef( - "logits", - [-1], - reduciable=False, - r_differentiable=False, - c_differentiable=False, - ), - ] - ) - - def forward( - self, - pair_weights, - diff, - nlist_mask, - features, - sw, - masked_tokens: Optional[torch.Tensor] = None, - ): - """Calculate the updated coord. - Args: - - coord: Input noisy coord with shape [nframes, nloc, 3]. - - pair_weights: Input pair weights with shape [nframes, nloc, nnei, head]. - - diff: Input pair relative coord list with shape [nframes, nloc, nnei, 3]. - - nlist_mask: Input nlist mask with shape [nframes, nloc, nnei]. - - Returns - ------- - - denoised_coord: Denoised updated coord with shape [nframes, nloc, 3]. - """ - # [nframes, nloc, nnei, 1] - logits = self.lm_head(features, masked_tokens=masked_tokens) - if not isinstance(self.attn_head, list): - attn_probs = self.pair2coord_proj(pair_weights) - out_coord = (attn_probs * diff).sum(dim=-2) / ( - sw.sum(dim=-1).unsqueeze(-1) + 1e-6 - ) - else: - assert len(self.prefactor) == self.ndescriptor - all_coord_update = [] - assert len(pair_weights) == len(diff) == len(nlist_mask) == self.ndescriptor - for ii in range(self.ndescriptor): - _attn_probs = self.pair2coord_proj[ii](pair_weights[ii]) - _coord_update = (_attn_probs * diff[ii]).sum(dim=-2) / ( - nlist_mask[ii].sum(dim=-1).unsqueeze(-1) + 1e-6 - ) - all_coord_update.append(_coord_update) - out_coord = self.prefactor[0] * all_coord_update[0] - for ii in range(self.ndescriptor - 1): - out_coord += self.prefactor[ii + 1] * all_coord_update[ii + 1] - return { - "updated_coord": out_coord, - "logits": logits, - } From f334119148fe8e74c3b8f212b1f3f1256523932d Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 13:49:56 +0000 Subject: [PATCH 20/31] delete kwargs in property head --- deepmd/pt/model/task/property.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 9f3a7a2373..6adacbb464 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -55,7 +55,6 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, - **kwargs, ): self.task_num = task_num super().__init__( @@ -71,7 +70,6 @@ def __init__( activation_function=activation_function, precision=precision, mixed_types=mixed_types, - **kwargs, ) @classmethod From 479b106c405d13f7b11e32488adcff266309dcc7 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Tue, 19 Mar 2024 14:11:32 +0000 Subject: [PATCH 21/31] fix pre-commit --- deepmd/pt/model/model/property_model.py | 6 ++++-- deepmd/pt/model/task/property.py | 2 ++ examples/property/train/input_torch.json | 2 +- 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/model/property_model.py b/deepmd/pt/model/model/property_model.py index 74b258013b..8ac62e6106 100644 --- a/deepmd/pt/model/model/property_model.py +++ b/deepmd/pt/model/model/property_model.py @@ -40,7 +40,8 @@ def forward( ) model_predict = {} model_predict["atom_property"] = model_ret["property"] - model_predict["property"] = model_ret["property_redu"] / atype.shape[-1] + natoms = model_predict["atom_property"].shape[1] + model_predict["property"] = model_ret["property_redu"] / natoms if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict @@ -67,7 +68,8 @@ def forward_lower( ) model_predict = {} model_predict["atom_property"] = model_ret["property"] - model_predict["property"] = model_ret["property_redu"] / atype.shape[-1] + natoms = model_predict["atom_property"].shape[1] + model_predict["property"] = model_ret["property_redu"] / natoms if "mask" in model_ret: model_predict["mask"] = model_ret["mask"] return model_predict diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 6adacbb464..9f3a7a2373 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -55,6 +55,7 @@ def __init__( activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, + **kwargs, ): self.task_num = task_num super().__init__( @@ -70,6 +71,7 @@ def __init__( activation_function=activation_function, precision=precision, mixed_types=mixed_types, + **kwargs, ) @classmethod diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json index 915edb46a1..6f914cb905 100644 --- a/examples/property/train/input_torch.json +++ b/examples/property/train/input_torch.json @@ -192,7 +192,7 @@ "../data/3", "../data/6" ], - "batch_size": 1, + "batch_size": 2, "_comment": "that's all" }, "validation_data": { From 5cffff9d1909c44aa94f784179c355641ae2dcb9 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 06:20:47 +0000 Subject: [PATCH 22/31] task_num->task_dim --- deepmd/dpmodel/fitting/property_fitting.py | 8 ++++---- deepmd/pt/loss/property.py | 12 ++++++------ deepmd/pt/model/task/property.py | 8 ++++---- deepmd/pt/train/training.py | 4 ++-- deepmd/utils/argcheck.py | 4 ++-- examples/property/train/input_torch.json | 2 +- source/tests/pt/model/test_property_fitting.py | 12 ++++++------ 7 files changed, 25 insertions(+), 25 deletions(-) diff --git a/deepmd/dpmodel/fitting/property_fitting.py b/deepmd/dpmodel/fitting/property_fitting.py index 056a1af861..725108688d 100644 --- a/deepmd/dpmodel/fitting/property_fitting.py +++ b/deepmd/dpmodel/fitting/property_fitting.py @@ -39,7 +39,7 @@ def __init__( self, ntypes: int, dim_descrpt: int, - task_num: int = 1, + task_dim: int = 1, neuron: List[int] = [128, 128, 128], resnet_dt: bool = True, numb_fparam: int = 0, @@ -58,12 +58,12 @@ def __init__( # not used seed: Optional[int] = None, ): - self.task_num = task_num + self.task_dim = task_dim super().__init__( var_name="property", ntypes=ntypes, dim_descrpt=dim_descrpt, - dim_out=task_num, + dim_out=task_dim, neuron=neuron, resnet_dt=resnet_dt, numb_fparam=numb_fparam, @@ -91,7 +91,7 @@ def deserialize(cls, data: dict) -> "GeneralFitting": def serialize(self) -> dict: """Serialize the fitting to dict.""" - return {**super().serialize(), "type": "property", "task_num": self.task_num} + return {**super().serialize(), "type": "property", "task_dim": self.task_dim} def output_def(self) -> FittingOutputDef: return FittingOutputDef( diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index e6e82e3f90..c088bcf1ef 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -23,7 +23,7 @@ class PropertyLoss(TaskLoss): def __init__( self, - task_num, + task_dim, loss_func: str = "smooth_mae", metric: list = ["mae"], **kwargs, @@ -32,7 +32,7 @@ def __init__( Parameters ---------- - task_num : float + task_dim : float The learning rate at the start of the training. loss_func : str The loss function, such as "smooth_mae", "mae", "rmse" @@ -44,7 +44,7 @@ def __init__( super().__init__() self.loss_func = loss_func self.metric = metric - self.task_num = task_num + self.task_dim = task_dim self.mean = kwargs.get("mean", 0) self.std = kwargs.get("std", 1) self.beta = kwargs.get("beta", 1.00) @@ -68,8 +68,8 @@ def forward(self, model_pred, label, natoms, learning_rate, mae=False): more_loss: dict[str, torch.Tensor] Other losses for display. """ - assert label["property"].shape[-1] == self.task_num - assert model_pred["property"].shape[-1] == self.task_num + assert label["property"].shape[-1] == self.task_dim + assert model_pred["property"].shape[-1] == self.task_dim loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] more_loss = {} @@ -149,7 +149,7 @@ def label_requirement(self) -> List[DataRequirementItem]: label_requirement.append( DataRequirementItem( "property", - ndof=self.task_num, + ndof=self.task_dim, atomic=False, must=False, high_prec=True, diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 9f3a7a2373..85873abc0a 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -46,7 +46,7 @@ def __init__( self, ntypes: int, dim_descrpt: int, - task_num: int = 1, + task_dim: int = 1, neuron: List[int] = [128, 128, 128], bias_atom_e: Optional[torch.Tensor] = None, resnet_dt: bool = True, @@ -57,12 +57,12 @@ def __init__( mixed_types: bool = True, **kwargs, ): - self.task_num = task_num + self.task_dim = task_dim super().__init__( var_name="property", ntypes=ntypes, dim_descrpt=dim_descrpt, - dim_out=task_num, + dim_out=task_dim, neuron=neuron, bias_atom_e=bias_atom_e, resnet_dt=resnet_dt, @@ -84,7 +84,7 @@ def deserialize(cls, data: dict) -> "GeneralFitting": def serialize(self) -> dict: """Serialize the fitting to dict.""" - return {**super().serialize(), "type": "property", "task_num": self.task_num} + return {**super().serialize(), "type": "property", "task_dim": self.task_dim} def output_def(self) -> FittingOutputDef: return FittingOutputDef( diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index b061a986e8..bd0e6b38f3 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -299,8 +299,8 @@ def get_loss(loss_params, start_lr, _ntypes, _model): loss_params["label_name"] = label_name return TensorLoss(**loss_params) elif loss_type == "property": - task_num = _model.model_output_def()["property"].output_size - loss_params["task_num"] = task_num + task_dim = _model.model_output_def()["property"].output_size + loss_params["task_dim"] = task_dim return PropertyLoss(**loss_params) else: raise NotImplementedError diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 1a9421161b..165ef61732 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1183,7 +1183,7 @@ def fitting_property(): doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection' doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision." doc_seed = "Random seed for parameter initialization of the fitting net" - doc_task_num = "The dimension of outputs of fitting net" + doc_task_dim = "The dimension of outputs of fitting net" return [ Argument( "neuron", @@ -1203,7 +1203,7 @@ def fitting_property(): Argument("resnet_dt", bool, optional=True, default=True, doc=doc_resnet_dt), Argument("precision", str, optional=True, default="default", doc=doc_precision), Argument("seed", [int, None], optional=True, doc=doc_seed), - Argument("task_num", int, optional=True, default=1, doc=doc_task_num), + Argument("task_dim", int, optional=True, default=1, doc=doc_task_dim), ] diff --git a/examples/property/train/input_torch.json b/examples/property/train/input_torch.json index 6f914cb905..bc5ca389e0 100644 --- a/examples/property/train/input_torch.json +++ b/examples/property/train/input_torch.json @@ -157,7 +157,7 @@ }, "fitting_net": { "type": "property", - "task_num": 3, + "task_dim": 3, "neuron": [ 240, 240, diff --git a/source/tests/pt/model/test_property_fitting.py b/source/tests/pt/model/test_property_fitting.py index 00800a6251..a62f6ffa32 100644 --- a/source/tests/pt/model/test_property_fitting.py +++ b/source/tests/pt/model/test_property_fitting.py @@ -69,7 +69,7 @@ def test_consistency( ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num=5, + task_dim=5, numb_fparam=nfp, numb_aparam=nap, mixed_types=mixed_types, @@ -137,7 +137,7 @@ def test_jit( ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num=5, + task_dim=5, numb_fparam=nfp, numb_aparam=nap, mixed_types=mixed_types, @@ -178,7 +178,7 @@ def test_rot(self): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, # dim_descrpt - task_num=9, + task_dim=9, numb_fparam=nfp, numb_aparam=nap, mixed_types=True, @@ -230,7 +230,7 @@ def test_permu(self): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num=8, + task_dim=8, numb_fparam=0, numb_aparam=0, mixed_types=True, @@ -275,7 +275,7 @@ def test_trans(self): ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num=11, + task_dim=11, numb_fparam=0, numb_aparam=0, mixed_types=True, @@ -321,7 +321,7 @@ def setUp(self): self.ft0 = PropertyFittingNet( self.nt, self.dd0.dim_out, - task_num=3, + task_dim=3, numb_fparam=0, numb_aparam=0, mixed_types=True, From 5bda82e794a47771c4786c0b48a7d006961a0ab4 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 06:29:00 +0000 Subject: [PATCH 23/31] delete loss --- deepmd/pt/loss/__init__.py | 4 - deepmd/pt/loss/property.py | 158 ------------------------------------ deepmd/pt/train/training.py | 5 -- 3 files changed, 167 deletions(-) delete mode 100644 deepmd/pt/loss/property.py diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 0656db8dbe..9c8bbc9a2a 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -11,9 +11,6 @@ from .loss import ( TaskLoss, ) -from .property import ( - PropertyLoss, -) from .tensor import ( TensorLoss, ) @@ -24,5 +21,4 @@ "EnergySpinLoss", "TensorLoss", "TaskLoss", - "PropertyLoss", ] diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py deleted file mode 100644 index c088bcf1ef..0000000000 --- a/deepmd/pt/loss/property.py +++ /dev/null @@ -1,158 +0,0 @@ -# SPDX-License-Identifier: LGPL-3.0-or-later -import logging -from typing import ( - List, -) - -import torch -import torch.nn.functional as F - -from deepmd.pt.loss.loss import ( - TaskLoss, -) -from deepmd.pt.utils import ( - env, -) -from deepmd.utils.data import ( - DataRequirementItem, -) - -log = logging.getLogger(__name__) - - -class PropertyLoss(TaskLoss): - def __init__( - self, - task_dim, - loss_func: str = "smooth_mae", - metric: list = ["mae"], - **kwargs, - ): - r"""Construct a layer to compute loss on property. - - Parameters - ---------- - task_dim : float - The learning rate at the start of the training. - loss_func : str - The loss function, such as "smooth_mae", "mae", "rmse" - metric : list - The metric such as mae,rmse which will be printed. - **kwargs - Other keyword arguments. - """ - super().__init__() - self.loss_func = loss_func - self.metric = metric - self.task_dim = task_dim - self.mean = kwargs.get("mean", 0) - self.std = kwargs.get("std", 1) - self.beta = kwargs.get("beta", 1.00) - - def forward(self, model_pred, label, natoms, learning_rate, mae=False): - """Return loss on properties . - - Parameters - ---------- - model_pred : dict[str, torch.Tensor] - Model predictions. - label : dict[str, torch.Tensor] - Labels. - natoms : int - The local atom number. - - Returns - ------- - loss: torch.Tensor - Loss for model to minimize. - more_loss: dict[str, torch.Tensor] - Other losses for display. - """ - assert label["property"].shape[-1] == self.task_dim - assert model_pred["property"].shape[-1] == self.task_dim - loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] - more_loss = {} - - label_mean = torch.tensor( - self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE - ) - label_std = torch.tensor( - self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE - ) - - # loss - if self.loss_func == "smooth_mae": - loss += F.smooth_l1_loss( - (label["property"] - label_mean) / label_std, - model_pred["property"], - reduction="sum", - beta=self.beta, - ) - elif self.func == "mae": - loss += F.l1_loss( - (label["property"] - label_mean) / label_std, - model_pred["property"], - reduction="sum", - ) - elif self.func == "mse": - loss += F.mse_loss( - (label["property"] - label_mean) / label_std, - model_pred["property"], - reduction="sum", - ) - elif self.func == "rmse": - loss += torch.sqrt( - F.mse_loss( - (label["property"] - label_mean) / label_std, - model_pred["property"], - reduction="mean", - ) - ) - else: - raise RuntimeError(f"Unknown loss function : {self.func}") - - # more loss - if "smooth_mae" in self.metric: - more_loss["smooth_mae"] = F.smooth_l1_loss( - label["property"], - (model_pred["property"] * label_std) + label_mean, - reduction="mean", - beta=self.beta, - ).detach() - if "mae" in self.metric: - more_loss["mae"] = F.l1_loss( - label["property"], - (model_pred["property"] * label_std) + label_mean, - reduction="mean", - ).detach() - if "mse" in self.metric: - more_loss["mse"] = F.mse_loss( - label["property"], - (model_pred["property"] * label_std) + label_mean, - reduction="mean", - ).detach() - if "rmse" in self.metric: - more_loss["rmse"] = torch.sqrt( - F.mse_loss( - label["property"], - (model_pred["property"] * label_std) + label_mean, - reduction="mean", - ) - ).detach() - - return loss, more_loss - - @property - def label_requirement(self) -> List[DataRequirementItem]: - """Return data label requirements needed for this loss calculation.""" - label_requirement = [] - label_requirement.append( - DataRequirementItem( - "property", - ndof=self.task_dim, - atomic=False, - must=False, - high_prec=True, - ) - ) - return label_requirement diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index bd0e6b38f3..2056b9b305 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -27,7 +27,6 @@ DenoiseLoss, EnergySpinLoss, EnergyStdLoss, - PropertyLoss, TensorLoss, ) from deepmd.pt.model.model import ( @@ -298,10 +297,6 @@ def get_loss(loss_params, start_lr, _ntypes, _model): label_name = "polarizability" loss_params["label_name"] = label_name return TensorLoss(**loss_params) - elif loss_type == "property": - task_dim = _model.model_output_def()["property"].output_size - loss_params["task_dim"] = task_dim - return PropertyLoss(**loss_params) else: raise NotImplementedError From 00ec2560b98e811ce58ad3ee941e6b47b2c60922 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 16:38:42 +0000 Subject: [PATCH 24/31] Add property loss --- deepmd/pt/loss/__init__.py | 4 + deepmd/pt/loss/property.py | 163 ++++++++++++++++++++++++++++++++++++ deepmd/pt/train/training.py | 5 ++ 3 files changed, 172 insertions(+) create mode 100644 deepmd/pt/loss/property.py diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index 9c8bbc9a2a..c23df6b6d3 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -14,6 +14,9 @@ from .tensor import ( TensorLoss, ) +from .property import ( + PropertyLoss, +) __all__ = [ "DenoiseLoss", @@ -21,4 +24,5 @@ "EnergySpinLoss", "TensorLoss", "TaskLoss", + "PropertyLoss", ] diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py new file mode 100644 index 0000000000..2e5cd102a8 --- /dev/null +++ b/deepmd/pt/loss/property.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import logging +from typing import ( + List, +) + +import torch +import torch.nn.functional as F + +from deepmd.pt.loss.loss import ( + TaskLoss, +) +from deepmd.pt.utils import ( + env, +) +from deepmd.utils.data import ( + DataRequirementItem, +) + +log = logging.getLogger(__name__) + + +class PropertyLoss(TaskLoss): + def __init__( + self, + task_dim, + loss_func: str = "smooth_mae", + metric: list = ["mae"], + **kwargs, + ): + r"""Construct a layer to compute loss on property. + + Parameters + ---------- + task_dim : float + The learning rate at the start of the training. + loss_func : str + The loss function, such as "smooth_mae", "mae", "rmse" + metric : list + The metric such as mae,rmse which will be printed. + **kwargs + Other keyword arguments. + """ + super().__init__() + self.loss_func = loss_func + self.metric = metric + self.task_dim = task_dim + self.mean = kwargs.get("mean", 0) + self.std = kwargs.get("std", 1) + self.beta = kwargs.get("beta", 1.00) + + def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): + """Return loss on properties . + + Parameters + ---------- + input_dict : dict[str, torch.Tensor] + Model inputs. + model : torch.nn.Module + Model to be used to output the predictions. + label : dict[str, torch.Tensor] + Labels. + natoms : int + The local atom number. + + Returns + ------- + model_pred: dict[str, torch.Tensor] + Model predictions. + loss: torch.Tensor + Loss for model to minimize. + more_loss: dict[str, torch.Tensor] + Other losses for display. + """ + model_pred = model(**input_dict) + assert label["property"].shape[-1] == self.task_dim + assert model_pred["property"].shape[-1] == self.task_dim + loss = torch.zeros(1, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE)[0] + more_loss = {} + + label_mean = torch.tensor( + self.mean, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + label_std = torch.tensor( + self.std, dtype=env.GLOBAL_PT_FLOAT_PRECISION, device=env.DEVICE + ) + + # loss + if self.loss_func == "smooth_mae": + loss += F.smooth_l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + beta=self.beta, + ) + elif self.func == "mae": + loss += F.l1_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) + elif self.func == "mse": + loss += F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="sum", + ) + elif self.func == "rmse": + loss += torch.sqrt( + F.mse_loss( + (label["property"] - label_mean) / label_std, + model_pred["property"], + reduction="mean", + ) + ) + else: + raise RuntimeError(f"Unknown loss function : {self.func}") + + # more loss + if "smooth_mae" in self.metric: + more_loss["smooth_mae"] = F.smooth_l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + beta=self.beta, + ).detach() + if "mae" in self.metric: + more_loss["mae"] = F.l1_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() + if "mse" in self.metric: + more_loss["mse"] = F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ).detach() + if "rmse" in self.metric: + more_loss["rmse"] = torch.sqrt( + F.mse_loss( + label["property"], + (model_pred["property"] * label_std) + label_mean, + reduction="mean", + ) + ).detach() + + return model_pred, loss, more_loss + + @property + def label_requirement(self) -> List[DataRequirementItem]: + """Return data label requirements needed for this loss calculation.""" + label_requirement = [] + label_requirement.append( + DataRequirementItem( + "property", + ndof=self.task_dim, + atomic=False, + must=False, + high_prec=True, + ) + ) + return label_requirement \ No newline at end of file diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 9fd675a8f2..2899638f88 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -28,6 +28,7 @@ EnergySpinLoss, EnergyStdLoss, TensorLoss, + PropertyLoss, ) from deepmd.pt.model.model import ( DPZBLModel, @@ -297,6 +298,10 @@ def get_loss(loss_params, start_lr, _ntypes, _model): label_name = "polarizability" loss_params["label_name"] = label_name return TensorLoss(**loss_params) + elif loss_type == "property": + task_dim = _model.model_output_def()["property"].output_size + loss_params["task_dim"] = task_dim + return PropertyLoss(**loss_params) else: raise NotImplementedError From 6396fb39da57739c7c9c83a4a7578fff62ab59b0 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 20 Mar 2024 16:39:13 +0000 Subject: [PATCH 25/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/loss/__init__.py | 6 +++--- deepmd/pt/loss/property.py | 2 +- deepmd/pt/train/training.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/deepmd/pt/loss/__init__.py b/deepmd/pt/loss/__init__.py index c23df6b6d3..0656db8dbe 100644 --- a/deepmd/pt/loss/__init__.py +++ b/deepmd/pt/loss/__init__.py @@ -11,12 +11,12 @@ from .loss import ( TaskLoss, ) -from .tensor import ( - TensorLoss, -) from .property import ( PropertyLoss, ) +from .tensor import ( + TensorLoss, +) __all__ = [ "DenoiseLoss", diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 2e5cd102a8..70384a3d7c 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -160,4 +160,4 @@ def label_requirement(self) -> List[DataRequirementItem]: high_prec=True, ) ) - return label_requirement \ No newline at end of file + return label_requirement diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py index 2899638f88..e55803954d 100644 --- a/deepmd/pt/train/training.py +++ b/deepmd/pt/train/training.py @@ -27,8 +27,8 @@ DenoiseLoss, EnergySpinLoss, EnergyStdLoss, - TensorLoss, PropertyLoss, + TensorLoss, ) from deepmd.pt.model.model import ( DPZBLModel, From 698ef2173bb864640db7137adf797587ef9cc2a4 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 16:45:08 +0000 Subject: [PATCH 26/31] fix eval --- deepmd/infer/deep_property.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/infer/deep_property.py b/deepmd/infer/deep_property.py index e3b1866eb5..ac88f0dd9c 100644 --- a/deepmd/infer/deep_property.py +++ b/deepmd/infer/deep_property.py @@ -119,7 +119,7 @@ def eval( **kwargs, ) atomic_property = results["property"].reshape(nframes, natoms, -1) - property = np.sum(atomic_property, axis=1) + property = results["property_redu"].reshape(nframes, -1) if atomic: return ( From 5fafabbac204217365dc091b118db6b0939314fa Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 16:54:16 +0000 Subject: [PATCH 27/31] resolve conversation --- deepmd/pt/model/task/property.py | 38 ++++++++++++++++++++++++++++++-- deepmd/utils/argcheck.py | 2 +- 2 files changed, 37 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index 85873abc0a..e25860453f 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -42,19 +42,52 @@ @Fitting.register("property") class PropertyFittingNet(InvarFitting): + """Construct a fitting net for energy. + + Parameters + ---------- + ntypes : int + Element count. + dim_descrpt : int + Embedding width per atom. + task_dim : int + The dimension of outputs of fitting net + neuron : List[int] + Number of neurons in each hidden layers of the fitting net. + bias_atom_p : torch.Tensor, optional + Average property 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. + mixed_types : bool + If true, use a uniform fitting net for all atom types, otherwise use + different fitting nets for different atom types. + seed : int, optional + Random seed. + + """ + def __init__( self, ntypes: int, dim_descrpt: int, task_dim: int = 1, neuron: List[int] = [128, 128, 128], - bias_atom_e: Optional[torch.Tensor] = None, + bias_atom_p: Optional[torch.Tensor] = None, resnet_dt: bool = True, numb_fparam: int = 0, numb_aparam: int = 0, activation_function: str = "tanh", precision: str = DEFAULT_PRECISION, mixed_types: bool = True, + seed: Optional[int] = None, **kwargs, ): self.task_dim = task_dim @@ -64,13 +97,14 @@ def __init__( dim_descrpt=dim_descrpt, dim_out=task_dim, neuron=neuron, - bias_atom_e=bias_atom_e, + bias_atom_e=bias_atom_p, resnet_dt=resnet_dt, numb_fparam=numb_fparam, numb_aparam=numb_aparam, activation_function=activation_function, precision=precision, mixed_types=mixed_types, + seed=seed, **kwargs, ) diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py index 415173a419..42ab8b8a00 100644 --- a/deepmd/utils/argcheck.py +++ b/deepmd/utils/argcheck.py @@ -1181,7 +1181,7 @@ def fitting_dipole(): ] -@fitting_args_plugin.register("property") +@fitting_args_plugin.register("property", doc=doc_only_pt_supported) def fitting_property(): doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built." doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.' From c76e23eb233361c554408359d9bb2418904850e6 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 20 Mar 2024 16:54:55 +0000 Subject: [PATCH 28/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/task/property.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/model/task/property.py b/deepmd/pt/model/task/property.py index e25860453f..b4c18ad378 100644 --- a/deepmd/pt/model/task/property.py +++ b/deepmd/pt/model/task/property.py @@ -71,9 +71,9 @@ class PropertyFittingNet(InvarFitting): different fitting nets for different atom types. seed : int, optional Random seed. - + """ - + def __init__( self, ntypes: int, From 97bd86f10459a0aa3f66ec92796a5bafa7d27092 Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Wed, 20 Mar 2024 16:57:13 +0000 Subject: [PATCH 29/31] Add example to tests --- source/tests/common/test_examples.py | 1 + 1 file changed, 1 insertion(+) diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index 91bb9c0174..f759440a52 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -49,6 +49,7 @@ p_examples / "water" / "se_e2_a" / "input_torch.json", p_examples / "water" / "se_atten" / "input_torch.json", p_examples / "water" / "dpa2" / "input_torch.json", + p_examples / "property" / "train"/ "input_torch.json", ) From cbb9c4bace01d47bf6c901beea8cd72853815aa9 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 20 Mar 2024 16:58:11 +0000 Subject: [PATCH 30/31] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- source/tests/common/test_examples.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/source/tests/common/test_examples.py b/source/tests/common/test_examples.py index f759440a52..9366e707fa 100644 --- a/source/tests/common/test_examples.py +++ b/source/tests/common/test_examples.py @@ -49,7 +49,7 @@ p_examples / "water" / "se_e2_a" / "input_torch.json", p_examples / "water" / "se_atten" / "input_torch.json", p_examples / "water" / "dpa2" / "input_torch.json", - p_examples / "property" / "train"/ "input_torch.json", + p_examples / "property" / "train" / "input_torch.json", ) From 18182717baa30566bc988ca65182b5e75d3eae1b Mon Sep 17 00:00:00 2001 From: Chengqian-Zhang <2000011006@stu.pku.edu.cn> Date: Fri, 22 Mar 2024 10:02:27 +0000 Subject: [PATCH 31/31] fix bug of loss_func --- deepmd/pt/loss/property.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deepmd/pt/loss/property.py b/deepmd/pt/loss/property.py index 70384a3d7c..a03dc82173 100644 --- a/deepmd/pt/loss/property.py +++ b/deepmd/pt/loss/property.py @@ -93,19 +93,19 @@ def forward(self, input_dict, model, label, natoms, learning_rate, mae=False): reduction="sum", beta=self.beta, ) - elif self.func == "mae": + elif self.loss_func == "mae": loss += F.l1_loss( (label["property"] - label_mean) / label_std, model_pred["property"], reduction="sum", ) - elif self.func == "mse": + elif self.loss_func == "mse": loss += F.mse_loss( (label["property"] - label_mean) / label_std, model_pred["property"], reduction="sum", ) - elif self.func == "rmse": + elif self.loss_func == "rmse": loss += torch.sqrt( F.mse_loss( (label["property"] - label_mean) / label_std,