一键操作，十倍加速——DP模型压缩功能新鲜出炉
DP Model compression
DP已经在物理、化学、材料、生物、地质等学科领域产生了广泛的影响。然而，对于许多需要大系统规模或长时间尺度的重要问题，DP模拟的计算成本仍然非常高。例如，在一台具有150个V100 GPU卡的集群上使用DP对100万个铜原子进行1ns时长的模拟，大约需要耗费半周时间。
最近，DeePMD-kit核心开发团队实现了Model Compression这一新功能。只需要一个简单的指令dp compress，Model Compression就能在数分钟内对训练好的DP模型进行压缩处理。在和原模型保持一致精度的情况下，Compress模型的模拟速度提升可达10倍以上，同时压缩模型的模拟体系可扩大最多20倍！Model compression paper已经上架arxiv：https://arxiv.org/abs/2107.02103，感兴趣的小伙伴欢迎点击查看。
Model Compression的精度以及性能
图1说明了原模型和Compress模型在 Cu、H2O 和 Al-Cu-Mg 系统上的性能表现。对于大多数 DP应用来说，step = 0.01的误差小于训练误差两到三个数量级，已经足够准确，其中step为model compression参数，详见参数介绍部分。此外我们通过使用DP模型进行MD模拟进一步计算液态水的结构特性：如图2所示，所有具有不同步长的压缩模型预测的O-O、O-H和H-H的径向分布函数与AIMD结果都非常吻合。
进一步对比原模型和Compress性能表现可以发现，在几乎没有精度损失的情况下，根据不同的模拟体系和训练参数，Compress模型在CPU和GPU设备的加速可以达到10倍以上。同时，Model Compresssion可以极大地减少内存使用量，在相同的硬件条件下压缩后的模型内存使用量减少多达20倍。
Model compression的使用方式
用户可访问github：https://github.com/deepmodeling/deepmd-kit/主页快速获取模型压缩功能的使用帮助。一旦经过DP训练获得冻结的模型后，我们可以通过以下命令对其进行压缩（需要编译安装最新的DP-devel分支代码或者使用DP Release的2.0.0以后的最新版本conda安装包），简单直接：
dp compress input.json -i graph.pb -o graph-compress.pb
其中，input.json表示和原模型一致的训练脚本，graph.pb表示原模型，graph-compress.pb表示输出的压缩模型。在未来迭代的版本中，我们会优先考虑去掉对于input.json的依赖。其他支持的额外的输入参数可以通dp compress -h命令查看。
Model Compression实现简要介绍
一、使用五阶多项式拟合Embedding net
当使用一个训练好的DP模型进行模型推理的时候，embedding net的训练参数已经固化，如图3(b)所示：可以被看作是一个把一维向量映射到高维空间的矩阵的光滑函数 ，其中M表示embedding net的最外层神经网络大小。我们考虑使用多项式插值的方式拟合embedding net,在确保精度的情况下，同时也为了避免产生龙格现象 以及插值参数过大的问题，我们使用分段五阶多项式来对embedding net进行打表插值。插值范围也就是的值域由程序遍历训练数据集检测得到。
二、高效的核融合操作
我们观察到，在未经过模型压缩处理的DP模型的执行过程中，embedding matrix 通常消耗主机/设备总内存使用量的 90% 以上。具体来说，对于原子i的执行过程，首先通过embedding net产生存放到主机/设备内存，然后被加载到寄存器和执行矩阵乘法操作。上述计算过程导致了寄存器和主机/设备内存之间的巨大的数据访问开销。
在图3(c)所示的压缩DP工作流程中，在使用五阶多项式将插值后，我们执行有效的核融合操作来产生。以GPU的实现为例，对于每个原子i，矩阵乘法由一个单独的CUDA线程块处理。详细地说，每次计算的一列并将其存储在寄存器中（不放回全局内存）。接下来，将的对应行加载到寄存器中，与的列执行外积 运算。外积运算最多必须进行，最终产生维度为 的矩阵。
三、冗余性去除
在使用DP模型进行推理时，在一定的参数设置条件下，可能会导致environment matrix 出现冗余，从而使其具有一定的稀疏性 ：如图3（c）所示的维度为，其中是训练时给定的经验参数 的加和：。在训练过程中，为了兼顾训练数据集的多样性，一般设置的都比较大。但其中的有效邻居数量则和设置以及LAMMPS给定的邻居表相关，实际模型推理过程中原子的有效邻居数量往往会小于，从而给带来了一定的稀疏性。
如图3(a)所示，在原有的DP模型中需要经过embedding net，而embedding net的实现依赖于Eigen数学库和cuBLAS等基础库，对于矩阵乘法的维度有着明确的要求，难以利用的冗余性。而新的实现使用五阶多项式拟合embedding net，基于自定义的C++和CUDA函数的新实现，可以方便的利用的稀疏性，只计算的有效邻居信息，这进一步优化了程序的性能。
关于DP Model Compression的补充解释
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关键参数解释(参照GITHUB MARKDOWN)
在实际使用的过程中step(-s)以及extrapolate(-e)参数可能会被用户使用到，现将参数使用介绍如下：
模型压缩使用五阶多项式对embedding net进行插值。它引入了了两个具有不同步长的表格来存储多项式的系数参数。第一个表的插值范围覆盖了训练数据的范围，而第二个表的覆盖范围可以被看作是训练数据的外推。每个表格的插值范围被给定的步长均匀划分，而step(-s)参数表示第一个表格的插值步长，而第二个表格的插值步长将使用 10 * step。通常情况下step的取值范围从0.1到0.001。更小的步长意味着更高的精度和更大的参数表格尺寸（参数默认值：0.01）。
前文提到，第一个表格的插值覆盖范围是由代码自动检测得到的：。而第二表格插值覆盖范围从第一个表的上边界到,即。其中extrapolate是输入参数（参数默认值：5）
下面是dp compress -h命令的输出，可以作为参考理解：