如何使用深度學(xué)習(xí)包在LAMMPS中驅(qū)動MD模擬

幾十年來，分子模擬界在模擬勢能面和原子間作用力時面臨著精度與效率的兩難選擇。深勢，人工神經(jīng)網(wǎng)絡(luò)力場，通過結(jié)合經(jīng)典分子動力學(xué)（ MD ）模擬的速度和密度泛函理論（ DFT ）計算的準(zhǔn)確性來解決這個問題。 這是通過使用 GPU – 優(yōu)化包 DeePMD-kit 實現(xiàn)的，這是一個用于多體勢能表示和 MD 模擬的深度學(xué)習(xí)包。

這篇文章提供了一個端到端的演示，演示如何為二維材料石墨烯訓(xùn)練神經(jīng)網(wǎng)絡(luò)潛力，并使用它在開源平臺大型原子/分子大規(guī)模并行模擬器（ LAMMPS ）中驅(qū)動 MD 模擬。 培訓(xùn)數(shù)據(jù)可從維也納從頭算模擬軟件包（ VASP ）獲得 ，或量子濃縮咖啡（ QE ）。

分子建模、機器學(xué)習(xí)和高性能計算（ HPC ）的無縫集成通過分子動力學(xué)和從頭算準(zhǔn)確性—這完全是通過基于容器的工作流來實現(xiàn)的。利用人工智能技術(shù)擬合 DFT 產(chǎn)生的原子間作用力，可以通過線性標(biāo)度將可訪問的時間和尺寸標(biāo)度提高幾個數(shù)量級。

深度潛能本質(zhì)上是機器學(xué)習(xí)和物理原理的結(jié)合，它開啟了一種新的計算范式，如圖 1 所示。

圖 1 。由分子建模、人工智能和高性能計算組成的新計算范式。（圖提供：張林峰博士， DP 技術(shù)）

整個工作流如圖 2 所示。數(shù)據(jù)生成步驟由 VASP 和 QE 完成。數(shù)據(jù)準(zhǔn)備、模型訓(xùn)練、測試和壓縮步驟使用 DeePMD 工具包完成。模型部署在 LAMMPS 中。

圖 2 。 DeePMD 工作流程圖。

為什么是集裝箱？

容器是一個可移植的軟件單元，它將應(yīng)用程序及其所有依賴項組合到一個與底層主機操作系統(tǒng)無關(guān)的包中。

本文中的工作流程涉及 AIMD 、 DP 培訓(xùn)和 LAMMPS MD 模擬。使用正確的編譯器設(shè)置、 MPI 、 GPU 庫和優(yōu)化標(biāo)志從源代碼安裝每個軟件包是非常重要和耗時的。

容器通過為每個步驟提供一個高度優(yōu)化的 GPU 支持的計算環(huán)境來解決這個問題，并且消除了安裝和測試軟件的時間。

NGC 目錄是 GPU 優(yōu)化的 HPC 和 AI 軟件的集線器，它攜帶了整個 HPC 和 AI 容器 ，可以很容易地部署在任何 GPU 系統(tǒng)上。 NGC 目錄中的 HPC 和 AI 容器經(jīng)常更新，并進(jìn)行可靠性和性能測試，這對于加快解決時間是必要的。

還將掃描這些容器的常見漏洞和暴露（ CVE ），確保它們沒有任何開放端口和惡意軟件。此外， HPC 容器支持 Docker 和 Singularity 運行時，并且可以部署在云中或本地運行的多［ZFBB］和多節(jié)點系統(tǒng)上。

訓(xùn)練數(shù)據(jù)生成

模擬的第一步是數(shù)據(jù)生成。我們將向您展示如何使用 VASP 和 Quantum ESPRESSO 來運行 AIMD 模擬并為 DeePMD 生成訓(xùn)練數(shù)據(jù)集?？梢允褂靡韵旅顝?GitHub 存儲庫下載所有輸入文件：

git clone https://github.com/deepmodeling/SC21_DP_Tutorial.git

VASP

如圖 3 所示，使用具有 98 個原子的二維石墨烯系統(tǒng)。 為了生成訓(xùn)練數(shù)據(jù)集，在 300K 下進(jìn)行 0 。 5ps NVT AIMD 模擬。選擇的時間步長為 0 。 5fs 。 DP 模型是使用固定溫度下 0 。 5ps MD 軌跡的 1000 個時間步長創(chuàng)建的。

由于仿真時間較短，訓(xùn)練數(shù)據(jù)集包含連續(xù)的系統(tǒng)快照，這些快照高度相關(guān)。通常，訓(xùn)練數(shù)據(jù)集應(yīng)從與各種系統(tǒng)條件和配置不相關(guān)的快照中采樣。對于這個例子，我們使用了一個簡化的訓(xùn)練數(shù)據(jù)方案。對于生產(chǎn) DP 培訓(xùn)，建議使用 DP-GEN 利用并行學(xué)習(xí)方案，以有效探索更多的條件組合。

