跳转至

DeePMD-kit 2.x 使用入门

简介

DeePMD-kit是一个训练神经网络势能(Machine Learning Potential)的代码包。该包主要由张林峰(普林斯顿大学),王涵(北京应用物理与计算数学研究所)开发。黄剑兴和庄永斌曾经短时间参与开发。如有问题,可以向他们询问。

Danger

我们已经舍弃了1.x版本的教程。

以下为参考信息:

Warning

此页面仅限提供贡献者对于该软件的理解,如有任何问题请联系贡献者

第一次尝试

运行第一次机器学习

如果你正在使用 Zeus 集群,请使用 slurm 脚本来提交 DeePMD-kit 任务。

请从 Github 下载 DeePMD-kit 的代码,我们将会使用里面的水模型做为例子。

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

首先进入含有水模型的例子的目录

cd <deepmd repositoy>/examples/water/se_e2_a/

你会看到input.json文件,这是DeePMD-kit使用的输入文件。现在复制/data/share/base/script/deepmd.lsf到当前文件夹,并且修改它。

cp /data/share/base/script/deepmd.lsf ./
vim deepmd.lsf

Warning

如果调用的是1.0的版本,需要在learning_rate下加入decay_rate关键词,一般设为0.95.

你现在仅需要修改 slurm 脚本中的输入文件名称即可。把脚本中的input.json替换成water_se_a.json

#!/bin/bash

#BSUB -q gpu
#BSUB -W 24:00
#BSUB -J train
#BSUB -o %J.stdout
#BSUB -e %J.stderr
#BSUB -n 8
#BSUB -R "span[ptile=8]"
# ============================================
# modify the number of cores to use
# according to the number of GPU you select
# for example, 8 cores for one GPU card
# while there are 32 cores in total
# ============================================

# add modulefiles
module add deepmd/2.2.7

# automatic select the gpu
source /data/share/base/script/find_gpu.sh

dp train input.json -l train.log

使用如下命令提交任务:

#submit your job
bsub < deepmd.lsf
#check your job by
bjobs 

当任务执行中,当前目录会生成以下文件:

  • train.log: 训练的记录文件
  • lcurve.out: 机器学习的学习曲线
  • model.ckpt.data-00000-of-00001, model.ckpt.index, checkpoint, model.ckpt.meta: 以上三个为训练存档点

非常好!已经成功开始第一次机器学习训练了!

浏览输出文件

使用 less 命令来浏览输出文件

less train.log

你将会看到如下内容

# DEEPMD: initialize model from scratch
# DEEPMD: start training at lr 1.00e-03 (== 1.00e-03), final lr will be 3.51e-08
2019-12-07 00:03:49.659876: I tensorflow/stream_executor/platform/default/dso_loader.cc:42] Successfully opened dynamic library libcublas.so.10.0
# DEEPMD: batch     100 training time 5.95 s, testing time 0.18 s
# DEEPMD: batch     200 training time 4.58 s, testing time 0.20 s
# DEEPMD: batch     300 training time 4.56 s, testing time 0.14 s
# DEEPMD: batch     400 training time 4.49 s, testing time 0.13 s
# DEEPMD: batch     500 training time 4.60 s, testing time 0.14 s
# DEEPMD: batch     600 training time 4.61 s, testing time 0.15 s
# DEEPMD: batch     700 training time 4.43 s, testing time 0.18 s
# DEEPMD: batch     800 training time 4.59 s, testing time 0.13 s
# DEEPMD: batch     900 training time 4.41 s, testing time 0.17 s
# DEEPMD: batch    1000 training time 4.66 s, testing time 0.11 s
# DEEPMD: saved checkpoint model.ckpt
# DEEPMD: batch    1100 training time 4.45 s, testing time 0.15 s
# DEEPMD: batch    1200 training time 4.37 s, testing time 0.14 s

batch后面的数字表明程序已经放入了多少数据进行训练。这个数字的显示间隔,即100,是在输入文件的"disp_freq": 100 设置的。

现在来看看你的学习曲线 lcurve.out

less lcurve.out

你将会看到:

