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

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

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

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

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

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

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

圖 2 。 DeePMD 工作流程圖。

為什么是集裝箱？

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

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

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

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

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

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

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

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

VASP

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

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

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

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

量子濃縮咖啡

AIMD 模擬也可以使用 Quantum ESPRESSO （ NGC 目錄中的container提供）進(jìn)行。 Quantum ESPRESSO 是一套基于密度泛函理論、平面波和贗勢(shì)的開(kāi)放源代碼，用于 Nan oscale 的電子結(jié)構(gòu)計(jì)算和材料建模。 QE 計(jì)算中使用了相同的石墨烯結(jié)構(gòu)。以下命令可用于啟動(dòng) 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 工具包格式。對(duì)于深勢(shì)訓(xùn)練，必須提供物理系統(tǒng)的以下信息：原子類(lèi)型、盒邊界、坐標(biāo)、力、病毒和系統(tǒng)能量。

快照或系統(tǒng)框架在一個(gè)時(shí)間步中包含所有原子的所有這些數(shù)據(jù)點(diǎn)，可以以?xún)煞N格式存儲(chǔ)，即raw和npy。

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

force.raw的一個(gè)示例：

$ 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

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

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

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

$ cat type.raw
0 1

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

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

對(duì)于 VASP 輸出，我們準(zhǔn)備了一個(gè)outcartodata.py腳本來(lái)處理 VASP OUTCAR 文件。通過(guò)運(yùn)行以下命令：

$ 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的文件夾并將其移動(dòng)到培訓(xùn)目錄。它生成五組0/set.000, 1/set.000, 2/set.000, 3/set.000, 4/set.000，每組包含 200 幀。不需要處理每個(gè) set .*目錄中的二進(jìn)制數(shù)據(jù)文件。包含set.*文件夾和type.raw文件的路徑稱(chēng)為系統(tǒng)。如果要訓(xùn)練非周期系統(tǒng)，應(yīng)在系統(tǒng)目錄下放置一個(gè)空nopbc文件。box.raw不是必需的，因?yàn)樗欠侵芷谙到y(tǒng)。

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

深勢(shì)模型訓(xùn)練

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

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

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

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

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

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

輸出的一個(gè)示例：

# 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)練后的模型在測(cè)試數(shù)據(jù)集上進(jìn)行了測(cè)試，并與 AIMD 仿真結(jié)果進(jìn)行了比較。測(cè)試命令是：

$ 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 能量和力測(cè)試訓(xùn)練后的 DP 模型的預(yù)測(cè)精度。

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

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

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

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

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

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

LAMMPS 中的模型部署

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

pair_style tersoff
pair_coeff * * BNC.tersoff C

若要使用深電位，請(qǐng)將以前的線(xiàn)路替換為：

pair_style deepmd graphene-compress.pb
pair_coeff * *

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

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

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

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

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

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

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

圖 9 。分別用 AIMD 、 DP 和 Tersoff 勢(shì)計(jì)算徑向分布函數(shù)。可以觀(guān)察到， DP 計(jì)算的 RDF 與 AIMD 非常接近。

結(jié)論

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

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

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

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

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

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

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

關(guān)于作者

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

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

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

審核編輯：郭婷