用投影增強波贗勢描述了價電子與凍結(jié)核之間的相互作用。廣義梯度近似交換? Perdew 的相關(guān)泛函?伯克?恩澤霍夫。在所有系統(tǒng)中，只有 Γ-point 用于 k-space 采樣。

圖 3 AIMD 模擬中使用了由 98 個碳原子組成的石墨烯系統(tǒng)。

量子濃縮咖啡

AIMD 模擬也可以使用 Quantum ESPRESSO （ NGC 目錄中的container提供）進(jìn)行。 Quantum ESPRESSO 是一套基于密度泛函理論、平面波和贗勢的開放源代碼，用于 Nan oscale 的電子結(jié)構(gòu)計算和材料建模。 QE 計算中使用了相同的石墨烯結(jié)構(gòu)。以下命令可用于啟動 AIMD 模擬：

$ singularity exec --nv docker://nvcr.io/hpc/quantum_espresso:qe-6.8 cp.x < c.md98.cp.in

培訓(xùn)數(shù)據(jù)準(zhǔn)備

一旦從 AIMD 仿真中獲得訓(xùn)練數(shù)據(jù)，我們希望使用?dpdata因此，它可以作為深層神經(jīng)網(wǎng)絡(luò)的輸入。dpdata包是 AIMD 、 Classic MD 和 DeePMD 工具包之間的格式轉(zhuǎn)換工具包。

您可以使用方便的工具dpdata將數(shù)據(jù)直接從 first principles 軟件包的輸出轉(zhuǎn)換為 DeePMD 工具包格式。對于深勢訓(xùn)練，必須提供物理系統(tǒng)的以下信息：原子類型、盒邊界、坐標(biāo)、力、病毒和系統(tǒng)能量。

快照或系統(tǒng)框架在一個時間步中包含所有原子的所有這些數(shù)據(jù)點，可以以兩種格式存儲，即raw和npy。

第一種格式raw是純文本，所有信息都在一個文件中，文件的每一行表示一個快照。不同的系統(tǒng)信息存儲在名為box.raw, coord.raw, force.raw, energy.raw和virial.raw的不同文件中。我們建議您在準(zhǔn)備培訓(xùn)文件時遵循這些命名約定。

force.raw的一個示例：

$ cat force.raw
-0.724 2.039 -0.951 0.841 -0.464 0.363 6.737 1.554 -5.587 -2.803 0.062 2.222
-1.968 -0.163 1.020 -0.225 -0.789 0.343

這個force.raw包含三個框架，每個框架具有兩個原子的力，形成三條線和六列。每條線在一幀中提供兩個原子的所有三個力分量。前三個數(shù)字是第一個原子的三個力分量，而下三個數(shù)字是第二個原子的力分量。

坐標(biāo)文件coord.raw的組織方式類似。在box.raw中，應(yīng)在每行上提供盒向量的九個分量。在virial.raw中，維里張量的九個分量應(yīng)按XX XY XZ YX YY YZ ZX ZY ZZ的順序提供在每一行上。所有原始文件的行數(shù)應(yīng)相同。我們假設(shè)原子類型不會在所有幀中改變。它由type.raw提供，它有一行原子類型，一行一行地寫。

原子類型應(yīng)該是整數(shù)。例如，一個系統(tǒng)的type.raw有兩個原子，分別為零和一：

$ cat type.raw
0 1

將數(shù)據(jù)格式轉(zhuǎn)換為raw不是一項要求，但此過程應(yīng)能說明可作為培訓(xùn)用 DeePMD 工具包輸入的數(shù)據(jù)類型。

將第一原理結(jié)果轉(zhuǎn)換為訓(xùn)練數(shù)據(jù)的最簡單方法是將其保存為 NumPy 二進(jìn)制數(shù)據(jù)。

對于 VASP 輸出，我們準(zhǔn)備了一個outcartodata.py腳本來處理 VASP OUTCAR 文件。通過運行以下命令：

$ cd SC21_DP_Tutorial/AIMD/VASP/
$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 python outcartodata.py
$ mv deepmd_data ../../DP/

量化寬松產(chǎn)出：

$ cd SC21_DP_Tutorial/AIMD/QE/
$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 python logtodata.py
$ mv deepmd_data ../../DP/

生成名為deepmd_data的文件夾并將其移動到培訓(xùn)目錄。它生成五組0/set.000, 1/set.000, 2/set.000, 3/set.000, 4/set.000，每組包含 200 幀。不需要處理每個 set .*目錄中的二進(jìn)制數(shù)據(jù)文件。包含set.*文件夾和type.raw文件的路徑稱為系統(tǒng)。如果要訓(xùn)練非周期系統(tǒng)，應(yīng)在系統(tǒng)目錄下放置一個空nopbc文件。box.raw不是必需的，因為它是非周期系統(tǒng)。

