Skip to content

v3.0.0

Latest
Latest
Compare
Choose a tag to compare
@njzjz njzjz released this 23 Nov 08:10
v3.0.0
e695a91
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Learn about vigilant mode.

DeePMD-kit v3: Multiple-backend Framework, DPA-2 Large Atomic Model, and Plugin Mechanisms

After eight months of public tests, we are excited to present the first stable version of DeePMD-kit v3, an advanced version that enables deep potential models with TensorFlow, PyTorch, or JAX backends. Additionally, DeePMD-kit v3 introduces support for the DPA-2 model, a novel architecture optimized for large atomic models. This release enhances plugin mechanisms, making integrating and developing new models easier.

Highlights

Multiple-backend framework: TensorFlow, PyTorch, and JAX support

image

DeePMD-kit v3 adds a versatile, pluggable framework providing consistent training and inference experience across multiple backends. Version 3.0.0 includes:

  • TensorFlow backend: Known for its computational efficiency with a static graph design.
  • PyTorch backend: A dynamic graph backend that simplifies model extension and development.
  • DP backend: Built with NumPy and Array API, a reference backend for development without heavy deep-learning frameworks.
  • JAX backend: Based on the DP backend via Array API, a static graph backend.
Features TensorFlow PyTorch JAX DP
Descriptor local frame ✅
Descriptor se_e2_a ✅ ✅ ✅ ✅
Descriptor se_e2_r ✅ ✅ ✅ ✅
Descriptor se_e3 ✅ ✅ ✅ ✅
Descriptor se_e3_tebd ✅ ✅ ✅
Descriptor DPA1 ✅ ✅ ✅ ✅
Descriptor DPA2 ✅ ✅ ✅
Descriptor Hybrid ✅ ✅ ✅ ✅
Fitting energy ✅ ✅ ✅ ✅
Fitting dipole ✅ ✅ ✅ ✅
Fitting polar ✅ ✅ ✅ ✅
Fitting DOS ✅ ✅ ✅ ✅
Fitting property ✅ ✅ ✅
ZBL ✅ ✅ ✅ ✅
DPLR ✅
DPRc ✅ ✅ ✅ ✅
Spin ✅ ✅ ✅
Gradient calculation ✅ ✅ ✅
Model training ✅ ✅
Model compression ✅ ✅
Python inference ✅ ✅ ✅ ✅
C++ inference ✅ ✅ ✅

Critical features of the multiple-backend framework include the ability to:

  • Train models using different backends with the same training data and input script, allowing backend switching based on your efficiency or convenience needs.
# Training a model using the TensorFlow backend
dp --tf train input.json
dp --tf freeze
dp --tf compress

# Training a model using the PyTorch backend
dp --pt train input.json
dp --pt freeze
dp --pt compress
  • Convert models between backends using dp convert-backend, with backend-specific file extensions (e.g., .pb for TensorFlow and .pth for PyTorch).
# Convert from a TensorFlow model to a PyTorch model
dp convert-backend frozen_model.pb frozen_model.pth
# Convert from a PyTorch model to a TensorFlow model
dp convert-backend frozen_model.pth frozen_model.pb
# Convert from a PyTorch model to a JAX model
dp convert-backend frozen_model.pth frozen_model.savedmodel
# Convert from a PyTorch model to the backend-independent DP format
dp convert-backend frozen_model.pth frozen_model.dp
  • Run inference across backends via interfaces like dp test, Python/C++/C interfaces, or third-party packages (e.g., dpdata, ASE, LAMMPS, AMBER, Gromacs, i-PI, CP2K, OpenMM, ABACUS, etc.).
# In a LAMMPS file:
# run LAMMPS with a TensorFlow backend model
pair_style deepmd frozen_model.pb
# run LAMMPS with a PyTorch backend model
pair_style deepmd frozen_model.pth
# run LAMMPS with a JAX backend model
pair_style deepmd frozen_model.savedmodel
# Calculate model deviation using different models
pair_style deepmd frozen_model.pb frozen_model.pth frozen_model.savedmodel out_file md.out out_freq 100
  • Add a new backend to DeePMD-kit much more quickly if you want to contribute to DeePMD-kit.