#  step      rmse_val    rmse_trn    rmse_e_val  rmse_e_trn    rmse_f_val  rmse_f_trn         lr
      0      1.69e+01    1.58e+01      1.52e+00    5.69e-01      5.35e-01    5.00e-01    1.0e-03
   1000      4.74e+00    4.68e+00      3.88e-02    4.02e-01      1.50e-01    1.48e-01    1.0e-03
   2000      5.06e+00    3.93e+00      1.86e-01    1.54e-01      1.60e-01    1.24e-01    1.0e-03
   3000      4.73e+00    4.34e+00      9.08e-02    3.90e-01      1.49e-01    1.37e-01    1.0e-03
   4000      4.65e+00    6.09e+00      2.24e-01    1.92e-01      1.47e-01    1.93e-01    1.0e-03
   5000      3.84e+00    3.25e+00      5.26e-02    2.40e-02      1.25e-01    1.06e-01    9.4e-04
   6000      4.17e+00    2.78e+00      6.35e-02    3.89e-02      1.36e-01    9.03e-02    9.4e-04
   7000      3.24e+00    3.00e+00      5.55e-02    8.58e-03      1.05e-01    9.76e-02    9.4e-04
   8000      2.97e+00    2.83e+00      2.97e-02    2.46e-02      9.68e-02    9.22e-02    9.4e-04
   9000      1.01e+01    6.92e+00      1.36e-01    1.89e-01      3.28e-01    2.25e-01    9.4e-04
  10000      3.73e+00    3.39e+00      4.38e-02    3.23e-02      1.25e-01    1.14e-01    8.9e-04
  11000      3.51e+00    2.76e+00      1.31e-01    3.47e-01      1.17e-01    8.98e-02    8.9e-04
  12000      2.59e+00    2.89e+00      1.35e-01    1.18e-01      8.57e-02    9.65e-02    8.9e-04
  13000      5.65e+00    4.68e+00      3.08e-01    3.28e-01      1.88e-01    1.55e-01    8.9e-04

这些数字展示了当前机器学习模型对于数据预测的误差有多大。 rmse_e_trn 意味着在测试集上使用机器学习模型预测的能量误差会有多大。 rmse_e_val 意味着在训练集上使用机器学习模型预测的能量误差会有多大。 rmse_f_tst and rmse_f_trn 表示相同意义,不过是对于力的预测. 你可以使用Matplotlib Python包进行作图。

使用进阶

准备训练数据

前半部分仅仅是让你运行DeePMD-kit进行训练。为了训练一个针对你的体系的模型,你需要自己来准备数据。这些数据都是第一性原理计算得到的数据。这些数据可以是单点能计算得到的数据,或者是分子动力学模拟得到的数据。作为数据集需要的数据有:

  • 体系的结构文件:coord.npy
  • 体系的结构文件对应的元素标记:type.raw
  • 体系的结构文件对应的能量:energy.npy
  • 体系的结构文件对应的力:force.npy
  • 体系的结构文件对应的晶胞大小,如果是非周期性体系,请在训练文件里准备一个超大周期边界条件:box.npy

代码块里的文件名为DeePMD-kit使用的命名。npy后缀为Python的numpy代码包生成的文件,请在此之前学习numpy。如果你使用cp2k得到数据,你会有 *pos-1.xyz*frc-1.xyz 文件。你可以使用帮助的脚本转化成DeePMD-kit的数据集格式。

现在我们来看看DeePMD-kit的训练数据格式。之前我们训练的水模型的数据集储存在 <deepmd repository>/examples/water/data/data_0. 让我们来看看数据集的目录结构:

# directory structre for training data
.
├── data_0
│   ├── set.000
│      ├── box.npy
│      ├── coord.npy
│      ├── energy.npy
│      └── force.npy
│   ├── type.raw
│   └── type_map.raw
├── data_1
│   ├── set.000
│      ├── box.npy
│      ├── coord.npy
│      ├── energy.npy
│      └── force.npy
│   ├── set.001
│      ├── box.npy
│      ├── coord.npy
│      ├── energy.npy
│      └── force.npy
│   ├── type.raw
│   └── type_map.raw
├── data_2
│   ├── set.000
│      ├── box.npy
│      ├── coord.npy
│      ├── energy.npy
│      └── force.npy
│   ├── type.raw
│   └── type_map.raw
└── data_3
    ├── set.000
       ├── box.npy
       ├── coord.npy
       ├── energy.npy
       └── force.npy
    ├── type.raw
    └── type_map.raw

显然,我们会看到type.raw文件和一堆以set开头的目录。type.raw文件记录了体系的元素信息。如果你打开你会发现它仅仅记录了一堆数字。这些数字对应着你在water_se_a.json"type_map":["O","H"]的信息。此时0代表O,1代表H。对应着["O","H"]中的位置,其中第一位为0。

0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1

