DeePMD-kit v3：多后端框架、大原子模型支持
The following article is from 深度势能 Author 曾晋哲
在DeePMD-kit v2发布三年之后，今天，我们很高兴地发布DeePMD-kit v3的首个alpha版本（v3.0.0a0）。DeePMD-kit v3允许您在TensorFlow或PyTorch任意一个框架之上训练和运行深度势能模型。DeePMD-kit v3还新增了对DPA-2模型的支持，翻开了大原子模型（LAM）的新篇章。
多后端框架
DeePMD-kit v3添加了多后端框架的支持，可以在不同后端之间提供一致的训练和推理体验。您可以：
使用相同的训练数据和输入脚本，利用不同的后端来训练深度势能模型，并基于效率、功能或便利性切换后端：
# 用TensorFlow后端训练模型dp --tf train input.jsondp --tf freeze
# 用PyTorch后端训练模型dp --pt train input.jsondp --pt freeze通过任何现有接口，使用任何模型进行推理。现有接口包括dp test、Python/C++/C接口和第三方包（dpdata、ASE、LAMMPS、AMBER、Gromacs、i-PI、CP2K、OpenMM、ABACUS等）。以LAMMPS为例：
# 用TensorFlow后端模型运行LAMMPSpair_style deepmd frozen_model.pb# 用PyTorch后端模型运行LAMMPSpair_style deepmd frozen_model.pth# 同时用两个后端的模型计算模型偏差pair_style deepmd frozen_model.pb frozen_model.pth out_file md.out out_freq 100如果两个后端都支持某个模型，可以使用dp convert-backend在后端之间转换模型：
dp convert-backend frozen_model.pb frozen_model.pthdp convert-backend frozen_model.pth frozen_model.pb如果您想为DeePMD-kit贡献一个新的后端，可以更快地实现这个后端。
PyTorch后端：用于大原子模型和科学探索的后端
我们在DeePMD-kit v3中添加了PyTorch后端，以支持新模型的开发，特别是针对大原子模型。
DPA-2模型：面向分子和材料模拟的通用大原子模型
DPA-2模型是大原子模型的一种新架构，能够准确表示多种化学和材料体系，与传统方法相比，进行高质量的模拟和预测所需的努力大为减少。DPA-2模型仅在PyTorch后端实现。示例配置位于examples/water/dpa2目录中。
DPA-2描述符包括两个主要部分：repinit和repformer。详细架构如下图所示。
大原子模型的训练策略
PyTorch后端已支持构建大原子模型的多种训练策略。
并行训练：大原子模型有大量超参数和复杂架构，因此有必要在多个GPU上训练模型。得益于PyTorch社区生态系统，PyTorch后端的并行训练可以通过torchrun来驱动，torchrun是用于分布式数据并行的PyTorch启动器。
torchrun --nproc_per_node=4 --no-python dp --pt train input.json多任务训练：大原子模型针对广泛的数据和不同的DFT方法进行训练，因而需要多任务训练。PyTorch后端支持多任务训练，在不同的任务中共享描述符。一个示例位于
examples/water_multi_task/pytorch_example/input_torch.json。
微调模型：微调可用于在更小的、特定任务的数据集上训练预训练的大模型。PyTorch后端已支持在dp --pt train命令行中使用--finetune参数。
使用Python和动态图开发新模型
TensorFlow后端的静态图和自定义C++算子牺牲了研究便利性以换取计算性能，但研究者可能会对这些设计感到困惑。PyTorch后端具有良好设计的代码结构，使用动态图编写，目前完全使用Python语言编写，因此在PyTorch后端扩展和调试新的深度势能模型比在TensorFlow后端更容易。
支持传统深度势能模型
研究者可能仍然想在PyTorch后端使用TensorFlow后端已经支持的传统模型，并在不同的后端之间比较同一个模型。我们在PyTorch后端几乎重写了所有传统模型，如下：
已支持的功能：
描述符：se_e2_a、se_e2_r、se_atten、hybrid；
拟合网络：energy、dipole、polar、fparam/apram支持
模型：standard、DPRc
Python推理接口
C++推理接口（仅限势能）
TensorBoard
尚未支持的功能：
描述符：se_e3、se_atten_v2、se_e2_a_mask
拟合：dos
模型：linear_ener、DPLR、pairtab、linear_ener、frozen、pairwise_dprc、ZBL、Spin
模型压缩
Python推理接口（仅DPLR）
C++推理接口（仅DeepTensor和DPLR）
使用Horovod进行并行训练
没有计划支持的功能：
描述符：loc_frame、se_e2_a + type embedding、se_a_ebd_v2
NVNMD
DP后端和格式：其它后端的参考后端
DP后端是用于开发的参考后端，使用纯NumPy实现模型，不使用任何重型深度学习框架。DP后端不能用于训练，只能用于Python推理。作为参考后端，DP后端的目标不是最佳的性能，而只是正确的结果。DP后端使用HDF5存储模型序列化数据，无需调用任何深度学习框架。
DP后端和序列化数据在单元测试中使用，以确保不同后端有一致的结果，并且可以相互转换。
在当前版本中，DP后端的支持状态与PyTorch后端类似，但DPA-1和DPA-2尚未被支持。
作者贡献
Hangrui Bi, Chun Cai, Junhan Chang, Yiming Du, Guolin Ke, Xinzijian Liu, Anyang Peng, Xuejian Qin, Han Wang, Jinzhe Zeng, Chengqian Zhang, Duo Zhang, Xiangyu Zhang对以上功能均有贡献。
完整贡献、完整更新内容、下载预安装包，可查阅下方网址，或点击阅读原文即可访问：
https://github.com/deepmodeling/deepmd-kit/releases/tag/v3.0.0a0
推荐关注
关于深势科技
我们开创性地提出了「多尺度建模+机器学习+高性能计算」的革命性科学研究新范式，并推出了Bohrium®科研云平台、Hermite®药物计算设计平台、RiDYMO®难成药靶标研发平台及 Piloteye®电池设计自动化平台等工业设计与仿真基础设施，颠覆了现有研发模式，打造“计算引导实验、实验优化设计”的全新范式。
深势科技是国家高新技术企业、国家专精特新“小巨人”企业，总部位于北京，并在上海、深圳等城市布局研发中心。科研技术团队由中国科学院院士领衔，汇集了超百位数学、物理、化学、生物、材料、计算机等多个领域的优秀青年科学家和工程师，其中公司的博士及博士后占比超过35%。核心成员获得过2020年全球计算机高性能计算领域的最高奖项“戈登贝尔奖”，相关工作当选2020年中国十大科技进展和全球AI领域十大技术突破。