DPA-2 model: a large atomic model as a multi-task learner

The DPA-2 model offers a robust architecture for large atomic models (LAM), accurately representing diverse chemical systems for high-quality simulations. In this release, DPA-2 can be trained using the PyTorch backend, supporting both single-task (see examples/water/dpa2) or multi-task (see examples/water_multi_task/pytorch_example) training schemes. DPA-2 is available for Python/C++ inference in the JAX backend.

The DPA-2 descriptor comprises repinit and repformer, as shown below.

DPA-2

The PyTorch backend supports training strategies for large atomic models, including:

  • Parallel training: Train large atomic models on multiple GPUs for efficiency.
torchrun --nproc_per_node=4 --no-python dp --pt train input.json
  • Multi-task training: For large atomic models trained across a broad range of data calculated on different DFT levels with shared descriptors. An example is given in examples/water_multi_task/pytorch_example/input_torch.json.
  • Finetune: Training a pre-train large atomic model on a smaller, task-specific dataset. The PyTorch backend has supported --finetune argument in the dp --pt train command line.

Plugin mechanisms for external models

In version 3.0.0, the plugin capabilities have been implemented to support the development and integration of potential energy models using TensorFlow, PyTorch, or JAX backends, leveraging DeePMD-kit's trainer, loss functions, and interfaces. A plugin example is deepmd-gnn, which supports training the MACE and NequIP models in the DeePMD-kit with the familiar commands.

dp --pt train mace.json
dp --pt freeze
dp --pt test -m frozen_model.pth -s ../data/

image

Other new features

  • Descriptor se_e3_tebd. (#4066)
  • Fitting the property (#3867).
  • New training parameters: max_ckpt_keep (#3441), change_bias_after_training (#3993), and stat_file.
  • New command line interface: dp change-bias (#3993) and dp show (#3796).
  • Support generating JSON schema for integration with VSCode (#3849).
  • The latest LAMMPS version (stable_29Aug2024_update1) is supported. (#4088, #4179)

Breaking changes

  • The deepmodeling conda channel is deprecated. Use the conda-forge channel instead. (#3462, #4385)
  • The offline package and conda packages for CUDA 11 are dropped.
  • Python 3.7 and 3.8 supports are dropped. (#3185, #4185)
  • The minimal versions of deep learning frameworks: TensorFlow 2.7, PyTorch 2.1, JAX 0.4.33, and NumPy 1.21.
  • We require all model files to have the correct filename extension for all interfaces so a corresponding backend can load them. TensorFlow model files must end with .pb extension.
  • Bias is removed by default from type embedding. (#3958)
  • The spin model is refactored, and its usage in the LAMMPS module has been changed. (#3301, #4321)
  • Multi-task training support is removed from the TensorFlow backend. (#3763)
  • The set_prefix key is deprecated. (#3753)
  • dp test now uses all sets for training and test. In previous versions, only the last set is used as the test set in dp test. (#3862)
  • The Python module structure is fully refactored. The old deepmd module was moved to deepmd.tf without other API changes, and deepmd_utils was moved to deepmd without other API changes. (#3177, #3178)
  • Python class DeepTensor (including DeepDiople and DeepPolar) now returns atomic tensor in the dimension of natoms instead of nsel_atoms. (#3390)
  • C++ 11 support is dropped. (#4068)

For other changes, refer to Full Changelog: v2.2.11...v3.0.0rc0

Contributors

The PyTorch backend was developed in the dptech-corp/deepmd-pytorch repository, and then it was fully merged into the deepmd-kit repository in #3180. Contributors to the deepmd-pytorch repository:

Contributors to the deepmd-kit repository:

We also thank everyone who did tests and reported bugs in the past eight months.

Contributors

amcadmus, chazeon, and 29 other contributors
Assets 7
Loading
5 Join discussion