box.npy, coord.npy, energy.npyforce.npy 储存的信息在上文已经说过。唯一需要注意的是这些文件都储存着一个超大的矩阵。如果我们有Y个结构,每个结构有X个原子。box.npy, coord.npy, energy.npyforce.npy 对应的矩阵形状分别是 (Y, 9), (Y, X*3), (Y, 1), (Y, X*3)。

设置你的输入文件

输入文件是json文件。你可以使用之前我们的json文件进行细微改动就投入到自己体系的训练中。这些需要修改的关键词如下:

  • type": "se_a": 设置描述符(descriptor)类型。一般使用se_a
  • "sel": [46, 92]: 设置每个原子的截断半径内所拥有的最大原子数。注意这里的两个数字46,92分别对应的是O原子和H原子。与你在type_map里设置的元素类型是相对应的。

"descriptor" :{
         "type":     "se_a",
         "sel":      [46, 92],
         "rcut_smth":    0.50,
         "rcut":     6.00,
         "neuron":       [25, 50, 100],
         "resnet_dt":    false,
         "axis_neuron":  16,
         "seed":     1,
         "_comment":     " that's all"
     },
在"training"的"training_data"下 - "systems": ["../data/data_0/", "../data/data_1/", "../data/data_2/"]: 设置包含训练数据的目录。 - "batch_size": auto, 这个会根据体系原子数进行分配,不过我们自己通常设置为1,因为体系原子数有400-800个左右。

    "training_data": {
        "systems":      ["../data/data_0/", "../data/data_1/", "../data/data_2/"],
        "batch_size":   "auto",
        "_comment":     "that's all"
    }
在"training"的"validation_data"下 - "systems": ["../data/data_3"]: 设置包含测试数据的目录。 - "batch_size": 1, 这个会根据体系原子数进行分配,不过我们自己通常设置为1,因为体系原子数有400-800个左右。 - "numb_btch": 3 , 每次迭代中,测试的结构数量为batch_size乘以numb_btch。 - 更多参数说明,请参考官方文档:https://deepmd.readthedocs.io/en/latest/train-input.html

Warning

记住在集群上训练,请使用lsf脚本。

开始你的训练

使用如下命令开始:

dp train input.json

Warning

记住在集群上训练,请使用 Slurm 脚本。

重启你的训练

使用以下命令重启:

dp train input.json --restart model.ckpt

Warning

记住在集群上训练,请使用 Slurm 脚本。

使用生成的势能函数进行分子动力学(MD)模拟

当我们完成训练之后,我们需要根据节点文件(model.ckpt*)冻结(Freeze)出一个模型来。

利用如下命令,可以冻结模型:

dp freeze

你将会得到一个*.pb文件。利用此文件可以使用LAMMPS, ASE, CP2K 等软件进行分子动力学模拟。

利用压缩模型进行产出(Production)

机器学习势能*.pb文件进行MD模拟虽然已经非常迅速了。但是还有提升的空间。首先我们需要用2.0以上版本的deepmd进行训练势能函数,并得到*.pb文件。利用1.2/1.3版本的deepmd训练得到势能函数也不用担心,可以利用以下命令对旧版本的势能函数进行转换。例如想要从1.2转换的话:

dp convert-from 1.2 -i old_frozen_model.pb -o new_frozen_model.pb

关于兼容性的说明

关于目前势函数的兼容性,请参考官方文档。 目前DeePMD-kit支持从 v0.12, v1.0, v1.1, v1.2, v1.3 版本到新版本的转换。

建议将原训练文件夹备份后复制,我们利用如下命令进行压缩(文件夹下应该含有对应的input.json文件和checkpoint文件):

module load deepmd/2.0-cuda11.3
dp compress -i normal-model.pb -o compressed-model.pb -l compress.log

适用范围

注意模型压缩仅适用于部分模型,如 se_e2_a, se_e3, se_e2_r 和上述模型的 Hybrid 模型。

若使用其他模型,如 se_attn 模型 (DPA-1),模型压缩尚未被支持,可能会报错。

另外请注意,压缩模型是通过使用 5 次多项式拟合 Embedding-net 从而换取性能提升,这一改动 几乎 不会对预测精度产生影响,但实际上部分牺牲了精度。 因而使用时请务必注意观察默认参数是否适用于当前体系的情况,如是否出现误差漂移,并针对修改参数,如拟合时采用的步数 --step

压缩模型与原始模型对比

测试2080Ti, 显存11G

体系 原子数 提速前 (ns/day) 提速后(ns/day) 提升倍率
LIGePS 5000 0.806 3.569 4.42
SnO2/water interface 6021 0.059 0.355 6.01
SnO2/water interface 5352 0.067 0.382 5.70
SnO2/water interface 2676 0.132 0.738 5.59
SnO2/water interface 1338 0.261 1.367 5.23
SnO2/water interface 669 0.501 2.236 4.46
LiGePS 400 7.461 23.992 3.21
Cu13 13 51.268 65.944 1.28

