DeePMD-kit实现对PaddlePaddle机器学习框架的支持

近期DeePMD-kit和开源深度学习框架飞桨（PaddlePaddle）的开发者们进行了跨领域深度合作，依托飞桨框架成熟的底层功能和动静统一的开发体验实现DeePMD算法的构建和运行，促进了双方在开发者生态上的双向融合。

DeePMD算法包括两部分，用神经网络深度学习方案来模拟材料势能面的DeepPot（深度势模型）和用深度势模型跑MD（分子动力学）。DeePMD算法中，神经网络深度学习部分采用的是Tensorflow静态图实现。这里，我们用DeePMD当中较为典型的SeA势模型为例来为大家介绍一下用飞桨实现DeePMD算法中深度学习方案的开发背景和过程。

1. 技术背景

SeA指DeepPot smooth edition constructed from all information including angle and radial information of atomic configurations。从原子坐标信息（包括原子间的相对径向坐标和角向坐标）出发，通过神经网络构造复杂材料体系的势能面。

SeA模型由两部分构成:

1. Embeding网络，输入中心原子i与其所有近邻原子j的相对坐标信息，输出Feature Matrix。

2. Fitting网络，输入Feature Matrix，输出中心原子i的局域能量。

DeepPot已被用于准确描述包括分子结构、各种周期性结构，以及表面、高熵、电池等等多种复杂结构。

2. 使用飞桨框架来实现SeA

飞桨依托于百度多年的深度学习技术研究和业务应用，是国内第一个自主研发、功能非常完备、并向开发者全面开源开放的产业级深度学习平台。飞桨采用动态图模式和自定义算子方案。接下来为大家介绍下用飞桨实现DeePMD的具体方法。

2.1 构建SeA

2.1.1 SeA算法结构

图1 DeepPot-SeA算法结构

1. Embedding网络：要想让随机性的神经网络能准确描述物理体系，Embedding输出的Feature Matrix应该满足体系的物理对称性，包括平移、置换和旋转对称性。

   因此，Embedding网络的构造基于以下两个数学原理：

   a. 矩阵满足体系的所有平移和旋转对称性；

   b. 所有满足置换对称性的函数形式可以变换成如下形式：，比如函数对于任意标量函数都满足置换对称性。

   基于以上两个原理，Feature Matrix的构造过程如上图：将weighting function矩阵输入Embedding网络（其中），输出维矩阵，并取的前列得到。最后通过矩阵相乘构造Feature Matrix：。

2. Fitting网络：将Embedding网络输出的Feature Matrix输入Fitting网络，输出i原子的局域能量E_i。

具体实现参考：

https://github.com/deepmodeling/deepmd-kit

2.1.2 使用飞桨框架动态图来构建SeA

飞桨框架动态图使用了可以支持更加贴近Python语法和执行逻辑的执行机制，使得算法实现可以按照顺序的执行逻辑来进行。

通过继承 paddle.nn.Layer，可以使用飞桨框架内部的各类操作网络的方法，具体参见：

https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/fluid/dygraph/layers/Layer_cn.html#layer

具体步骤如下：

在__init__方法中设置网络结构（包括层数和每层的维度）；

在forward()方法中实现加权求和+激活函数逻辑；

而深度学习中反向的计算逻辑，使用autograd技术自动完成，只要在整个网络完成后使用backward()方法执行反向网络。

