Skip to content

v3.0.0rc0

Pre-release
Pre-release
Compare
Choose a tag to compare
@njzjz njzjz released this 14 Nov 19:36
· 106 commits to devel since this release
v3.0.0rc0
0ad4289
This commit was created on GitHub.com and signed with GitHub’s verified signature.
GPG key ID: B5690EEEBB952194
Verified
Learn about vigilant mode.

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

We are excited to present the first release candidate 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: Towards a universal large atomic model for molecular and material simulation

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 is trainable in the PyTorch backend, with an example configuration available in examples/water/dpa2. DPA-2 is available for Python 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 v3.0.0, plugin capabilities allow you to develop models with TensorFlow, PyTorch, or JAX, 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

  • Python 3.7 and 3.8 supports are dropped. (#3185, #4185)
  • 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:

Contributors

amcadmus, chazeon, and 29 other contributors
Assets 9
Loading