SnO2/water interface: 原始模型Maximum 6021 ——> 压缩模型Maximum 54189个原子

Trouble Shooting

warning: loc idx out of lower bound

Solution: https://github.com/deepmodeling/deepmd-kit/issues/21

ValueError: NodeDef missing attr 'T' from ...

当一个模型使用 deepmd/1.2 训练,但是用更高版本的 deepmd-kit (> v1.3) 进行 lammps 任务的时候经常会报这个错,例子:

但是,现在发现这个报错在压缩 v1.3 版本模型的时候也会出现。使用下列命令:

dp compress ${input} --checkpoint-folder ${ckpt} 1.3-model.pb -o compressed-model.pb -l compress.log

其中${input}${ckpt}分别是对应模型的输入脚本所在路径和检查点目录。在这个例子里,我们仅把需要压缩的模型复制到了工作文件夹下,输入脚本所在路径和检查点目录人工指认。至于为什么这样会报错 ‘ValueError’,目前还没有找到原因。

因此,我们建议 备份之前的训练文件夹,在训练文件夹的一个 copy 下进行压缩任务

Extra Support

Script for convertion from cp2k xyz to numpy set

from ase.io import read
import numpy as np
import os, sys
import glob
import shutil


#############################
# USER INPUT PARAMETER HERE #
#############################

# input data path here, string, this directory should contains
#   ./data/*frc-1.xyz ./data/*pos-1.xyz
data_path = "./data"

#input the number of atom in system
atom_num = 189

#input cell paramter here
cell = [[10.0,0,0],[0,10.0,0],[0,0,10.0]]

# conversion unit here, modify if you need
au2eV = 2.72113838565563E+01
au2A = 5.29177208590000E-01


####################
# START OF PROGRAM #
####################

def xyz2npy(pos, atom_num, output, unit_convertion=1.0):
    total = np.empty((0,atom_num*3), float)
    for single_pos in pos:
        tmp=single_pos.get_positions()
        tmp=np.reshape(tmp,(1,atom_num*3))
        total = np.concatenate((total,tmp), axis=0)
    total = total * unit_convertion
    np.save(output, total)

def energy2npy(pos, output, unit_convertion=1.0):
     total = np.empty((0), float)
     for single_pos in pos:
         tmp=single_pos.info.pop('E')
         tmp=np.array(tmp,dtype="float")
         tmp=np.reshape(tmp,1)
         total = np.concatenate((total,tmp), axis=0)
     total = total * unit_convertion
     np.save(output,total)

def cell2npy(pos, output, cell, unit_convertion=1.0):
    total = np.empty((0,9),float)
    frame_num = len(pos)
    cell = np.array(cell, dtype="float")
    cell = np.reshape(cell, (1,9))
    for frame in range(frame_num):
        total = np.concatenate((total,cell),axis=0)
    total = total * unit_convertion
    np.save(output,total)

def type_raw(single_pos, output):
    element = single_pos.get_chemical_symbols()
    element = np.array(element)
    tmp, indice = np.unique(element, return_inverse=True)
    np.savetxt(output, indice, fmt='%s',newline=' ')


# read the pos and frc
data_path = os.path.abspath(data_path)
pos_path = os.path.join(data_path, "*pos-1.xyz")
frc_path = os.path.join(data_path, "*frc-1.xyz")
#print(data_path)
pos_path = glob.glob(pos_path)[0]
frc_path = glob.glob(frc_path)[0]
#print(pos_path)
#print(frc_path)
pos = read(pos_path, index = ":" )
frc = read(frc_path, index = ":" )

# numpy path
set_path = os.path.join(data_path, "set.000")
if os.path.isdir(set_path):
    print("detect directory exists\n now remove it")
    shutil.rmtree(set_path)
    os.mkdir(set_path)
else:
    print("detect directory doesn't exist\n now create it")
    os.mkdir(set_path)
type_path = os.path.join(data_path, "type.raw")
coord_path = os.path.join(set_path, "coord.npy")
force_path = os.path.join(set_path, "force.npy")
box_path = os.path.join(set_path, "box.npy")
energy_path = os.path.join(set_path, "energy.npy")


#tranforrmation
xyz2npy(pos, atom_num, coord_path)
xyz2npy(frc, atom_num, force_path, au2eV/au2A)
energy2npy(pos, energy_path, au2eV)
cell2npy(pos, box_path, cell)
type_raw(pos[0], type_path)

