選單

【清流Family】「深度勢能」聯手Paddle Paddle,推動分子模擬新生態建設

近期DeePMD-kit和開源深度學習框架飛槳(PaddlePaddle)的開發者們進行了跨領域深度合作,依託飛槳框架成熟的底層功能和動靜統一的開發體驗實現DeePMD演算法的構建和執行,促進了雙方在開發者生態上的雙向融合。

DeePMD演算法包括兩部分,用神經網路深度學習方案來模擬材料勢能面的DeepPot(深度勢模型)和用深度勢模型跑MD(分子動力學)。DeePMD演算法中,神經網路深度學習部分採用的是Tensorflow靜態圖實現。這裡,我們用DeePMD當中較為典型的SeA勢模型為例來為大家介紹一下用飛槳實現DeePMD演算法中深度學習方案的開發背景和過程。

技術背景

SeA指DeepPot smooth edition constructed from all information including angle and radial information of atomic configurations。從原子座標資訊(包括原子間的相對徑向座標和角向座標)出發,透過神經網路構造複雜材料體系的勢能面。

SeA模型由兩部分構成:

Embeding網路,輸入中心原子i與其所有近鄰原子j的相對座標資訊,輸出Feature Matrix。

Fitting網路,輸入Feature Matrix,輸出中心原子i的局域能量。

DeepPot已被用於準確描述包括分子結構、各種週期性結構,以及表面、高熵、電池等等多種複雜結構。

使用飛槳框架來實現SeA

飛槳依託於百度多年的深度學習技術研究和業務應用,是國內第一個自主研發、功能非常完備、並向開發者全面開源開放的產業級深度學習平臺。飛槳採用動態圖模式和自定義運算元方案。接下來為大家介紹下用飛槳實現DeePMD的具體方法。

2.1 構建SeA

2.1.1 SeA演算法結構

【清流Family】「深度勢能」聯手Paddle Paddle,推動分子模擬新生態建設

圖1 DeepPot-SeA演算法結構

Embedding網路:要想讓隨機性的神經網路能準確描述物理體系,Embedding輸出的Feature Matrix應該滿足體系的物理對稱性,包括平移、置換和旋轉對稱性。

因此,Embedding網路的構造基於以下兩個數學原理:

a。 矩陣滿足體系的所有平移和旋轉對稱性;

b。 所有滿足置換對稱性的函式形式可以變換成如下形式:,比如函式對於任意標量函式都滿足置換對稱性。

基於以上兩個原理,Feature Matrix的構造過程如上圖:將weighting function矩陣輸入Embedding網路(其中),輸出維矩陣,並取的前列得到。最後透過矩陣相乘構造Feature Matrix:。

Fitting網路:將Embedding網路輸出的Feature Matrix輸入Fitting網路,輸出i原子的局域能量Ei。

具體實現參考:

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()方法執行反向網路。

2.1.3 使用飛槳框架自定義運算元

我們使用飛槳框架的自定義運算元方案來實現neighbor list採集、descriptor生成,以及力與維裡應力的計算,替代原本的C++程式碼塊。具體實現如下:構造

具體實現參考:

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

編譯運算元:飛槳框架支援jit的方式來編譯運算元,可以使用Python程式碼完成運算元的編譯和整合。

使用飛槳框架自定義運算元的介紹參見:

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

2.2 最終效果

使用飛槳框架實現的計算精度如下:

【清流Family】「深度勢能」聯手Paddle Paddle,推動分子模擬新生態建設

其中,訓練後模型計算能量和力的精度都與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。