Backend

1. Backend#

1.1. Supported backends#

DeePMD-kit supports multiple backends: TensorFlow and PyTorch. To use DeePMD-kit, you must install at least one backend. Each backend does not support all features. In the documentation, TensorFlow TensorFlow and PyTorch PyTorch icons are used to mark whether a backend supports a feature.

1.1.1. TensorFlow TensorFlow#

  • Model filename extension: .pb

  • Checkpoint filename extension: .meta, .index, .data-00000-of-00001

TensorFlow 2.7 or above is required, since NumPy 1.21 or above is required. DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (tf.compat.v1) in the graph mode.

1.1.2. PyTorch PyTorch#

  • Model filename extension: .pth

  • Checkpoint filename extension: .pt

PyTorch 2.1 or above is required. While .pth and .pt are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint.

1.1.3. JAX JAX#

  • Model filename extension: .xlo, .savedmodel

  • Checkpoint filename extension: .jax

JAX 0.4.33 (which requires Python 3.10 or above) or above is required. Both .xlo and .jax are customized format extensions defined in DeePMD-kit, since JAX has no convention for file extensions. .savedmodel is the TensorFlow SavedModel format generated by JAX2TF, which needs the installation of TensorFlow. Only the .savedmodel format supports C++ inference, which needs the TensorFlow C++ interface. The model is device-specific, so that the model generated on the GPU device cannot be run on the CPUs. Currently, this backend is developed actively, and has no support for training.

1.1.4. DP DP#

Note

This backend is only for development and should not take into production.

  • Model filename extension: .dp, .yaml, .yml

DP is a reference backend for development, which uses pure NumPy to implement models without using any heavy deep-learning frameworks. Due to the limitation of NumPy, it doesn’t support gradient calculation and thus cannot be used for training. As a reference backend, it is not aimed at the best performance, but only the correct results. The DP backend has two formats, both of which are backend-independent: The .dp format uses HDF5 to store model serialization data, which has good performance. The .yaml or .yml use YAML to save the data as plain texts, which is easy to read for human beings. Only Python inference interface can load these formats.

NumPy 1.21 or above is required.

1.2. Switch the backend#

1.2.1. Training#

When training and freezing a model, you can use dp --tf or dp --pt in the command line to switch the backend.

1.2.2. Inference#

When doing inference, DeePMD-kit detects the backend from the model filename. For example, when the model filename ends with .pb (the ProtoBuf file), DeePMD-kit will consider it using the TensorFlow backend.

1.3. Convert model files between backends#

If a model is supported by two backends, one can use dp convert-backend to convert the model file between these two backends.