简单操作，训练加速3倍｜DP压缩训练功能上线
上周六，DP团队正式发布了了DeePMD-kit的新版本v2.0.0，本次大版本更新也是对这款备受关注的软件全方位的升级，作为AI+分子模拟领域的领跑者，DeePMD-kit v2.0.0不仅为社区开发者提供了更灵活的硬件支持、更丰富的模型架构，还大大提高了训练推断的效率，其中，DP模型压缩为推断的速度带来了一个数量级的提升，而本文要介绍的DP压缩训练则进一步提高了模型训练的效率，为用户在使用DP的第一步带来更丝滑的体验。
DP Compressed Training
DP在物理、化学、材料、生物、地质等各个领域已经产生了广泛的影响，推动着第一性原理精度的分子动力学模拟。对于许多需要大系统规模或长时间尺度的重要问题，在拥有训练好的DP模型之后，通过之前实现的DP模型压缩功能，我们已经可以实现高速推断。然而，获取高精度的DP模型本身需要经过长时间多步数的训练拟合，训练所需要的时间与计算成本会随着体系规模、元素种类数的扩大而增加，为使用DP的第一步带来较大的困难。例如，在单卡V100 GPU上，用DP拟合500原子左右的蛋白主链体系，训练10万步已经需要花费半周的时间；对于一些较为复杂的合金体系，后期需要经过长时间的拟合微调，带来的计算成本也是巨大的。
近期，DeePMD-kit开发团队在Model Compression这一功能的基础上，进一步实现了Compressed Training，只需在dp train正常训练指令的末尾加上--init-frz-model graph-compress.pb 即可将已通过dp compress命令压缩好的模型graph-compress.pb作为训练的初始化模型，由于压缩后的模型可学习参数比原模型要少很多，从而可以实现对压缩模型的快速训练。压缩训练基本上可以取代正常训练中后期的微调过程，在小体系已经能实现3倍以上的加速。
Compressed Training
使用场景与使用方法
DP模型本身有两部分网络参数，分别为产生原子描述符(descriptor)的嵌入网络(embedding net)参数和真正用来拟合能量等目标的拟合网络(fitting net)参数。压缩训练目前主要有两种使用场景，分别是标准训练后期的模型微调以及从头开始训练模型，两者虽然在整体加速倍数上有差别（原因是后者仍然需要一段时间的标准训练），但本质上均为固定并转换了模型的一部分参数(embedding net部分)，只去训练另一部分参数(fitting net部分)的过程。
【场景一 训练后期的模型微调】
对于合金等体系，想要得到高精度的模型，需要经过长时间的训练与后续的微调，而快速的压缩训练基本上可以完全取代正常训练后期的微调过程。我们首先需要通过标准训练并冻结得到一个基础模型graph.pb，然后我们运行指令：
dp compress -i graph.pb -o graph-compress.pb输入原模型graph.pb，输出压缩后的模型graph-compress.pb（需要注意的是dp compress指令在新版本中去除了部分对训练脚本json文件的显式依赖，详情见文末补充信息），接着开始对压缩后的模型进行继续训练：
dp train input.json --init-frz-model graph-compress.pb其中input.json为训练脚本，即可实现以压缩模型为初始化的微调过程。
dp compress的本质是用多项式去拟合模型中的embedding网络参数，相当于固定了这部分参数并加速了运算；而继续进行的压缩训练则只需微调剩余的fitting网络部分，由于正反向运算显著减少，所以速度会有较大提升，在小体系（如example/water）已经能实现3倍以上的加速。
之所以能够这样选择性地微调，是因为网络两部分参数具有不同的训练性质：embedding网络规模较小且运算步骤较多，fitting网络较大但拟合能力要更强，一般在训练过程中，embedding网络较早趋于收敛，其实完全可以将其参数固定，将大量时间花在学习fitting网络参数上，这样会节省很多学习成本。实验表明在embedding网络接近收敛时，仅提高fitting网络的拟合能力也完全够用，所以压缩训练基本上可以完全取代正常训练后期的这一微调过程，从而实现加速。
【场景二 从头开始训练】
与场景一类似，区别不过是缺少embedding部分接近收敛的基础模型，我们只需要在标准训练到某一时刻停止，然后开始压缩训练即可，按经验来说，在达到整体训练总步数的20%左右停止是个不错的选择。举例来说，若给定体系训练100万步能达到不错的效果，我们可以先设定训练总步数numb_steps为20万步，通过：
dp train input_1.json && dp freeze -o graph.pb进行标准训练并冻结得到基础模型，之后重复场景一中的步骤：
dp compress -i graph.pb -o graph-compress.pbdp train input_2.json --init-frz-model graph-compress.pb设定继续训练80万步停止，按经验来说，此时冻结后得到的模型应该和直接训练100万步的模型效果接近，通过这样的流程从头训练也可以在大体系实现整体3倍左右的加速。
需要注意的是上述input_1.json和input_2.json训练脚本中，除了numb_steps参数不同，start_lr和stop_lr也要根据训练阶段调整，使其达到最佳训练效果，当然20%这个停止比例对于不同体系可能有所差别，欢迎大家尝试并反馈。
Compressed Training结果展示
实验选择了两个体系：样例水体系（即example/water，含64个水分子）以及一个蛋白主链体系（约500个原子、20+种embedding网络），实验设备为单卡V100并进行双精度计算，图1(a)展示了分别在上述两个体系进行标准训练、压缩训练的训练速度（其中‘custom ops on GPU’表示开启了自定义算子的GPU版本，详情见文末补充信息）。可以看到使用压缩训练加速比随体系规模与元素种类的增加而增加，在开启自定义算子的GPU版本后，配合压缩训练可以达到3倍以上的加速比，在训练后期模型微调的场景下，可以在几乎不损失训练效果的情况下享受这一加速比。
对于从头训练的场景，图1(b)展示了使用不同训练策略在上述两个体系进行从头训练的时间对比，即完全进行标准训练和进行总步数20%的标准训练加80%的压缩训练两种策略，同样可以达到2倍以上的加速效果。图2也对比了这两种策略在样例水体系训练的学习曲线，可以看到在总训练步数为100万步时，两种策略训练出来的模型精度几乎没有差别，甚至压缩训练策略的误差降低得更快。
图1 压缩训练相较于标准训练的性能比较
图2 压缩训练策略的精度比较
Compressed Training实现简介
模型压缩的实现介绍可见往期推送，本质上是用多项式来拟合DeePMD模型中的embedding网络参数，从而固定参数并实现了计算加速。如图3(a)所示是DeePMD标准模型的运算过程，包括前向运算得出能量E、反向求导得到受力F（或者维力V）以及整体再反向一次进行梯度更新的过程，其中黄色箭头代表仅使用了Tensorflow框架的标准算子的运算过程，蓝色箭头代表使用我们自定义算子的运算过程。
图3（a）DeePMD标准模型的运算过程
图3(b)所示则为DeePMD压缩模型的运算过程，我们从运算图可以看出，压缩模型的推断与训练只是用自定义算子替代了一部分标准的Tensorflow运算（即由环境矩阵
图3（b）DeePMD压缩模型的运算过程
另一方面，之前实现的模型压缩只能实现推断过程的加速，是因为只实现了这部分两个方向的自定义算子（即正向与一阶导），缺少图3(b)中红色所圈出的二阶导算子，将其补全即可实现压缩训练的功能。从运算图也可以看出，压缩训练并不会改变前一部分多项式运算所含参数，只会更新后续fitting网络的参数，而对于DeePMD模型而言，embedding网络较早便趋于收敛，这种情况下，采用压缩训练相当于把训练重心完全转移到了优化fitting网络拟合能力的过程上来，在保证训练效果的前提下提升了训练效率。
补充介绍
【Model compression接口变化】
我们在最新发布的DeePMD-kit v2.0.0中优化了模型压缩dp compress命令的接口，去除了对训练脚本的显式依赖。若graph.pb是由v2.0.0直接训练得到的模型，脚本相关信息在训练时便直接保存到了模型中，可以直接用以下命令来进行模型压缩：
dp compress -i graph.pb -o graph-compress.pb若基础模型为之前v2.0beta版本训练得到，需要增加--training-script 选项来显式指定当初所用训练脚步，即：
dp compress -i graph.pb -o graph-compress.pb --training-script input.json需要注意input.json中要包含正确的数据路径。
若基础模型为v1.3及之前版本训练得到，则需要先用dp convert-from命令将其转换为v2.0beta版本模型：
dp convert-from 1.3 -i graph_old.pb -o graph.pb接着再进行上述操作来完成模型压缩。
【Custom op开启GPU运算】
如图3(a)所示，其中蓝色箭头为我们实现的自定义算子部分，即使网络整体部分在GPU环境运行，DeePMD-kit v2.0.0自定义算子这部分会默认使用CPU计算，用户需要根据机器环境来修改setup.py中的dp_variant变量，选择开启自定义算子的GPU版本：
目前自定义算子支持CPU、CUDA、ROCm三种版本，在选择开启CUDA版本之后，如之前图1(a)中前两项所示，可以直接带来1.4倍左右的加速。
【补充资料】
1.DeePMD原理paper
Zhang L, Han J, Wang H, et al. Deep potential molecular dynamics: a scalable model with the accuracy of quantum mechanics[J]. Physical review letters, 2018, 120(14): 143001.
2.Model compression原理paper
https://arxiv.org/abs/2107.02103
3.Model compression上期推送：点此查看