我們將使用五套中的三套進(jìn)行培訓(xùn)，一套用于驗證，另一套用于測試。

深勢模型訓(xùn)練

深勢模型的輸入是包含前面提到的系統(tǒng)信息的描述符向量。神經(jīng)網(wǎng)絡(luò)包含幾個隱藏層，由線性和非線性變換組成。在這篇文章中，使用了一個三層神經(jīng)網(wǎng)絡(luò)，每層有 25 個、 50 個和 100 個神經(jīng)元。神經(jīng)網(wǎng)絡(luò)學(xué)習(xí)的目標(biāo)值或標(biāo)簽是原子能。訓(xùn)練過程通過最小化損失函數(shù)來優(yōu)化權(quán)重和偏差向量。

訓(xùn)練由命令啟動，其中input.json包含訓(xùn)練參數(shù)：

$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 dp train input.json

DeePMD 工具包打印培訓(xùn)和驗證數(shù)據(jù)集的詳細(xì)信息。數(shù)據(jù)集由輸入腳本的training部分中定義的training_data和validation_data確定。訓(xùn)練數(shù)據(jù)集由三個數(shù)據(jù)系統(tǒng)組成，而驗證數(shù)據(jù)集由一個數(shù)據(jù)系統(tǒng)組成。原子數(shù)、批次大小、系統(tǒng)中的批次數(shù)以及使用系統(tǒng)的概率均如圖 4 所示。最后一列顯示系統(tǒng)是否假設(shè)周期邊界條件。

圖 4 DP 培訓(xùn)輸出的屏幕截圖。

在培訓(xùn)期間，每disp_freq培訓(xùn)步驟都會使用用于培訓(xùn)模型的批次和驗證數(shù)據(jù)中的numb_btch批次測試模型的錯誤。在文件disp_file中相應(yīng)地打印訓(xùn)練錯誤和驗證錯誤（默認(rèn)為lcurve.out）?？稍谳斎肽_本中通過訓(xùn)練和驗證數(shù)據(jù)集相應(yīng)部分中的鍵batch_size設(shè)置批量大小。

輸出的一個示例：

# step rmse_val rmse_trn rmse_e_val rmse_e_trn rmse_f_val rmse_f_trn lr 0 3.33e+01 3.41e+01 1.03e+01 1.03e+01 8.39e-01 8.72e-01 1.0e-03 100 2.57e+01 2.56e+01 1.87e+00 1.88e+00 8.03e-01 8.02e-01 1.0e-03 200 2.45e+01 2.56e+01 2.26e-01 2.21e-01 7.73e-01 8.10e-01 1.0e-03 300 1.62e+01 1.66e+01 5.01e-02 4.46e-02 5.11e-01 5.26e-01 1.0e-03 400 1.36e+01 1.32e+01 1.07e-02 2.07e-03 4.29e-01 4.19e-01 1.0e-03 500 1.07e+01 1.05e+01 2.45e-03 4.11e-03 3.38e-01 3.31e-01 1.0e-03

如圖 5 所示，訓(xùn)練誤差隨著訓(xùn)練步驟單調(diào)減少。訓(xùn)練后的模型在測試數(shù)據(jù)集上進(jìn)行了測試，并與 AIMD 仿真結(jié)果進(jìn)行了比較。測試命令是：

$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 dp test -m frozen_model.pb -s deepmd_data/4/ -n 200 -d detail.out

圖 5 有步驟的訓(xùn)練損失

結(jié)果如圖 6 所示。

圖 6 用 AIMD 能量和力測試訓(xùn)練后的 DP 模型的預(yù)測精度。

模型導(dǎo)出和壓縮

模型訓(xùn)練完成后，生成一個凍結(jié)模型，用于 MD 仿真中的推理。從檢查點保存神經(jīng)網(wǎng)絡(luò)的過程稱為“凍結(jié)”模型：

$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 dp freeze -o graphene.pb

生成凍結(jié)模型后，可以在不犧牲精度的情況下對模型進(jìn)行壓縮；在 MD 中大大加快推理性能的同時，根據(jù)仿真和訓(xùn)練設(shè)置，模型壓縮可以將性能提高 10 倍，在 GPU 上運行時將內(nèi)存消耗減少 20 倍。

可以使用以下命令壓縮凍結(jié)模型，-i表示凍結(jié)模型，-o表示壓縮模型的輸出名稱：

$ singularity exec --nv docker://nvcr.io/hpc/deepmd-kit:v2.0.3 dp compress -i graphene.pb -o graphene-compress.pb