class EnerModel(paddle.nn.Layer): model_type = 'ener' def __init__( self, descrpt, fitting, type_map: List[str] = None, data_stat_nbatch: int = 10, data_stat_protect: float = 1e-2, use_srtab: str = None, smin_alpha: float = None, sw_rmin: float = None, sw_rmax: float = None ) -> None: """

2.1.3 使用飞桨框架自定义算子

我们使用飞桨框架的自定义算子方案来实现neighbor list采集、descriptor生成，以及力与维里应力的计算，替代原本的C++代码块。具体实现如下：

实现算子：在飞桨框架的自定义算子方案中，计算逻辑实现是通过像写函数一样，将输入和计算逻辑作为函数参数和执行逻辑加入。例如我们输入体系中原子坐标，找出给定截断距离cutoff内中心原子i的所有近邻j，并构造神经网络的输入矩阵descriptor：

std::vector<paddle::Tensor> PdProdEnvMatAOpCPUForward( const paddle::Tensor &coord_tensor, const paddle::Tensor &type_tensor, const paddle::Tensor &natoms_tensor, const paddle::Tensor &box_tensor, const paddle::Tensor &mesh_tensor, const paddle::Tensor &avg_tensor, const paddle::Tensor &std_tensor, float rcut_a, float rcut_r, float rcut_r_smth, std::vector<int> sel_a, std::vector<int> sel_r){ CHECK_INPUT(coord_tensor); CHECK_INPUT(type_tensor);

template <typename data_t>void PdProdEnvMatAOpCPUForwardKernel( int nsamples, int nloc, int ndescrpt, int nnei, int nall, int mem_cpy, int mem_nnei, int max_nbor_size, const int *mesh_tensor_data, int nei_mode, float rcut_a, float rcut_r, float rcut_r_smth, int max_cpy_trial, int max_nnei_trial, bool b_nlist_map, const std::vector<int>& sec_a, const std::vector<int>& sec_r, data_t *p_em, data_t *p_em_deriv, data_t *p_rij, int *p_nlist, const data_t *p_coord, const data_t *p_box, const data_t *avg, const data_t *std, const int *p_type){ for (size_t ff = 0; ff < nsamples; ++ff) { data_t *em = p_em + ff * nloc * ndescrpt; data_t *em_deriv = p_em_deriv + ff * nloc * ndescrpt * 3; data_t *rij = p_rij + ff * nloc * nnei * 3; int *nlist = p_nlist + ff * nloc * nnei; const data_t *coord = p_coord + ff * nall * 3; const data_t *box = p_box + ff * 9;)

具体实现参考：

https://github.com/deepmodeling/deepmd-kit/tree/paddle/source/op/paddle_ops/srcs

编译算子：飞桨框架支持jit的方式来编译算子，可以使用Python代码完成算子的编译和集成。

from paddle.utils.cpp_extension import CppExtension, CUDAExtension, setupimport site
site_package_dir = site.getsitepackages()[0]
setup( name='paddle_ops', ext_modules=CUDAExtension( sources=['../source/op/paddle_ops/srcs/pd_prod_env_mat_multi_devices_cpu.cc', '../source/op/paddle_ops/srcs/pd_prod_env_mat_multi_devices_cuda.cc', '../source/op/paddle_ops/srcs/pd_prod_force_se_a_multi_devices_cpu.cc', '../source/op/paddle_ops/srcs/pd_prod_force_se_a_multi_devices_cuda.cc', '../source/op/paddle_ops/srcs/pd_prod_virial_se_a_multi_devices_cpu.cc', '../source/op/paddle_ops/srcs/pd_prod_virial_se_a_multi_devices_cuda.cc'], include_dirs=["../source/lib/include/"], library_dirs=[site_package_dir+"/deepmd/op"], extra_link_args=["-ldeepmd"] ))

使用飞桨框架自定义算子的介绍参见：

https://www.paddlepaddle.org.cn/documentation/docs/zh/guides/07_new_op/new_custom_op.html

2.2 最终效果

使用飞桨框架实现的计算精度如下：

其中，训练后模型计算能量和力的精度都与Tensorflow版本一致。

代码节选自：

https://github.com/deepmodeling/deepmd-kit/tree/paddle

未来展望

本次百度飞桨深度学习平台与DeePMD-kit的合作，是开源社区交流、融合的重要实践，也是DeePMD-kit支持多机器学习框架的有益探索。后续实现飞桨机器学习框架的DeePMD-kit软件包将在算法优化、高性能适配等方面持续迭代。同时深度势能开源社区始终保持开源、开放、共享的精神，期待更多开源生态融合与交流的出现。

致谢

感谢飞桨开源社区贡献者，和深度势能社区的开发者：来自南京大学的王俊杰和王勇。

参考文献

1. L Zhang, J Han, H Wang, R Car, W E, Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics, Physical review letters 120 (14), 143001.

2. W Jia, H Wang, M Chen, D Lu, J Liu, L Lin, R Car, L Zhang, Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning, Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC '20), IEEE Press, Article 5, 1–14.

End

DP视角

感谢开源社区的高手朋友们！