升级到DeePMD-kit 2.0

目前 DeePMD-kit 2.0 正式版已经发布,相比旧版已有众多提升,且压缩模型为正式版特性。目前我们集群上已安装 DeePMD-kit 2.0.3。

输入文件

DeePMD-kit 2.0 相比 1.x 在输入文件上做了一定改动,以下给出一个 DeePMD-kit 2.0 输入文件的例子:

{
    "_comment": " model parameters",
    "model": {
        "type_map": [
            "O",
            "H"
        ],
        "descriptor": {
            "type": "se_e2_a",
            "sel": [
                46,
                92
            ],
            "rcut_smth": 0.50,
            "rcut": 6.00,
            "neuron": [
                25,
                50,
                100
            ],
            "resnet_dt": false,
            "axis_neuron": 16,
            "seed": 1,
            "_comment": " that's all"
        },
        "fitting_net": {
            "neuron": [
                240,
                240,
                240
            ],
            "resnet_dt": true,
            "seed": 1,
            "_comment": " that's all"
        },
        "_comment": " that's all"
    },
    "learning_rate": {
        "type": "exp",
        "decay_steps": 5000,
        "start_lr": 0.001,
        "stop_lr": 3.51e-8,
        "_comment": "that's all"
    },
    "loss": {
        "type": "ener",
        "start_pref_e": 0.02,
        "limit_pref_e": 1,
        "start_pref_f": 1000,
        "limit_pref_f": 1,
        "start_pref_v": 0,
        "limit_pref_v": 0,
        "_comment": " that's all"
    },
    "training": {
        "training_data": {
            "systems": [
                "../data/data_0/",
                "../data/data_1/",
                "../data/data_2/"
            ],
            "batch_size": "auto",
            "_comment": "that's all"
        },
        "validation_data": {
            "systems": [
                "../data/data_3"
            ],
            "batch_size": 1,
            "numb_btch": 3,
            "_comment": "that's all"
        },
        "numb_steps": 1000000,
        "seed": 10,
        "disp_file": "lcurve.out",
        "disp_freq": 100,
        "save_freq": 1000,
        "_comment": "that's all"
    },
    "_comment": "that's all"
}

DeePMD-kit 2.0 提供了对验证集(Validation Set)的支持,因而用户可指定某一数据集作为验证集,并输出模型在该数据集上的误差。 相比旧版而言,新版输入文件参数的具体含义变化不大,除了对数据集的定义外,大部分参数含义保持一致。

以下列出一些需要注意的事项:

  1. 训练数据集不再直接写在 training 下,而是写在 training 的子键 training_data 下,格式如下所示:
    "training_data": {
             "systems": [
                 "../data/data_0/",
                 "../data/data_1/",
                 "../data/data_2/"
             ],
             "batch_size": "auto"
         }
    
    默认情况下,每一训练步骤中,DeePMD-kit随机从数据集中挑选结构加入本轮训练,这一步骤加入数据的多少取决于 batch_size 的大小,此时,各 system 中数据被使用的概率是均等的。 若希望控制各 system 数据的权重,可使用 auto_prob 来控制,其参数选项如下所示
    • prob_uniform: 各 system 数据权重均等。
    • prob_sys_size: 各 system 数据的权重取决于其各自的大小。
    • prob_sys_size: 写法示例如下:sidx_0:eidx_0:w_0; sidx_1:eidx_1:w_1;...。 该参数中,sidx_ieidx_i 表示第 i 组数据的起止点,规则同 Python 语法中的切片,w_i 则表示该组数据的权重。在同一组中,各 system 数据的权重取决于各自的大小。 batch_size 的值可手动设定,根据经验一般根据“乘以原子数≤32”的规则设定。新版则支持自动设定,若设定为"auto"则表示按照此规则自动设置,若设定为"auto:N"则根据“乘以原子数≤N”的规则设定。
  2. save_ckpt, load_ckpt, decay_rate 等为过时参数,若由 1.x 迁移,请删除这些参数,否则会导致报错。
  3. n_neuron 更名为 neuronstop_batch 更名为 numb_steps,请注意更改。对应地,decay rate 由 start_lrstop_lr 决定。
  4. lcurve.out 中删除了测试数据的 RMSE 值,因此旧版作图脚本需要对应修改,减少列数(能量在第3列,力在第4列)。若指定了验证集,则会输出模型在验证集上的 RMSE。

更多详细说明,请参见官方文档

评论