LAMMPS 中的模型部署

在 LAMMPS 中實現(xiàn)了一種新的配對方式，以便在前面的步驟中部署經(jīng)過訓(xùn)練的神經(jīng)網(wǎng)絡(luò)。對于熟悉 LAMMPS 工作流程的用戶，只需進(jìn)行最小的更改即可切換到深度潛力。例如，具有 Tersoff 電位的傳統(tǒng) LAMMPS 輸入具有以下電位設(shè)置：

pair_style tersoff
pair_coeff * * BNC.tersoff C

若要使用深電位，請將以前的線路替換為：

pair_style deepmd graphene-compress.pb
pair_coeff * *

輸入文件中的pair_style命令使用 DeePMD 模型來描述石墨烯系統(tǒng)中的原子相互作用。

• graphene-compress.pb文件表示用于推斷的凍結(jié)和壓縮模型。
• MD 模擬中的石墨烯系統(tǒng)包含 1560 個原子。
• 周期性邊界條件應(yīng)用于橫向x和y方向，自由邊界應(yīng)用于z方向。
• 時間步長設(shè)置為 1 fs 。
• 將系統(tǒng)置于溫度為 300 K 的 NVT 系綜下進(jìn)行松弛，這與 AIMD 設(shè)置一致。

NVT 松弛后的系統(tǒng)配置如圖 7 所示?？梢杂^察到，深勢可以描述原子結(jié)構(gòu)，在橫平面方向上有小的波紋。在 10ps NVT 松弛后，將系統(tǒng)置于 NVE 系綜下以檢查系統(tǒng)穩(wěn)定性。

圖 7 深勢弛豫后石墨烯體系的原子構(gòu)型。

系統(tǒng)溫度如圖 8 所示。

圖 8 NVT 和 NVE 組合下的系統(tǒng)溫度。深勢驅(qū)動的分子動力學(xué)系統(tǒng)在弛豫后非常穩(wěn)定。

為了驗證經(jīng)過訓(xùn)練的 DP 模型的準(zhǔn)確性，從 AIMD 、 DP 和 Tersoff 計算出的徑向分布函數(shù)（ RDF ）如圖 9 所示。 DP 模型生成的 RDF 與 AIMD 模型非常接近，這表明 DP 模型可以很好地描述石墨烯的晶體結(jié)構(gòu)。

圖 9 。分別用 AIMD 、 DP 和 Tersoff 勢計算徑向分布函數(shù)?？梢杂^察到， DP 計算的 RDF 與 AIMD 非常接近。

結(jié)論

這篇文章展示了在給定條件下石墨烯的一個簡單案例研究。 DeePMD-kit 軟件包簡化了從 AIMD 到經(jīng)典 MD 的工作流程，具有很大的潛力，提供了以下關(guān)鍵優(yōu)勢：

在 TensorFlow 框架中實現(xiàn)高度自動化和高效的工作流。

使用流行的 DFT 和 MD 包（如 VASP 、 QE 和 LAMMPS ）的 API 。

廣泛應(yīng)用于有機分子、金屬、半導(dǎo)體、絕緣體等。

具有 MPI 和［ZFBB］支持的高效 HPC 代碼。

模塊化，便于其他深度學(xué)習(xí)潛在模型采用。

此外，使用 NGC 目錄中的 GPU – 優(yōu)化容器簡化并加速了整個工作流程，省去了安裝和配置軟件的步驟。

關(guān)于作者

Jingchao Zhang 是英偉達(dá)高等教育和研究團(tuán)隊的資深解決方案設(shè)計師。他也是佛羅里達(dá)大學(xué) Nvidia AI 技術(shù)中心（ NVAITC ）的常駐科學(xué)家。他獲得了博士學(xué)位。愛荷華州立大學(xué)機械工程專業(yè)。在加入 NVIDIA 之前，他在研究計算領(lǐng)域工作了 8 年。他的研究興趣包括計算材料科學(xué)和機器學(xué)習(xí)。

Yifan Li 是普林斯頓大學(xué)化學(xué)系一年級研究生。伊凡于 2020 年 7 月在北京大學(xué)獲得化學(xué)學(xué)士學(xué)位。他的研究興趣包括分子動力學(xué)中的機器學(xué)習(xí)力場、核量子動力學(xué)模擬和自動區(qū)分量子化學(xué)程序。

Akhil Docca 是 NVIDIA NGC 的高級產(chǎn)品營銷經(jīng)理，專注于 HPC 和 DL 容器。 Akhil 擁有加州大學(xué)洛杉磯分校安德森商學(xué)院工商管理碩士學(xué)位，圣何塞州立大學(xué)機械工程學(xué)士學(xué)位。

審核編輯：郭婷