助力前沿基础科学研究，飞桨联手深度势能推动分子模拟新生态建设

开源开放已经成为驱动技术创新和加速产业发展的核心动能。开源开放对人类社会过去几百年科学和技术的发展起到非常重要的作用,让我们可以更快追踪到最新的技术进展。

百度飞桨深度学习平台作为百度“AI大生产平台”的基础底座,以大规模分布式训练、全硬件平台支持、端到端全流程工具,助力企业和开发者加速AI生产。无需从头学习难度高、迭代快的前沿科技,避免重复“造轮子”,应用AI的门槛更低、成本更低、效率更高。

近期飞桨(PaddlePaddle)开源深度学习框架和深度势能开源社区的开发者们就深度势能分子动力学开源项目(DeePMD-kit)进行了跨领域深度合作,依托飞桨框架成熟的底层功能和动静统一的开发体验实现深度势能分子动力学模型的构建和运行,促进了双方在开发者生态上的双向融合,拓展了DeePMD-kit原有能力边界,增强在国产化硬件和大规模分布式训练上的能力,也让Paddle支持的科学计算开源项目中增加了重量级的一员。

技术解读

1.技术背景

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

SeA模型由两部分构成:

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

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

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

2.使用PaddlePaddle来实现SeA网络

1.构建网络

1.网络结构

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

具体实现方案参见https://dl.acm.org/doi/10.5555/3327345.3327356

2.使用PaddlePaddle动态图来构建SeA网络

PaddlePaddle动态图使用了相比tensorflow 1.x更加灵活的可以支持更加贴近Python语法和执行逻辑的执行机制,使得算法实现可以按照顺序的执行逻辑来进行。

1.构建能量模型:

我们只需要简单的继承 paddle.nn.Layer就可以使用paddle内部的各类操作网络的方法具体可以参见:

在__init__方法中我们像使用python一样设置网络结构(包括层数和每层的维度)

在forward()方法中我们就可以实现加权求和+激活函数 逻辑

而深度学习中反向的计算逻辑paddle的动态图此时已经使用autograd技术自动为我们完成了。我们只需要在整个网络完成后使用backward()方法即可执行反向网络。

飞桨框架2.0支持用户使用动态图完成深度学习相关领域全类别的模型算法开发。动态图模式下可以让开发者随时查看变量的输入、输出,方便快捷的调试程序,带来最佳的开发体验。为了解决动态图的部署问题,飞桨提供了全面完备的动转静支持,在Python语法支持覆盖度上达到领先水平。开发者在动态图编程调试的过程中,仅需添加一个装饰器paddle.jit.to_static,即可无缝平滑地自动实现静态图训练或模型保存。同时飞桨框架2.0还做到了模型存储和加载的接口统一,保证动转静之后保存的模型文件能够被纯动态图加载和使用。

3.使用PaddlePaddle自定义算子

我们使用PaddlePaddle的自定义算子方案来实现这种和深度势能领域强相关的计算单元。相比tensorflow 1.x当中需要实现一个内部算子,PaddlePaddle的自定义算子方案可以让我们更加专注于实现相关的计算逻辑而尽量脱离了繁琐的编译算子和集成算子到框架的过程。

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

具体的实现可以参考:

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

2.编译算子:PaddlePaddle甚至支持jit的方式来编译算子使得我们可以使用简单的几行Python代码就可以完成算子的编译和集成

具体使用PaddlePaddle自定义Op的介绍可以参见:

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

3. DoubleGrad的使用:

训练模型所用的label包括力和维里应力,两者都要通过对能量关于坐标求导得出。加上back propagation,就对模型的输出值能量进行了两次求导。

在paddlepaddle的动态图中可以简单的使用paddle.grad来实现

https://www.paddlepaddle.org.cn/documentation/docs/zh/api/paddle/fluid/dygraph/base/grad_cn.html#grad

Paddle.grad可以自动实现对给定函数的求导。对于输入多个自变量x,并输出多个因变量y的函数,paddle.grad对每个x分别返回这个x对所有y的导数之和,也就是说返回值个数与x的个数相同。

使用PaddlePaddle我们可以很完美的重现C++算法实现的能量关于坐标求导的功能,并且我们可以看到相关的实现精度也可以达到以下的效果:

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

以上图片选自

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

代码节选自:

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

关于DeePMD-kit

以深度学习为代表的机器学习方法在力场开发中的应用为发展高效精确的分子动力学方法带来了新机遇和新思路。深度势能【1】(Deep Potential)开源方法通过运用“机器学习+多尺度建模+高性能计算”的方式解决了传统分子模拟过程中难以兼顾精度与速度的痛点,成功实现在第一性原理精度基础上的上亿原子的分子动力学模拟【2】。其中DeePMD-kit和DP-GEN是围绕Deep Potential Molecular Dynamics(DeePMD,深度势能分子动力学)开发的开源科学软件包。发布以来,在凝聚态物理、化学物理、材料理性设计、生物物理、药物分子设计等领域均得到广泛应用。

关于飞桨

飞桨(PaddlePaddle)以百度多年的深度学习技术研究和业务应用为基础,是中国首个自主研发、功能完备、开源开放的产业级深度学习平台,集深度学习核心训练和推理框架、基础模型库、端到端开发套件和丰富的工具组件于一体。目前,飞桨已凝聚超265万开发者,服务企业10万家,基于飞桨开源深度学习平台产生了 34万个模型。飞桨助力开发者快速实现AI想法,快速上线AI业务。帮助越来越多的行业完成AI赋能,实现产业智能化升级。

当前在科学计算领域,飞桨联合各高校、科研机构和企业已为开发者建设支持了量子机器学习开发工具PaddleQuantum、生物计算开源工具集PaddleHelix和贝叶斯深度学习ZhuSuan,本次又在分子动力学研究领域新增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.

参考链接

DeePMD-kit项目地址:

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

如果您想详细了解更多飞桨的相关内容,请参阅以下文档。

飞桨官网地址·

https://www.paddlepaddle.org.cn/

飞桨开源框架项目地址·

GitHub: https://github.com/PaddlePaddle/Paddle

Gitee: https://gitee.com/paddlepaddle/Paddle