From 0000512a247d6f02d194d22d1c29108072906480 Mon Sep 17 00:00:00 2001 From: HydrogenSulfate <490868991@qq.com> Date: Mon, 27 Nov 2023 05:22:30 +0000 Subject: [PATCH] add unitest code and fix for custom op installation in python --- .../{setup_ins.py => custom_op_install.py} | 6 +- source/lib/paddle_src/custom_op_test.py | 272 ++++++++++++++++++ 2 files changed, 276 insertions(+), 2 deletions(-) rename source/lib/paddle_src/{setup_ins.py => custom_op_install.py} (91%) create mode 100644 source/lib/paddle_src/custom_op_test.py diff --git a/source/lib/paddle_src/setup_ins.py b/source/lib/paddle_src/custom_op_install.py similarity index 91% rename from source/lib/paddle_src/setup_ins.py rename to source/lib/paddle_src/custom_op_install.py index 402b47cbc5..be5b88e36e 100644 --- a/source/lib/paddle_src/setup_ins.py +++ b/source/lib/paddle_src/custom_op_install.py @@ -1,5 +1,7 @@ from paddle.utils import cpp_extension +# NOTE: 请在本文件的目录下执行：python custom_op_install.py install，以安装自定义算子 + cpp_extension.setup( name="paddle_deepmd_lib", ext_modules=cpp_extension.CUDAExtension( @@ -50,9 +52,9 @@ "./paddle_neighbor_stat.cc", ], include_dirs=[ - "/workspace/hesensen/deepmd_backend/deepmd-kit-tf/source/lib/include" + "../../lib/include/", ], library_dirs=["/usr/local/cuda-11/lib64"], - # extra_link_args=["-ldeepmd"] + define_macros=[("GOOGLE_CUDA", "1")], ), ) diff --git a/source/lib/paddle_src/custom_op_test.py b/source/lib/paddle_src/custom_op_test.py new file mode 100644 index 0000000000..18403c22c5 --- /dev/null +++ b/source/lib/paddle_src/custom_op_test.py @@ -0,0 +1,272 @@ +import os +from os import path as osp + +import paddle +import paddle_deepmd_lib + +unitest_dir = os.getenv("UNITTEST_DIR", None) + +if unitest_dir is None: + raise ValueError( + "Please download unitest data and set env with 4 scipts below:\n" + "1. wget -nc https://paddle-org.bj.bcebos.com/paddlescience/deepmd/deepmd_custom_op_test_data.tar\n" + "2. tar -xf deepmd_custom_op_test_data.tar\n" + "3. export UNITTEST_DIR=$PWD/deepmd_custom_op_test_data\n" + "4. python ./custom_op_test.py\n" + ) + + +def test_neighbor_stat(place="cpu"): + print("=" * 10, f"test_neighbor_stat [place={place}]", "=" * 10) + import numpy as np + + coord = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "neighbor_stat/coord.npy")) + ) + type = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "neighbor_stat/type.npy")) + ) + natoms = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "neighbor_stat/natoms_vec.npy")) + ) + box = np.ascontiguousarray(np.load(osp.join(unitest_dir, "neighbor_stat/box.npy"))) + default_mesh = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "neighbor_stat/default_mesh.npy")) + ) + + rcut = 6.0 + + coord = paddle.to_tensor(coord, "float32", place=place) + type = paddle.to_tensor(type, "int32", place=place) + natoms = paddle.to_tensor(natoms, "int64", place=place) + box = paddle.to_tensor(box, "float32", place=place) + default_mesh = paddle.to_tensor(default_mesh, "int32", place=place) + + mn, dt = paddle_deepmd_lib.neighbor_stat( + coord, + type, + natoms, + box, + default_mesh, + rcut=rcut, + ) + + mn_load = np.load(osp.join(unitest_dir, "neighbor_stat/mn.npy")) + dt_load = np.load(osp.join(unitest_dir, "neighbor_stat/dt.npy")) + + # print(mn.shape, mn.min().item()); print(mn.max().item()); print(mn.mean().item()); print(mn.var().item()) + # print(mn_load.shape); print(mn_load.min().item()); print(mn_load.max().item()); print(mn_load.mean().item()); print(mn_load.var().item()) + # print(dt.shape, dt.min().item(), dt.max().item(), dt.mean().item(), dt.var().item()) + # print(dt_load.shape, dt_load.min().item(), dt_load.max().item(), dt_load.mean().item(), dt_load.var().item()) + + print(np.allclose(mn.numpy(), mn_load, 1e-1)) + print(np.allclose(dt.numpy(), dt_load, 1e-1)) + + +def test_prod_env_mat_a(place="cpu"): + print("=" * 10, f"test_prod_env_mat_a [place={place}]", "=" * 10) + import numpy as np + + # "coord", "atype", "natoms", "box", "mesh", "t_avg", "t_std" + coord = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.coord.npy")) + ) + atype = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.atype.npy")) + ) + natoms = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.natoms.npy")) + ) + box = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.box.npy")) + ) + mesh = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.mesh.npy")) + ) + t_avg = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.t_avg.npy")) + ) + t_std = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.t_std.npy")) + ) + t_std = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.t_std.npy")) + ) + + coord = paddle.to_tensor(coord, place=place) + atype = paddle.to_tensor(atype, place=place) + natoms = paddle.to_tensor(natoms, place="cpu") + box = paddle.to_tensor(box, place=place) + mesh = paddle.to_tensor(mesh, place=place) + t_avg = paddle.to_tensor(t_avg, place=place) + t_std = paddle.to_tensor(t_std, place=place) + + rcut_a = -1 + rcut_r = 6.0 + rcut_r_smth = 0.5 + sel_a = [46, 92] + sel_r = [0, 0] + + # print(coord.shape, coord.dtype, coord.place) + # print(atype.shape, atype.dtype, atype.place) + # print(box.shape, box.dtype, box.place) + # print(mesh.shape, mesh.dtype, mesh.place) + # print(t_avg.shape, t_avg.dtype, t_avg.place) + # print(t_std.shape, t_std.dtype, t_std.place) + # print(natoms.shape, natoms.dtype, natoms.place) + # print(rcut_a) + # print(rcut_r) + # print(rcut_r_smth) + # print(sel_a) + # print(sel_r) + + descrpt, descrpt_deriv, rij, nlist = paddle_deepmd_lib.prod_env_mat_a( + coord, + atype, + box, + mesh, + t_avg, + t_std, + natoms, + rcut_a, + rcut_r, + rcut_r_smth, + sel_a, + sel_r, + ) + descrpt_load = np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.descrpt.npy")) + descrpt_deriv_load = np.load( + osp.join(unitest_dir, "prod_env_mat_a/descrpt.descrpt_deriv.npy") + ) + rij_load = np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.rij.npy")) + nlist_load = np.load(osp.join(unitest_dir, "prod_env_mat_a/descrpt.nlist.npy")) + # print(descrpt.shape) + # print(descrpt_deriv.shape) + # print(rij.shape) + # print(nlist.shape) + # print(descrpt_load.shape) # (1, 576) + # print(descrpt_deriv_load.shape) # (1, 192) + # print(rij_load.shape) # (4,) + # print(nlist_load.shape) # (1, 9) + + print(np.allclose(descrpt.numpy(), descrpt_load)) + print(np.allclose(descrpt_deriv.numpy(), descrpt_deriv_load)) + print(np.allclose(rij.numpy(), rij_load)) + print(np.allclose(nlist.numpy(), nlist_load)) + + +def test_prod_force_se_a(place="cpu"): + print("=" * 10, f"test_prod_force_se_a [place={place}]", "=" * 10) + import numpy as np + + # "coord", "atype", "natoms", "box", "mesh", "t_avg", "t_std" + net_deriv_reshape = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_force_se_a/descrpt.net_deriv_reshape.npy")) + ) + descrpt_deriv = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_force_se_a/descrpt.descrpt_deriv.npy")) + ) + nlist = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_force_se_a/descrpt.nlist.npy")) + ) + natoms = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_force_se_a/descrpt.natoms.npy")) + ) + + nnei_a = 138 + nnei_r = 0 + + net_deriv_reshape = paddle.to_tensor( + net_deriv_reshape, stop_gradient=False, place=place + ) + descrpt_deriv = paddle.to_tensor(descrpt_deriv, place=place) + nlist = paddle.to_tensor(nlist, place=place) + natoms = paddle.to_tensor(natoms, place="cpu") # [192, 192, 64 , 128] + force = paddle_deepmd_lib.prod_force_se_a( + net_deriv_reshape, + descrpt_deriv, + nlist, + natoms, + n_a_sel=nnei_a, + n_r_sel=nnei_r, + ) + force.sum().backward() + # print(f"net_deriv_reshape.grad.shape = {net_deriv_reshape.grad.shape}") + + force_load = np.load(osp.join(unitest_dir, "prod_force_se_a/descrpt.force.npy")) + # print(force.shape) # (1, 9) + # print(force_load.shape) # (1, 9) + + print(np.allclose(force.numpy(), force_load)) + + +def test_prod_virial_se_a(place="cpu"): + print("=" * 10, f"test_prod_virial_se_a [place={place}]", "=" * 10) + import numpy as np + + # "coord", "atype", "natoms", "box", "mesh", "t_avg", "t_std" + net_deriv_reshape = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.net_deriv_reshape.npy")) + ) + descrpt_deriv = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.descrpt_deriv.npy")) + ) + rij = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.rij.npy")) + ) + nlist = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.nlist.npy")) + ) + natoms = np.ascontiguousarray( + np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.natoms.npy")) + ) + + nnei_a = 138 + nnei_r = 0 + + net_deriv_reshape = paddle.to_tensor( + net_deriv_reshape, stop_gradient=False, place=place + ) + descrpt_deriv = paddle.to_tensor(descrpt_deriv, place=place) + rij = paddle.to_tensor(rij, place=place) + nlist = paddle.to_tensor(nlist, place=place) + # natoms = natoms.tolist() + natoms = paddle.to_tensor(natoms, place="cpu") # [192, 192, 64 , 128] + # print(natoms.place) + # print(natoms.dtype) + # print(natoms) + virial, atom_virial = paddle_deepmd_lib.prod_virial_se_a( + net_deriv_reshape, + descrpt_deriv, + rij, + nlist, + natoms, + n_a_sel=nnei_a, + n_r_sel=nnei_r, + ) + virial.sum().backward() + # print(f"net_deriv_reshape.grad.shape = {net_deriv_reshape.grad.shape}") + + virial_load = np.load(osp.join(unitest_dir, "prod_virial_se_a/descrpt.virial.npy")) + atom_virial_load = np.load( + osp.join(unitest_dir, "prod_virial_se_a/descrpt.atom_virial.npy") + ) + # print(virial.shape) # (1, 9) + # print(virial_load.shape) # (1, 9) + # print(atom_virial.shape) # (1, 9) + # print(atom_virial_load.shape) # (1, 9) + + print(np.allclose(virial.numpy(), virial_load)) + print(np.allclose(atom_virial.numpy(), atom_virial_load)) + + +if __name__ == "__main__": + test_neighbor_stat() + + test_prod_env_mat_a("gpu") + test_prod_force_se_a("gpu") + test_prod_virial_se_a("gpu") + + test_prod_env_mat_a("cpu") + test_prod_force_se_a("cpu") + test_prod_virial_se_a("cpu")