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.2 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.0 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. DP DP

Note

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

  • Model filename extension: .dp

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 uses HDF5 to store model serialization data, which is backend-independent. Only Python inference interface can load this format.

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.

Warning

Currently, only the se_e2_a model fully supports the backend conversion between TensorFlow TensorFlow and PyTorch PyTorch.