From 4d3602514fe8d2321c5b7645dc686198f32a1213 Mon Sep 17 00:00:00 2001 From: zhangxinfeng3 Date: Mon, 7 Dec 2020 20:46:10 +0800 Subject: [PATCH] update force calculation in MD --- mindspore/ops/_op_impl/akg/ascend/__init__.py | 1 + .../_op_impl/akg/ascend/prod_force_se_a.py | 33 +++++++++++++++++++ mindspore/ops/operations/__init__.py | 3 +- mindspore/ops/operations/_thor_ops.py | 19 +++++++++++ .../research/hpc/molecular_dynamics/README.md | 19 ++++++----- .../research/hpc/molecular_dynamics/eval.py | 14 ++++---- .../hpc/molecular_dynamics/src/__init__.py | 0 .../hpc/molecular_dynamics/src/network.py | 13 +++++--- 8 files changed, 81 insertions(+), 21 deletions(-) create mode 100644 mindspore/ops/_op_impl/akg/ascend/prod_force_se_a.py create mode 100644 model_zoo/research/hpc/molecular_dynamics/src/__init__.py diff --git a/mindspore/ops/_op_impl/akg/ascend/__init__.py b/mindspore/ops/_op_impl/akg/ascend/__init__.py index d1029e0bcb..61e9dea9db 100644 --- a/mindspore/ops/_op_impl/akg/ascend/__init__.py +++ b/mindspore/ops/_op_impl/akg/ascend/__init__.py @@ -43,5 +43,6 @@ from .select import _select_akg from .sqrt import _sqrt_akg from .square import _square_akg from .sub import _sub_akg +from .prod_force_se_a import _prod_force_se_a_akg # Please insert op register in lexicographical order of the filename. diff --git a/mindspore/ops/_op_impl/akg/ascend/prod_force_se_a.py b/mindspore/ops/_op_impl/akg/ascend/prod_force_se_a.py new file mode 100644 index 0000000000..4e0adbdcd7 --- /dev/null +++ b/mindspore/ops/_op_impl/akg/ascend/prod_force_se_a.py @@ -0,0 +1,33 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""ProdForceSeA op""" +from mindspore.ops.op_info_register import op_info_register, AkgAscendRegOp, DataType as DT + +op_info = AkgAscendRegOp("ProdForceSeA") \ + .fusion_type("ELEMWISE") \ + .attr("natoms", "required", "int") \ + .input(0, "net_deriv_tensor") \ + .input(1, "in_deriv_tensor") \ + .input(2, "nlist_tensor") \ + .output(0, "output") \ + .dtype_format(DT.F32_Default, DT.F32_Default, DT.I32_Default, DT.F32_Default) \ + .get_op_info() + + +@op_info_register(op_info) +def _prod_force_se_a_akg(): + """ProdForceSeA Akg register""" + return diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 63bdf01b67..9acd289e43 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -88,7 +88,8 @@ from .other_ops import (Assign, InplaceAssign, IOU, BoundingBoxDecode, BoundingB from ._thor_ops import (CusBatchMatMul, CusCholeskyTrsm, CusFusedAbsMax1, CusImg2Col, CusMatMulCubeDenseLeft, CusMatMulCubeFraczRightMul, CusMatMulCube, CusMatrixCombine, CusTranspose02314, CusMatMulCubeDenseRight, - CusMatMulCubeFraczLeftCast, Im2Col, UpdateThorGradient, Cholesky, CholeskyTrsm, DetTriangle) + CusMatMulCubeFraczLeftCast, Im2Col, UpdateThorGradient, Cholesky, CholeskyTrsm, DetTriangle, + ProdForceSeA) from .sparse_ops import SparseToDense from ._cache_ops import CacheSwapHashmap, SearchCacheIdx, CacheSwapTable, UpdateCache, MapCacheIdx, SubAndFilter diff --git a/mindspore/ops/operations/_thor_ops.py b/mindspore/ops/operations/_thor_ops.py index 71edf6dcaa..77cd0ca81e 100644 --- a/mindspore/ops/operations/_thor_ops.py +++ b/mindspore/ops/operations/_thor_ops.py @@ -709,3 +709,22 @@ class DetTriangle(PrimitiveWithInfer): def infer_dtype(self, x1_dtype): validator.check_tensor_dtype_valid('x1', x1_dtype, [mstype.float32], self.name) return x1_dtype + + +class ProdForceSeA(PrimitiveWithInfer): + """ + ProdForceSeA. + """ + + @prim_attr_register + def __init__(self, natoms=192): + self.init_prim_io_names(inputs=['net_deriv_tensor', "in_deriv_tensor", "nlist_tensor"], outputs=['y']) + self.natoms = natoms + self.add_prim_attr('natoms', self.natoms) + + def infer_shape(self, x1_shape, x2_shape, x3_shape): + out_shape = [x3_shape[0], x3_shape[1], 3] + return out_shape + + def infer_dtype(self, x1_dtype, x2_dtype, x3_dtype): + return x1_dtype diff --git a/model_zoo/research/hpc/molecular_dynamics/README.md b/model_zoo/research/hpc/molecular_dynamics/README.md index ce6c436466..3c787b5e31 100644 --- a/model_zoo/research/hpc/molecular_dynamics/README.md +++ b/model_zoo/research/hpc/molecular_dynamics/README.md @@ -31,9 +31,9 @@ The overall network architecture of MD simulation is show below. Dataset used: deepmodeling/deepmd-kit/examples/water/data -The data is generated by Quantum Espresso and the input of Quantum Espresso is setted manually. +The data is generated by Quantum Espresso and the input of Quantum Espresso is set manually. -The directory structure of the data is as follows: +The directory structure of the dataset is as follows: ```text └─data @@ -50,10 +50,10 @@ The directory structure of the data is as follows: In `deepmodeling/deepmd-kit/source`: -- Use `train/DataSystem.py` to get coord and atype. +- Use `train/DataSystem.py` to get d_coord and atype. - Use function compute_input_stats in `train/DataSystem.py` to get avg and std. -- Use `op/descrpt_se_a.cc` to get nlist. -- Save coord, atype, avg, std and nlist as `Npz` file for infer. +- Use `op/descrpt_se_a.cc` to get d_nlist and nlist. +- Save d_coord, d_nlist, atype, avg, std and nlist as `Npz` file for inference. ## Environment Requirements @@ -96,12 +96,13 @@ python eval.py --dataset_path [DATASET_PATH] --checkpoint_path [CHECKPOINT_PATH] ### Result -推理的结果如下： +The infer result： ```text -atom_ener: -94.38766 -94.294426 -94.39194 -94.70758 -94.51311 -94.457954 ... -force: 1.64911175 -1.09822524 0.46055657 -1.34915102 -0.33827361 -0.97184098 ... -virial: -11.736662 -4.286214 2.8852937 -4.286209 -10.408775 -5.6738234 ... +energy: -29944.03 +atom_energy: -94.38766 -94.294426 -94.39194 -94.70758 -94.51311 -94.457954 ... +force: 1.649112 -1.0982257 0.46055675 -1.3491507 -0.3382736 -0.97184074 ... +virial: -11.736662 -4.2862144 2.8852935 -4.286209 -10.408775 -5.6738224 ... ``` ## ModelZoo Homepage diff --git a/model_zoo/research/hpc/molecular_dynamics/eval.py b/model_zoo/research/hpc/molecular_dynamics/eval.py index 4f156706ef..cc143e919b 100644 --- a/model_zoo/research/hpc/molecular_dynamics/eval.py +++ b/model_zoo/research/hpc/molecular_dynamics/eval.py @@ -32,11 +32,12 @@ context.set_context(mode=context.GRAPH_MODE, save_graphs=False, device_target="A if __name__ == '__main__': # get input data r = np.load(args_opt.dataset_path) - d_coord, d_nlist, avg, std, atype = r['d_coord'], r['d_nlist'], r['avg'], r['std'], r['atype'] + d_coord, d_nlist, avg, std, atype, nlist = r['d_coord'], r['d_nlist'], r['avg'], r['std'], r['atype'], r['nlist'] batch_size = 1 atype_tensor = Tensor(atype) avg_tensor = Tensor(avg) std_tensor = Tensor(std) + nlist_tensor = Tensor(nlist) d_coord_tensor = Tensor(np.reshape(d_coord, (1, -1, 3))) d_nlist_tensor = Tensor(d_nlist) frames = [] @@ -48,8 +49,9 @@ if __name__ == '__main__': param_dict = load_checkpoint(args_opt.checkpoint_path) load_param_into_net(net, param_dict) net.to_float(mstype.float32) - energy, atom_ener, virial = \ - net(d_coord_tensor, d_nlist_tensor, frames, avg_tensor, std_tensor, atype_tensor) - print(energy) - print(atom_ener) - print(virial) + energy, atom_ener, force, virial = \ + net(d_coord_tensor, d_nlist_tensor, frames, avg_tensor, std_tensor, atype_tensor, nlist_tensor) + print('energy:', energy) + print('atom_energy:', atom_ener) + print('force:', force) + print('virial:', virial) diff --git a/model_zoo/research/hpc/molecular_dynamics/src/__init__.py b/model_zoo/research/hpc/molecular_dynamics/src/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/model_zoo/research/hpc/molecular_dynamics/src/network.py b/model_zoo/research/hpc/molecular_dynamics/src/network.py index 841cd2a3a0..885842562c 100644 --- a/model_zoo/research/hpc/molecular_dynamics/src/network.py +++ b/model_zoo/research/hpc/molecular_dynamics/src/network.py @@ -22,8 +22,8 @@ from mindspore import Tensor from mindspore import nn from mindspore.ops import composite as C from mindspore.ops import operations as P -from virial import ProdVirialSeA -from descriptor import DescriptorSeA +from .virial import ProdVirialSeA +from .descriptor import DescriptorSeA natoms = [192, 192, 64, 128] rcut_a = -1 @@ -231,11 +231,12 @@ class Network(nn.Cell): self.descrpt_se_a = DescriptorSeA() self.process = Processing() self.prod_virial_se_a = ProdVirialSeA() + self.prod_force_se_a = P.ProdForceSeA() - def construct(self, coord, nlist, frames, avg, std, atype): + def construct(self, d_coord, d_nlist, frames, avg, std, atype, nlist): """construct function.""" rij, descrpt, descrpt_deriv = \ - self.descrpt_se_a(coord, nlist, frames, avg, std, atype) + self.descrpt_se_a(d_coord, d_nlist, frames, avg, std, atype) # calculate energy and atom_ener atom_ener = self.mdnet(descrpt) energy_raw = atom_ener @@ -247,4 +248,6 @@ class Network(nn.Cell): descrpt_deriv_reshape = self.reshape(descrpt_deriv, (-1, natoms[0], ndescrpt, 3)) # calculate virial virial = self.prod_virial_se_a(net_deriv_reshape, descrpt_deriv_reshape, rij, nlist) - return energy, atom_ener, virial + # calculate force + force = self.prod_force_se_a(net_deriv_reshape, descrpt_deriv_reshape, nlist) + return energy, atom_ener, force, virial